from __future__ import annotations import uuid import sqlalchemy from langbot_plugin.api.entities.builtin.provider import message as provider_message from ....cloud.model_catalog import LANGBOT_MODELS_PROVIDER_REQUESTER from ....core import app from ....entity.persistence import model as persistence_model from ....entity.persistence import pipeline as persistence_pipeline from ....provider.modelmgr import requester as model_requester from ....workspace.errors import WorkspaceNotFoundError from .secrets import mask_secret_value, redact_secrets, restore_secret_placeholders from .tenant import TenantContext, require_workspace_uuid, scope_statement def _parse_provider_api_keys(provider_dict: dict) -> dict: """Parse api_keys if it's a JSON string""" if isinstance(provider_dict.get('api_keys'), str): import json try: provider_dict['api_keys'] = json.loads(provider_dict['api_keys']) except Exception: provider_dict['api_keys'] = [] return provider_dict def _runtime_model_data(model_uuid: str, model_data: dict) -> dict: """Return model data for rebuilding runtime models after an update. Update payloads intentionally omit uuid before writing to the database. Runtime model entities still need the stable uuid so pipeline configs can resolve the in-memory model immediately after an edit, without requiring a process restart. """ return {**model_data, 'uuid': model_uuid} def _redact_model_secrets(model_data: dict) -> dict: """Return a copy with model args and embedded provider credentials masked.""" redacted = model_data.copy() if 'extra_args' in redacted: redacted['extra_args'] = redact_secrets(redacted['extra_args']) if isinstance(redacted.get('provider'), dict): provider = redacted['provider'].copy() # ModelProvider never contains another provider. Dropping this key also # makes the serializer robust to a reused/self-referential test double. provider.pop('provider', None) if 'api_keys' in provider: provider['api_keys'] = mask_secret_value(provider['api_keys']) redacted['provider'] = provider return redacted async def _validate_provider_supports( ap: app.Application, context: TenantContext, provider_uuid: str, model_type: str, ) -> None: """Validate that the provider's requester declares support for ``model_type``. ``model_type`` is one of the manifest ``support_type`` values: 'llm', 'text-embedding', 'rerank'. Raises ValueError when the requester manifest does not list the requested type. This is a server-side guard so a model cannot be attached to a provider that does not support it, even if the frontend tab restriction is bypassed. """ model_mgr = getattr(ap, 'model_mgr', None) if model_mgr is None: return get_provider = getattr(model_mgr, 'get_provider_by_uuid', None) if not callable(get_provider): return try: runtime_provider = await get_provider(context, provider_uuid) except ValueError: return requester_name = getattr(getattr(runtime_provider, 'provider_entity', None), 'requester', None) if not requester_name: return get_manifest = getattr(model_mgr, 'get_available_requester_manifest_by_name', None) if not callable(get_manifest): return manifest = get_manifest(requester_name) if manifest is None: return spec = getattr(manifest, 'spec', None) or {} support_type = spec.get('support_type') if isinstance(spec, dict) else None # When a manifest omits support_type, do not block (backward compatible). if not support_type: return if model_type not in support_type: raise ValueError(f'Provider requester "{requester_name}" does not support {model_type} models') async def _require_workspace_provider( ap: app.Application, context: TenantContext, provider_uuid: str, ) -> dict: """Require the referenced provider to belong to the active Workspace.""" provider = await ap.provider_service.get_provider(context, provider_uuid) if provider is None: raise WorkspaceNotFoundError('Provider not found') return provider def _is_cloud_runtime(ap: app.Application) -> bool: mode = getattr(ap.persistence_mgr, 'mode', None) return getattr(mode, 'value', None) == 'cloud_runtime' async def _assert_cloud_managed_provider_mutable( ap: app.Application, context: TenantContext, provider_uuid: str, ) -> None: if not _is_cloud_runtime(ap): return provider = await _require_workspace_provider(ap, context, provider_uuid) if provider.get('requester') == LANGBOT_MODELS_PROVIDER_REQUESTER: raise ValueError('LangBot Models is managed by Cloud and cannot be modified') async def _require_runtime_provider( ap: app.Application, context: TenantContext, provider_uuid: str, ) -> model_requester.RuntimeProvider: try: return await ap.model_mgr.get_provider_by_uuid(context, provider_uuid) except ValueError as exc: raise Exception('provider not found') from exc class LLMModelsService: ap: app.Application def __init__(self, ap: app.Application) -> None: self.ap = ap async def get_llm_models(self, context: TenantContext, include_secret: bool = False) -> list[dict]: """Get all LLM models with provider info""" result = await self.ap.persistence_mgr.execute_async( scope_statement(sqlalchemy.select(persistence_model.LLMModel), persistence_model.LLMModel, context) ) models = result.all() # Get all providers for lookup providers_result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.ModelProvider), persistence_model.ModelProvider, context ) ) providers = {p.uuid: p for p in providers_result.all()} models_list = [] for model in models: model_dict = self.ap.persistence_mgr.serialize_model(persistence_model.LLMModel, model) provider = providers.get(model.provider_uuid) if provider: provider_dict = self.ap.persistence_mgr.serialize_model(persistence_model.ModelProvider, provider) provider_dict = _parse_provider_api_keys(provider_dict) model_dict['provider'] = provider_dict if not include_secret: model_dict = _redact_model_secrets(model_dict) models_list.append(model_dict) return models_list async def get_llm_models_by_provider( self, context: TenantContext, provider_uuid: str, *, include_secret: bool = False, ) -> list[dict]: """Get LLM models by provider UUID""" await _require_workspace_provider(self.ap, context, provider_uuid) result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.LLMModel).where( persistence_model.LLMModel.provider_uuid == provider_uuid ), persistence_model.LLMModel, context, ) ) models = result.all() serialized = [self.ap.persistence_mgr.serialize_model(persistence_model.LLMModel, m) for m in models] return serialized if include_secret else [_redact_model_secrets(model) for model in serialized] async def create_llm_model( self, context: TenantContext, model_data: dict, preserve_uuid: bool = False, auto_set_to_default_pipeline: bool = True, ) -> str: """Create a new LLM model""" workspace_uuid = require_workspace_uuid(context) model_data = model_data.copy() if not preserve_uuid: model_data['uuid'] = str(uuid.uuid4()) model_data['workspace_uuid'] = workspace_uuid if 'extra_args' in model_data: model_data['extra_args'] = restore_secret_placeholders(model_data['extra_args']) # Handle provider creation if needed if 'provider' in model_data: provider_data = model_data.pop('provider') if provider_data.get('uuid'): model_data['provider_uuid'] = provider_data['uuid'] else: # Create new provider provider_uuid = await self.ap.provider_service.find_or_create_provider( context, requester=provider_data.get('requester', ''), base_url=provider_data.get('base_url', ''), api_keys=provider_data.get('api_keys', []), ) model_data['provider_uuid'] = provider_uuid await _require_workspace_provider(self.ap, context, model_data['provider_uuid']) await _assert_cloud_managed_provider_mutable(self.ap, context, model_data['provider_uuid']) await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'llm') await self.ap.persistence_mgr.execute_async(sqlalchemy.insert(persistence_model.LLMModel).values(**model_data)) runtime_provider = await _require_runtime_provider(self.ap, context, model_data['provider_uuid']) runtime_llm_model = await self.ap.model_mgr.load_llm_model_with_provider( context, persistence_model.LLMModel(**model_data), runtime_provider, ) await self.ap.model_mgr.cache_llm_model(context, runtime_llm_model) if auto_set_to_default_pipeline: # set the default pipeline model to this model result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_pipeline.LegacyPipeline).where( persistence_pipeline.LegacyPipeline.is_default == True ), persistence_pipeline.LegacyPipeline, workspace_uuid, ) ) pipeline = result.first() if pipeline is not None: model_config = pipeline.config.get('ai', {}).get('local-agent', {}).get('model', {}) if not model_config.get('primary', ''): pipeline_config = pipeline.config pipeline_config['ai']['local-agent']['model'] = { 'primary': model_data['uuid'], 'fallbacks': [], } pipeline_data = {'config': pipeline_config} await self.ap.pipeline_service.update_pipeline(context, pipeline.uuid, pipeline_data) return model_data['uuid'] async def get_llm_model( self, context: TenantContext, model_uuid: str, include_secret: bool = False, ) -> dict | None: """Get a single LLM model with provider info""" result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.LLMModel).where(persistence_model.LLMModel.uuid == model_uuid), persistence_model.LLMModel, context, ) ) model = result.first() if model is None: return None model_dict = self.ap.persistence_mgr.serialize_model(persistence_model.LLMModel, model) # Get provider provider_result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.ModelProvider).where( persistence_model.ModelProvider.uuid == model.provider_uuid ), persistence_model.ModelProvider, context, ) ) provider = provider_result.first() if provider: provider_dict = self.ap.persistence_mgr.serialize_model(persistence_model.ModelProvider, provider) provider_dict = _parse_provider_api_keys(provider_dict) model_dict['provider'] = provider_dict if not include_secret: model_dict = _redact_model_secrets(model_dict) return model_dict async def update_llm_model( self, context: TenantContext, model_uuid: str, model_data: dict, ) -> None: """Update an existing LLM model""" existing_model = await self.get_llm_model(context, model_uuid, include_secret=True) if existing_model is None: raise WorkspaceNotFoundError('Model not found') await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) model_data = model_data.copy() model_data.pop('uuid', None) model_data.pop('workspace_uuid', None) if 'extra_args' in model_data: model_data['extra_args'] = restore_secret_placeholders( model_data['extra_args'], existing_model.get('extra_args', {}), ) # Handle provider update if needed if 'provider' in model_data: provider_data = model_data.pop('provider') if provider_data.get('uuid'): model_data['provider_uuid'] = provider_data['uuid'] else: provider_uuid = await self.ap.provider_service.find_or_create_provider( context, requester=provider_data.get('requester', ''), base_url=provider_data.get('base_url', ''), api_keys=provider_data.get('api_keys', []), ) model_data['provider_uuid'] = provider_uuid provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid']) await _require_workspace_provider(self.ap, context, provider_uuid) await _assert_cloud_managed_provider_mutable(self.ap, context, provider_uuid) await _validate_provider_supports(self.ap, context, provider_uuid, 'llm') result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.update(persistence_model.LLMModel) .where(persistence_model.LLMModel.uuid == model_uuid) .values(**model_data), persistence_model.LLMModel, context, ) ) if getattr(result, 'rowcount', None) == 0: raise WorkspaceNotFoundError('Model not found') await self.ap.model_mgr.remove_llm_model(context, model_uuid) runtime_provider = await _require_runtime_provider(self.ap, context, provider_uuid) runtime_llm_model = await self.ap.model_mgr.load_llm_model_with_provider( context, persistence_model.LLMModel( **_runtime_model_data( model_uuid, { key: value for key, value in {**existing_model, **model_data, 'provider_uuid': provider_uuid}.items() if key not in {'provider', 'created_at', 'updated_at'} }, ) ), runtime_provider, ) await self.ap.model_mgr.cache_llm_model(context, runtime_llm_model) async def delete_llm_model(self, context: TenantContext, model_uuid: str) -> None: """Delete an LLM model""" if _is_cloud_runtime(self.ap): existing_model = await self.get_llm_model(context, model_uuid, include_secret=True) if existing_model is None: raise WorkspaceNotFoundError('Model not found') await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.delete(persistence_model.LLMModel).where(persistence_model.LLMModel.uuid == model_uuid), persistence_model.LLMModel, context, ) ) if getattr(result, 'rowcount', None) == 0: raise WorkspaceNotFoundError('Model not found') await self.ap.model_mgr.remove_llm_model(context, model_uuid) async def test_llm_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None: """Test an LLM model""" require_workspace_uuid(context) runtime_llm_model: model_requester.RuntimeLLMModel | None = None if model_uuid != '_': if await self.get_llm_model(context, model_uuid) is None: raise WorkspaceNotFoundError('Model not found') runtime_llm_model = await self.ap.model_mgr.get_model_by_uuid(context, model_uuid) else: runtime_llm_model = await self.ap.model_mgr.init_temporary_runtime_llm_model(context, model_data) extra_args = model_data.get('extra_args', {}) await runtime_llm_model.provider.invoke_llm( query=None, model=runtime_llm_model, messages=[provider_message.Message(role='user', content='Hello, world! Please just reply a "Hello".')], funcs=[], extra_args=extra_args, execution_context=runtime_llm_model.execution_context, ) class EmbeddingModelsService: ap: app.Application def __init__(self, ap: app.Application) -> None: self.ap = ap async def get_embedding_models(self, context: TenantContext, include_secret: bool = False) -> list[dict]: """Get all embedding models with provider info""" result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.EmbeddingModel), persistence_model.EmbeddingModel, context ) ) models = result.all() providers_result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.ModelProvider), persistence_model.ModelProvider, context ) ) providers = {p.uuid: p for p in providers_result.all()} models_list = [] for model in models: model_dict = self.ap.persistence_mgr.serialize_model(persistence_model.EmbeddingModel, model) provider = providers.get(model.provider_uuid) if provider: provider_dict = self.ap.persistence_mgr.serialize_model(persistence_model.ModelProvider, provider) provider_dict = _parse_provider_api_keys(provider_dict) model_dict['provider'] = provider_dict if not include_secret: model_dict = _redact_model_secrets(model_dict) models_list.append(model_dict) return models_list async def get_embedding_models_by_provider( self, context: TenantContext, provider_uuid: str, *, include_secret: bool = False, ) -> list[dict]: """Get embedding models by provider UUID""" await _require_workspace_provider(self.ap, context, provider_uuid) result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.EmbeddingModel).where( persistence_model.EmbeddingModel.provider_uuid == provider_uuid ), persistence_model.EmbeddingModel, context, ) ) models = result.all() serialized = [self.ap.persistence_mgr.serialize_model(persistence_model.EmbeddingModel, m) for m in models] return serialized if include_secret else [_redact_model_secrets(model) for model in serialized] async def create_embedding_model( self, context: TenantContext, model_data: dict, preserve_uuid: bool = False, ) -> str: """Create a new embedding model""" model_data = model_data.copy() if not preserve_uuid: model_data['uuid'] = str(uuid.uuid4()) model_data['workspace_uuid'] = require_workspace_uuid(context) if 'extra_args' in model_data: model_data['extra_args'] = restore_secret_placeholders(model_data['extra_args']) if 'provider' in model_data: provider_data = model_data.pop('provider') if provider_data.get('uuid'): model_data['provider_uuid'] = provider_data['uuid'] else: provider_uuid = await self.ap.provider_service.find_or_create_provider( context, requester=provider_data.get('requester', ''), base_url=provider_data.get('base_url', ''), api_keys=provider_data.get('api_keys', []), ) model_data['provider_uuid'] = provider_uuid await _require_workspace_provider(self.ap, context, model_data['provider_uuid']) await _assert_cloud_managed_provider_mutable(self.ap, context, model_data['provider_uuid']) await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'text-embedding') await self.ap.persistence_mgr.execute_async( sqlalchemy.insert(persistence_model.EmbeddingModel).values(**model_data) ) runtime_provider = await _require_runtime_provider(self.ap, context, model_data['provider_uuid']) runtime_embedding_model = await self.ap.model_mgr.load_embedding_model_with_provider( context, persistence_model.EmbeddingModel(**model_data), runtime_provider, ) await self.ap.model_mgr.cache_embedding_model(context, runtime_embedding_model) return model_data['uuid'] async def get_embedding_model( self, context: TenantContext, model_uuid: str, include_secret: bool = False, ) -> dict | None: """Get a single embedding model with provider info""" result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.EmbeddingModel).where( persistence_model.EmbeddingModel.uuid == model_uuid ), persistence_model.EmbeddingModel, context, ) ) model = result.first() if model is None: return None model_dict = self.ap.persistence_mgr.serialize_model(persistence_model.EmbeddingModel, model) provider_result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.ModelProvider).where( persistence_model.ModelProvider.uuid == model.provider_uuid ), persistence_model.ModelProvider, context, ) ) provider = provider_result.first() if provider: provider_dict = self.ap.persistence_mgr.serialize_model(persistence_model.ModelProvider, provider) provider_dict = _parse_provider_api_keys(provider_dict) model_dict['provider'] = provider_dict if not include_secret: model_dict = _redact_model_secrets(model_dict) return model_dict async def update_embedding_model( self, context: TenantContext, model_uuid: str, model_data: dict, ) -> None: """Update an existing embedding model""" existing_model = await self.get_embedding_model(context, model_uuid, include_secret=True) if existing_model is None: raise WorkspaceNotFoundError('Model not found') await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) model_data = model_data.copy() model_data.pop('uuid', None) model_data.pop('workspace_uuid', None) if 'extra_args' in model_data: model_data['extra_args'] = restore_secret_placeholders( model_data['extra_args'], existing_model.get('extra_args', {}), ) if 'provider' in model_data: provider_data = model_data.pop('provider') if provider_data.get('uuid'): model_data['provider_uuid'] = provider_data['uuid'] else: provider_uuid = await self.ap.provider_service.find_or_create_provider( context, requester=provider_data.get('requester', ''), base_url=provider_data.get('base_url', ''), api_keys=provider_data.get('api_keys', []), ) model_data['provider_uuid'] = provider_uuid provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid']) await _require_workspace_provider(self.ap, context, provider_uuid) await _assert_cloud_managed_provider_mutable(self.ap, context, provider_uuid) await _validate_provider_supports(self.ap, context, provider_uuid, 'text-embedding') result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.update(persistence_model.EmbeddingModel) .where(persistence_model.EmbeddingModel.uuid == model_uuid) .values(**model_data), persistence_model.EmbeddingModel, context, ) ) if getattr(result, 'rowcount', None) == 0: raise WorkspaceNotFoundError('Model not found') await self.ap.model_mgr.remove_embedding_model(context, model_uuid) runtime_provider = await _require_runtime_provider(self.ap, context, provider_uuid) runtime_embedding_model = await self.ap.model_mgr.load_embedding_model_with_provider( context, persistence_model.EmbeddingModel( **_runtime_model_data( model_uuid, { key: value for key, value in {**existing_model, **model_data, 'provider_uuid': provider_uuid}.items() if key not in {'provider', 'created_at', 'updated_at'} }, ) ), runtime_provider, ) await self.ap.model_mgr.cache_embedding_model(context, runtime_embedding_model) async def delete_embedding_model(self, context: TenantContext, model_uuid: str) -> None: """Delete an embedding model""" if _is_cloud_runtime(self.ap): existing_model = await self.get_embedding_model(context, model_uuid, include_secret=True) if existing_model is None: raise WorkspaceNotFoundError('Model not found') await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.delete(persistence_model.EmbeddingModel).where( persistence_model.EmbeddingModel.uuid == model_uuid ), persistence_model.EmbeddingModel, context, ) ) if getattr(result, 'rowcount', None) == 0: raise WorkspaceNotFoundError('Model not found') await self.ap.model_mgr.remove_embedding_model(context, model_uuid) async def test_embedding_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None: """Test an embedding model""" require_workspace_uuid(context) runtime_embedding_model: model_requester.RuntimeEmbeddingModel | None = None if model_uuid != '_': if await self.get_embedding_model(context, model_uuid) is None: raise WorkspaceNotFoundError('Model not found') runtime_embedding_model = await self.ap.model_mgr.get_embedding_model_by_uuid(context, model_uuid) else: runtime_embedding_model = await self.ap.model_mgr.init_temporary_runtime_embedding_model( context, model_data, ) await runtime_embedding_model.provider.invoke_embedding( model=runtime_embedding_model, input_text=['Hello, world!'], extra_args={}, execution_context=runtime_embedding_model.execution_context, ) class RerankModelsService: ap: app.Application def __init__(self, ap: app.Application) -> None: self.ap = ap async def get_rerank_models(self, context: TenantContext, include_secret: bool = False) -> list[dict]: """Get all rerank models with provider info""" result = await self.ap.persistence_mgr.execute_async( scope_statement(sqlalchemy.select(persistence_model.RerankModel), persistence_model.RerankModel, context) ) models = result.all() providers_result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.ModelProvider), persistence_model.ModelProvider, context ) ) providers = {p.uuid: p for p in providers_result.all()} models_list = [] for model in models: model_dict = self.ap.persistence_mgr.serialize_model(persistence_model.RerankModel, model) provider = providers.get(model.provider_uuid) if provider: provider_dict = self.ap.persistence_mgr.serialize_model(persistence_model.ModelProvider, provider) provider_dict = _parse_provider_api_keys(provider_dict) model_dict['provider'] = provider_dict if not include_secret: model_dict = _redact_model_secrets(model_dict) models_list.append(model_dict) return models_list async def get_rerank_models_by_provider( self, context: TenantContext, provider_uuid: str, *, include_secret: bool = False, ) -> list[dict]: """Get rerank models by provider UUID""" await _require_workspace_provider(self.ap, context, provider_uuid) result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.RerankModel).where( persistence_model.RerankModel.provider_uuid == provider_uuid ), persistence_model.RerankModel, context, ) ) models = result.all() serialized = [self.ap.persistence_mgr.serialize_model(persistence_model.RerankModel, m) for m in models] return serialized if include_secret else [_redact_model_secrets(model) for model in serialized] async def create_rerank_model( self, context: TenantContext, model_data: dict, preserve_uuid: bool = False, ) -> str: """Create a new rerank model""" model_data = model_data.copy() if not preserve_uuid: model_data['uuid'] = str(uuid.uuid4()) model_data['workspace_uuid'] = require_workspace_uuid(context) if 'extra_args' in model_data: model_data['extra_args'] = restore_secret_placeholders(model_data['extra_args']) if 'provider' in model_data: provider_data = model_data.pop('provider') if provider_data.get('uuid'): model_data['provider_uuid'] = provider_data['uuid'] else: provider_uuid = await self.ap.provider_service.find_or_create_provider( context, requester=provider_data.get('requester', ''), base_url=provider_data.get('base_url', ''), api_keys=provider_data.get('api_keys', []), ) model_data['provider_uuid'] = provider_uuid await _require_workspace_provider(self.ap, context, model_data['provider_uuid']) await _assert_cloud_managed_provider_mutable(self.ap, context, model_data['provider_uuid']) await _validate_provider_supports(self.ap, context, model_data['provider_uuid'], 'rerank') await self.ap.persistence_mgr.execute_async( sqlalchemy.insert(persistence_model.RerankModel).values(**model_data) ) runtime_provider = await _require_runtime_provider(self.ap, context, model_data['provider_uuid']) runtime_rerank_model = await self.ap.model_mgr.load_rerank_model_with_provider( context, persistence_model.RerankModel(**model_data), runtime_provider, ) await self.ap.model_mgr.cache_rerank_model(context, runtime_rerank_model) return model_data['uuid'] async def get_rerank_model( self, context: TenantContext, model_uuid: str, include_secret: bool = False, ) -> dict | None: """Get a single rerank model with provider info""" result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.RerankModel).where( persistence_model.RerankModel.uuid == model_uuid ), persistence_model.RerankModel, context, ) ) model = result.first() if model is None: return None model_dict = self.ap.persistence_mgr.serialize_model(persistence_model.RerankModel, model) provider_result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.select(persistence_model.ModelProvider).where( persistence_model.ModelProvider.uuid == model.provider_uuid ), persistence_model.ModelProvider, context, ) ) provider = provider_result.first() if provider: provider_dict = self.ap.persistence_mgr.serialize_model(persistence_model.ModelProvider, provider) provider_dict = _parse_provider_api_keys(provider_dict) model_dict['provider'] = provider_dict if not include_secret: model_dict = _redact_model_secrets(model_dict) return model_dict async def update_rerank_model( self, context: TenantContext, model_uuid: str, model_data: dict, ) -> None: """Update an existing rerank model""" existing_model = await self.get_rerank_model(context, model_uuid, include_secret=True) if existing_model is None: raise WorkspaceNotFoundError('Model not found') await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) model_data = model_data.copy() model_data.pop('uuid', None) model_data.pop('workspace_uuid', None) if 'extra_args' in model_data: model_data['extra_args'] = restore_secret_placeholders( model_data['extra_args'], existing_model.get('extra_args', {}), ) if 'provider' in model_data: provider_data = model_data.pop('provider') if provider_data.get('uuid'): model_data['provider_uuid'] = provider_data['uuid'] else: provider_uuid = await self.ap.provider_service.find_or_create_provider( context, requester=provider_data.get('requester', ''), base_url=provider_data.get('base_url', ''), api_keys=provider_data.get('api_keys', []), ) model_data['provider_uuid'] = provider_uuid provider_uuid = model_data.get('provider_uuid', existing_model['provider_uuid']) await _require_workspace_provider(self.ap, context, provider_uuid) await _assert_cloud_managed_provider_mutable(self.ap, context, provider_uuid) await _validate_provider_supports(self.ap, context, provider_uuid, 'rerank') result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.update(persistence_model.RerankModel) .where(persistence_model.RerankModel.uuid == model_uuid) .values(**model_data), persistence_model.RerankModel, context, ) ) if getattr(result, 'rowcount', None) == 0: raise WorkspaceNotFoundError('Model not found') await self.ap.model_mgr.remove_rerank_model(context, model_uuid) runtime_provider = await _require_runtime_provider(self.ap, context, provider_uuid) runtime_rerank_model = await self.ap.model_mgr.load_rerank_model_with_provider( context, persistence_model.RerankModel( **_runtime_model_data( model_uuid, { key: value for key, value in {**existing_model, **model_data, 'provider_uuid': provider_uuid}.items() if key not in {'provider', 'created_at', 'updated_at'} }, ) ), runtime_provider, ) await self.ap.model_mgr.cache_rerank_model(context, runtime_rerank_model) async def delete_rerank_model(self, context: TenantContext, model_uuid: str) -> None: """Delete a rerank model""" if _is_cloud_runtime(self.ap): existing_model = await self.get_rerank_model(context, model_uuid, include_secret=True) if existing_model is None: raise WorkspaceNotFoundError('Model not found') await _assert_cloud_managed_provider_mutable(self.ap, context, existing_model['provider_uuid']) result = await self.ap.persistence_mgr.execute_async( scope_statement( sqlalchemy.delete(persistence_model.RerankModel).where( persistence_model.RerankModel.uuid == model_uuid ), persistence_model.RerankModel, context, ) ) if getattr(result, 'rowcount', None) == 0: raise WorkspaceNotFoundError('Model not found') await self.ap.model_mgr.remove_rerank_model(context, model_uuid) async def test_rerank_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None: """Test a rerank model""" require_workspace_uuid(context) runtime_rerank_model: model_requester.RuntimeRerankModel | None = None if model_uuid != '_': if await self.get_rerank_model(context, model_uuid) is None: raise WorkspaceNotFoundError('Model not found') runtime_rerank_model = await self.ap.model_mgr.get_rerank_model_by_uuid(context, model_uuid) else: runtime_rerank_model = await self.ap.model_mgr.init_temporary_runtime_rerank_model( context, model_data, ) await runtime_rerank_model.provider.invoke_rerank( model=runtime_rerank_model, query='What is artificial intelligence?', documents=[ 'Artificial intelligence is a branch of computer science.', 'The weather is nice today.', ], execution_context=runtime_rerank_model.execution_context, )