From 64725ce33b4c125a07d55e3da73fe0fa9a77ac0b Mon Sep 17 00:00:00 2001 From: dadachann <185672915+dadachann@users.noreply.github.com> Date: Sat, 1 Aug 2026 07:40:42 +0000 Subject: [PATCH] feat(cloud): provision workspace model catalog --- src/langbot/pkg/api/http/service/model.py | 75 ++++- src/langbot/pkg/api/http/service/provider.py | 23 ++ src/langbot/pkg/cloud/bootstrap.py | 5 + src/langbot/pkg/cloud/model_catalog.py | 284 ++++++++++++++++++ src/langbot/pkg/core/app.py | 12 +- src/langbot/pkg/core/stages/build_app.py | 11 + .../test_cloud_managed_model_protection.py | 112 +++++++ .../api/service/test_provider_service.py | 54 ++++ tests/unit_tests/cloud/test_bootstrap.py | 5 + tests/unit_tests/cloud/test_model_catalog.py | 271 +++++++++++++++++ 10 files changed, 844 insertions(+), 8 deletions(-) create mode 100644 src/langbot/pkg/cloud/model_catalog.py create mode 100644 tests/unit_tests/api/service/test_cloud_managed_model_protection.py create mode 100644 tests/unit_tests/cloud/test_model_catalog.py diff --git a/src/langbot/pkg/api/http/service/model.py b/src/langbot/pkg/api/http/service/model.py index 88801c6b8..898a791e2 100644 --- a/src/langbot/pkg/api/http/service/model.py +++ b/src/langbot/pkg/api/http/service/model.py @@ -5,6 +5,7 @@ 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 @@ -113,6 +114,23 @@ async def _require_workspace_provider( 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, @@ -213,6 +231,7 @@ class LLMModelsService: 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)) @@ -291,11 +310,17 @@ class LLMModelsService: return model_dict - async def update_llm_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None: + 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) @@ -321,6 +346,7 @@ class LLMModelsService: 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( @@ -355,6 +381,11 @@ class LLMModelsService: 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), @@ -448,7 +479,10 @@ class EmbeddingModelsService: 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 + self, + context: TenantContext, + model_data: dict, + preserve_uuid: bool = False, ) -> str: """Create a new embedding model""" model_data = model_data.copy() @@ -472,6 +506,7 @@ class EmbeddingModelsService: 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( @@ -530,11 +565,17 @@ class EmbeddingModelsService: return model_dict - async def update_embedding_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None: + 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) @@ -559,6 +600,7 @@ class EmbeddingModelsService: 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( @@ -593,6 +635,11 @@ class EmbeddingModelsService: 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( @@ -685,7 +732,12 @@ class RerankModelsService: 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: + 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: @@ -708,6 +760,7 @@ class RerankModelsService: 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( @@ -766,11 +819,17 @@ class RerankModelsService: return model_dict - async def update_rerank_model(self, context: TenantContext, model_uuid: str, model_data: dict) -> None: + 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) @@ -795,6 +854,7 @@ class RerankModelsService: 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( @@ -829,6 +889,11 @@ class RerankModelsService: 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( diff --git a/src/langbot/pkg/api/http/service/provider.py b/src/langbot/pkg/api/http/service/provider.py index 74647d7ab..dc29d3858 100644 --- a/src/langbot/pkg/api/http/service/provider.py +++ b/src/langbot/pkg/api/http/service/provider.py @@ -5,6 +5,7 @@ import traceback import sqlalchemy +from ....cloud.model_catalog import LANGBOT_MODELS_PROVIDER_REQUESTER from ....core import app from ....entity.persistence import model as persistence_model from ....workspace.errors import WorkspaceNotFoundError @@ -20,6 +21,20 @@ class ModelProviderService: def __init__(self, ap: app.Application) -> None: self.ap = ap + def _is_cloud_runtime(self) -> bool: + mode = getattr(self.ap.persistence_mgr, 'mode', None) + return getattr(mode, 'value', None) == 'cloud_runtime' + + def _system_requester_is_reserved(self, requester: object) -> bool: + return self._is_cloud_runtime() and requester == LANGBOT_MODELS_PROVIDER_REQUESTER + + async def _assert_provider_mutable(self, context: TenantContext, provider_uuid: str) -> None: + if not self._is_cloud_runtime(): + return + provider = await self.get_provider(context, provider_uuid) + if provider is not None and self._system_requester_is_reserved(provider.get('requester')): + raise ValueError('LangBot Models is managed by Cloud and cannot be modified') + @staticmethod def _normalize_api_keys(api_keys: str | list[str] | tuple[str, ...] | None) -> list[str]: if api_keys is None: @@ -99,6 +114,8 @@ class ModelProviderService: async def create_provider(self, context: TenantContext, provider_data: dict) -> str: """Create a new provider""" provider_data = provider_data.copy() + if self._system_requester_is_reserved(provider_data.get('requester')): + raise ValueError('space-chat-completions is reserved for the Cloud-managed LangBot Models provider') provider_data['uuid'] = str(uuid.uuid4()) provider_data['workspace_uuid'] = require_workspace_uuid(context) provider_data['api_keys'] = self._normalize_api_keys( @@ -115,7 +132,10 @@ class ModelProviderService: async def update_provider(self, context: TenantContext, provider_uuid: str, provider_data: dict) -> None: """Update an existing provider""" + await self._assert_provider_mutable(context, provider_uuid) provider_data = provider_data.copy() + if self._system_requester_is_reserved(provider_data.get('requester')): + raise ValueError('space-chat-completions is reserved for the Cloud-managed LangBot Models provider') provider_data.pop('uuid', None) provider_data.pop('workspace_uuid', None) if 'api_keys' in provider_data: @@ -145,6 +165,7 @@ class ModelProviderService: async def delete_provider(self, context: TenantContext, provider_uuid: str) -> None: """Delete a provider (only if no models reference it)""" + await self._assert_provider_mutable(context, provider_uuid) workspace_uuid = require_workspace_uuid(context) # Check if any models use this provider llm_result = await self.ap.persistence_mgr.execute_async( @@ -245,6 +266,8 @@ class ModelProviderService: api_keys: list, ) -> str: """Find existing provider or create new one""" + if self._system_requester_is_reserved(requester): + raise ValueError('space-chat-completions is reserved for the Cloud-managed LangBot Models provider') workspace_uuid = require_workspace_uuid(context) api_keys = self._normalize_api_keys(restore_secret_placeholders(api_keys, sensitive=True)) diff --git a/src/langbot/pkg/cloud/bootstrap.py b/src/langbot/pkg/cloud/bootstrap.py index c8341e56c..2c7ddafd2 100644 --- a/src/langbot/pkg/cloud/bootstrap.py +++ b/src/langbot/pkg/cloud/bootstrap.py @@ -13,6 +13,7 @@ from typing import Any, Protocol, runtime_checkable from ..workspace.policy import CloudWorkspacePolicy, SingleWorkspacePolicy from .directory import DirectoryProjectionProvider, directory_projection_limits_from_config from .entitlements import EntitlementProvider, OpenSourceEntitlementProvider +from .model_catalog import CloudModelCatalogProvider CLOUD_BOOTSTRAP_ENTRY_POINT = 'langbot.cloud_bootstrap' @@ -50,6 +51,7 @@ class OpenSourceDeployment: ) directory_provider: None = None manifest_provider: None = None + model_catalog_provider: None = None persistence_mode: str = 'oss_compat' required_vector_backend: str | None = None @@ -80,6 +82,7 @@ class VerifiedCloudDeployment: entitlement_provider: EntitlementProvider directory_provider: DirectoryProjectionProvider manifest_provider: CloudManifestProvider + model_catalog_provider: CloudModelCatalogProvider verification_key_id: str mode: str = dataclasses.field(default='cloud', init=False) workspace_policy: CloudWorkspacePolicy = dataclasses.field(default_factory=CloudWorkspacePolicy, init=False) @@ -110,6 +113,8 @@ class VerifiedCloudDeployment: raise CloudBootstrapError('Verified Cloud bootstrap did not provide a directory adapter') if not isinstance(self.manifest_provider, CloudManifestProvider): raise CloudBootstrapError('Verified Cloud bootstrap did not provide a Manifest renewal adapter') + if not isinstance(self.model_catalog_provider, CloudModelCatalogProvider): + raise CloudBootstrapError('Verified Cloud bootstrap did not provide a model catalog adapter') def validate_instance_config(self, config: dict[str, Any]) -> None: try: diff --git a/src/langbot/pkg/cloud/model_catalog.py b/src/langbot/pkg/cloud/model_catalog.py new file mode 100644 index 000000000..6e926ad84 --- /dev/null +++ b/src/langbot/pkg/cloud/model_catalog.py @@ -0,0 +1,284 @@ +from __future__ import annotations + +import asyncio +import uuid +from datetime import datetime +from typing import Any, Literal, Protocol, runtime_checkable + +import sqlalchemy +from pydantic import BaseModel, ConfigDict, Field, SecretStr, field_validator + +from ..entity.persistence import model as persistence_model + + +LANGBOT_MODELS_PROVIDER_REQUESTER = 'space-chat-completions' +LANGBOT_MODELS_PROVIDER_NAME = 'LangBot Models' +_MODEL_RESOURCE_NAMESPACE = uuid.UUID('94c703ca-1df5-4e91-bcd3-74ac65cb7921') +_SUPPORTED_CATEGORIES = {'chat', 'embedding', 'rerank'} +_MODEL_TABLES = ( + persistence_model.LLMModel, + persistence_model.EmbeddingModel, + persistence_model.RerankModel, +) + + +class CloudModelCatalogItem(BaseModel): + model_config = ConfigDict(extra='forbid', frozen=True) + + uuid: str = Field(min_length=1, max_length=255) + model_id: str = Field(min_length=1, max_length=255) + category: Literal['chat', 'embedding', 'rerank'] + llm_abilities: tuple[str, ...] = () + is_featured: bool = False + featured_order: int = 0 + + @field_validator('llm_abilities') + @classmethod + def validate_abilities(cls, value: tuple[str, ...]) -> tuple[str, ...]: + if any(not item.strip() or len(item) > 64 for item in value): + raise ValueError('Model abilities must be non-empty strings of at most 64 characters') + if len(set(value)) != len(value): + raise ValueError('Model abilities must be unique') + return value + + +class CloudWorkspaceModelBilling(BaseModel): + model_config = ConfigDict(extra='forbid', frozen=True) + + workspace_uuid: str = Field(min_length=36, max_length=36) + owner_account_uuid: str | None = Field(default=None, min_length=36, max_length=36) + api_key: SecretStr | None = None + + @field_validator('workspace_uuid') + @classmethod + def validate_uuid(cls, value: str) -> str: + return str(uuid.UUID(value)) + + @field_validator('owner_account_uuid') + @classmethod + def validate_optional_uuid(cls, value: str | None) -> str | None: + return None if value is None else str(uuid.UUID(value)) + + +class CloudModelCatalogSnapshot(BaseModel): + model_config = ConfigDict(extra='forbid', frozen=True) + + instance_uuid: str = Field(min_length=1, max_length=255) + generated_at: datetime + base_url: str = Field(min_length=1, max_length=512) + models: tuple[CloudModelCatalogItem, ...] + workspaces: tuple[CloudWorkspaceModelBilling, ...] + + @field_validator('base_url') + @classmethod + def validate_base_url(cls, value: str) -> str: + normalized = value.rstrip('/') + if not normalized.startswith('https://'): + raise ValueError('Cloud model gateway base URL must use HTTPS') + return normalized + + @field_validator('models') + @classmethod + def validate_models(cls, value: tuple[CloudModelCatalogItem, ...]) -> tuple[CloudModelCatalogItem, ...]: + if len(value) > 500: + raise ValueError('Cloud model catalog exceeds 500 models') + identities = {(item.category, item.uuid) for item in value} + if len(identities) != len(value): + raise ValueError('Cloud model catalog contains duplicate model identities') + return value + + @field_validator('workspaces') + @classmethod + def validate_workspaces( + cls, value: tuple[CloudWorkspaceModelBilling, ...] + ) -> tuple[CloudWorkspaceModelBilling, ...]: + if len(value) > 10_000: + raise ValueError('Cloud model catalog exceeds 10000 Workspaces') + identities = {item.workspace_uuid for item in value} + if len(identities) != len(value): + raise ValueError('Cloud model catalog contains duplicate Workspaces') + return value + + +@runtime_checkable +class CloudModelCatalogProvider(Protocol): + async def fetch_model_catalog(self, instance_uuid: str) -> CloudModelCatalogSnapshot: + """Fetch and verify the complete model catalog and Workspace billing projection.""" + ... + + +def system_provider_uuid(workspace_uuid: str) -> str: + workspace = str(uuid.UUID(workspace_uuid)) + return str(uuid.uuid5(_MODEL_RESOURCE_NAMESPACE, f'{workspace}:provider:{LANGBOT_MODELS_PROVIDER_REQUESTER}')) + + +def system_model_uuid(workspace_uuid: str, category: str, upstream_uuid: str) -> str: + workspace = str(uuid.UUID(workspace_uuid)) + if category not in _SUPPORTED_CATEGORIES: + raise ValueError(f'Unsupported model category: {category}') + if not upstream_uuid: + raise ValueError('Upstream model UUID is required') + return str(uuid.uuid5(_MODEL_RESOURCE_NAMESPACE, f'{workspace}:model:{category}:{upstream_uuid}')) + + +class CloudModelCatalogSyncService: + """Reconcile Space-owned model catalog and Owner billing tokens into every Cloud Workspace.""" + + def __init__( + self, + ap: Any, + provider: CloudModelCatalogProvider, + instance_uuid: str, + *, + sync_interval_seconds: float = 3600.0, + ) -> None: + if not isinstance(provider, CloudModelCatalogProvider): + raise TypeError('Cloud model catalog sync requires a CloudModelCatalogProvider') + if sync_interval_seconds < 10: + raise ValueError('Cloud model catalog sync interval must be at least 10 seconds') + self.ap = ap + self.provider = provider + self.instance_uuid = instance_uuid + self.sync_interval_seconds = float(sync_interval_seconds) + + async def initialize(self) -> None: + await self.sync_once(reload_runtime=False) + + async def run(self) -> None: + while True: + await asyncio.sleep(self.sync_interval_seconds) + try: + await self.sync_once(reload_runtime=True) + except asyncio.CancelledError: + raise + except Exception as exc: + self.ap.logger.warning(f'Cloud model catalog synchronization failed: {exc}') + + async def sync_once(self, *, reload_runtime: bool = True) -> dict[str, int]: + snapshot = await self.provider.fetch_model_catalog(self.instance_uuid) + if snapshot.instance_uuid != self.instance_uuid: + raise ValueError('Cloud model catalog targets another LangBot instance') + + bindings = await self.ap.workspace_service.list_active_execution_bindings() + billing_by_workspace = {item.workspace_uuid: item for item in snapshot.workspaces} + missing = sorted( + binding.workspace_uuid for binding in bindings if binding.workspace_uuid not in billing_by_workspace + ) + if missing: + raise ValueError(f'Cloud model catalog is missing billing projections for {len(missing)} active Workspaces') + + summary = {'workspaces': 0, 'created': 0, 'updated': 0, 'deleted': 0} + changed = False + for binding in bindings: + counts = await self._sync_workspace( + binding.workspace_uuid, + snapshot, + billing_by_workspace[binding.workspace_uuid], + ) + summary['workspaces'] += 1 + for key in ('created', 'updated', 'deleted'): + summary[key] += counts[key] + changed = changed or counts[key] > 0 + + if changed and reload_runtime and getattr(self.ap, 'model_mgr', None) is not None: + await self.ap.model_mgr.load_models_from_db() + if changed: + self.ap.logger.info( + 'Cloud model catalog synchronized ' + f'({summary["workspaces"]} Workspaces, {len(snapshot.models)} models, ' + f'created={summary["created"]}, updated={summary["updated"]}, deleted={summary["deleted"]})' + ) + return summary + + async def _sync_workspace( + self, + workspace_uuid: str, + snapshot: CloudModelCatalogSnapshot, + billing: CloudWorkspaceModelBilling, + ) -> dict[str, int]: + counts = {'created': 0, 'updated': 0, 'deleted': 0} + provider_uuid = system_provider_uuid(workspace_uuid) + desired_keys = [billing.api_key.get_secret_value()] if billing.api_key is not None else [] + + async with self.ap.persistence_mgr.tenant_uow(workspace_uuid) as uow: + provider = await uow.session.scalar( + sqlalchemy.select(persistence_model.ModelProvider).where( + persistence_model.ModelProvider.uuid == provider_uuid + ) + ) + provider_values = { + 'workspace_uuid': workspace_uuid, + 'name': LANGBOT_MODELS_PROVIDER_NAME, + 'requester': LANGBOT_MODELS_PROVIDER_REQUESTER, + 'base_url': snapshot.base_url, + 'api_keys': desired_keys, + } + if provider is None: + provider = persistence_model.ModelProvider(uuid=provider_uuid, **provider_values) + uow.session.add(provider) + await uow.session.flush() + counts['created'] += 1 + elif self._update_entity(provider, provider_values): + counts['updated'] += 1 + + existing_by_table: dict[type, dict[str, Any]] = {} + for table in _MODEL_TABLES: + rows = ( + await uow.session.scalars(sqlalchemy.select(table).where(table.provider_uuid == provider_uuid)) + ).all() + existing_by_table[table] = {row.uuid: row for row in rows} + + desired_ids: dict[type, set[str]] = {table: set() for table in _MODEL_TABLES} + for item in snapshot.models: + table, values = self._model_values(workspace_uuid, provider_uuid, item) + model_uuid = system_model_uuid(workspace_uuid, item.category, item.uuid) + desired_ids[table].add(model_uuid) + existing = existing_by_table[table].get(model_uuid) + if existing is None: + uow.session.add(table(uuid=model_uuid, **values)) + counts['created'] += 1 + elif self._update_entity(existing, values): + counts['updated'] += 1 + + for table, entities in existing_by_table.items(): + for model_uuid, entity in entities.items(): + if model_uuid not in desired_ids[table]: + await uow.session.delete(entity) + counts['deleted'] += 1 + + return counts + + @staticmethod + def _update_entity(entity: Any, values: dict[str, Any]) -> bool: + changed = False + for key, value in values.items(): + if getattr(entity, key) != value: + setattr(entity, key, value) + changed = True + return changed + + @staticmethod + def _model_values( + workspace_uuid: str, + provider_uuid: str, + item: CloudModelCatalogItem, + ) -> tuple[type, dict[str, Any]]: + ranking = 100 - item.featured_order if item.is_featured else 0 + common = { + 'workspace_uuid': workspace_uuid, + 'name': item.model_id, + 'provider_uuid': provider_uuid, + 'extra_args': {}, + 'prefered_ranking': ranking, + } + if item.category == 'chat': + return persistence_model.LLMModel, { + **common, + 'abilities': list(item.llm_abilities), + 'context_length': None, + } + if item.category == 'embedding': + return persistence_model.EmbeddingModel, common + if item.category == 'rerank': + return persistence_model.RerankModel, common + raise ValueError(f'Unsupported model category: {item.category}') diff --git a/src/langbot/pkg/core/app.py b/src/langbot/pkg/core/app.py index bfd01387d..3f4a17329 100644 --- a/src/langbot/pkg/core/app.py +++ b/src/langbot/pkg/core/app.py @@ -54,6 +54,7 @@ from ..cloud import launch as cloud_launch_module from ..cloud import support_admin as cloud_support_admin_module from ..cloud import directory_projection as cloud_directory_projection_module from ..cloud import entitlements as cloud_entitlements_module +from ..cloud import model_catalog as cloud_model_catalog_module from ..api.http.context import ExecutionContext, PrincipalContext, PrincipalType @@ -142,13 +143,12 @@ class Application: deployment: cloud_bootstrap_module.OpenSourceDeployment | cloud_bootstrap_module.VerifiedCloudDeployment = None deployment_admission: cloud_bootstrap_module.DeploymentAdmissionGuard = None - + directory_projection_service: cloud_directory_projection_module.DirectoryProjectionService | None = None + cloud_model_catalog_service: cloud_model_catalog_module.CloudModelCatalogSyncService | None = None manifest_refresh_service: cloud_bootstrap_module.CloudManifestRefreshService | None = None entitlement_resolver: cloud_entitlements_module.EntitlementResolver | None = None - directory_projection_service: cloud_directory_projection_module.DirectoryProjectionService | None = None - vector_db_mgr: vectordb_mgr.VectorDBManager = None http_ctrl: http_controller.HTTPController = None @@ -306,6 +306,12 @@ class Application: name='cloud-directory-projection', scopes=[core_entities.LifecycleControlScope.APPLICATION], ) + if self.cloud_model_catalog_service is not None: + self.task_mgr.create_task( + self.cloud_model_catalog_service.run(), + name='cloud-model-catalog-sync', + scopes=[core_entities.LifecycleControlScope.APPLICATION], + ) if self.manifest_refresh_service is not None: self.task_mgr.create_task( self.manifest_refresh_service.run(), diff --git a/src/langbot/pkg/core/stages/build_app.py b/src/langbot/pkg/core/stages/build_app.py index b08fdfaf7..14d533dcc 100644 --- a/src/langbot/pkg/core/stages/build_app.py +++ b/src/langbot/pkg/core/stages/build_app.py @@ -46,6 +46,7 @@ from ...cloud import support_admin as cloud_support_admin_module from ...cloud.directory import directory_projection_limits_from_config from ...cloud.directory_projection import DirectoryProjectionService from ...cloud.entitlements import EntitlementResolver +from ...cloud.model_catalog import CloudModelCatalogSyncService from ...api.http.context import ExecutionContext, PrincipalContext, PrincipalType from ...api.http.authz import WorkspaceRequiredError @@ -176,6 +177,16 @@ class BuildAppStage(stage.BootingStage): # of repeating tenant validation for every manager. await workspace_service_inst.prime_startup_execution_bindings() + if not isinstance(deployment, cloud_bootstrap.VerifiedCloudDeployment): + raise RuntimeError('Multi-Workspace runtime requires a verified Cloud deployment') + cloud_model_catalog_service = CloudModelCatalogSyncService( + ap, + deployment.model_catalog_provider, + constants.instance_id, + ) + await cloud_model_catalog_service.initialize() + ap.cloud_model_catalog_service = cloud_model_catalog_service + ap.workspace_collaboration_service = workspace_collaboration_module.WorkspaceCollaborationService( ap, workspace_service_inst, diff --git a/tests/unit_tests/api/service/test_cloud_managed_model_protection.py b/tests/unit_tests/api/service/test_cloud_managed_model_protection.py new file mode 100644 index 000000000..62ed06049 --- /dev/null +++ b/tests/unit_tests/api/service/test_cloud_managed_model_protection.py @@ -0,0 +1,112 @@ +"""Cloud Runtime write protection for the managed LangBot Models catalog.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from langbot.pkg.api.http.service import model as model_service_module +from langbot.pkg.api.http.service.model import ( + EmbeddingModelsService, + LLMModelsService, + RerankModelsService, + _assert_cloud_managed_provider_mutable, +) +from langbot.pkg.cloud.model_catalog import LANGBOT_MODELS_PROVIDER_REQUESTER + + +WORKSPACE = 'workspace-a' +PROVIDER = 'managed-provider' +MODEL = 'managed-model' + + +@pytest.mark.asyncio +async def test_managed_provider_guard_is_cloud_only(monkeypatch) -> None: + async def managed_provider(_ap, _context, provider_uuid): + assert provider_uuid == PROVIDER + return {'uuid': PROVIDER, 'requester': LANGBOT_MODELS_PROVIDER_REQUESTER} + + monkeypatch.setattr(model_service_module, '_require_workspace_provider', managed_provider) + application = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='cloud_runtime'))) + + with pytest.raises(ValueError, match='managed by Cloud'): + await _assert_cloud_managed_provider_mutable( + application, + WORKSPACE, + PROVIDER, + ) + + application.persistence_mgr.mode.value = 'normal' + await _assert_cloud_managed_provider_mutable( + application, + WORKSPACE, + PROVIDER, + ) + + +@pytest.mark.parametrize( + ('service_type', 'create_method', 'model_data'), + [ + (LLMModelsService, 'create_llm_model', {'provider_uuid': PROVIDER, 'name': 'chat', 'abilities': []}), + (EmbeddingModelsService, 'create_embedding_model', {'provider_uuid': PROVIDER, 'name': 'embedding'}), + (RerankModelsService, 'create_rerank_model', {'provider_uuid': PROVIDER, 'name': 'rerank'}), + ], +) +@pytest.mark.asyncio +async def test_all_model_types_reject_creation_under_managed_provider( + monkeypatch, + service_type, + create_method: str, + model_data: dict, +) -> None: + guard = AsyncMock(side_effect=ValueError('LangBot Models is managed by Cloud and cannot be modified')) + monkeypatch.setattr(model_service_module, '_assert_cloud_managed_provider_mutable', guard) + application = SimpleNamespace( + persistence_mgr=SimpleNamespace(), + provider_service=SimpleNamespace( + get_provider=AsyncMock(return_value={'uuid': PROVIDER, 'requester': LANGBOT_MODELS_PROVIDER_REQUESTER}) + ), + model_mgr=None, + ) + service = service_type(application) + + with pytest.raises(ValueError, match='managed by Cloud'): + await getattr(service, create_method)(WORKSPACE, model_data) + + guard.assert_awaited_once() + + +@pytest.mark.parametrize( + ('service_type', 'get_method', 'write_method', 'payload'), + [ + (LLMModelsService, 'get_llm_model', 'update_llm_model', {'name': 'changed'}), + (LLMModelsService, 'get_llm_model', 'delete_llm_model', None), + (EmbeddingModelsService, 'get_embedding_model', 'update_embedding_model', {'name': 'changed'}), + (EmbeddingModelsService, 'get_embedding_model', 'delete_embedding_model', None), + (RerankModelsService, 'get_rerank_model', 'update_rerank_model', {'name': 'changed'}), + (RerankModelsService, 'get_rerank_model', 'delete_rerank_model', None), + ], +) +@pytest.mark.asyncio +async def test_all_model_types_reject_update_and_delete_for_managed_provider( + monkeypatch, + service_type, + get_method: str, + write_method: str, + payload: dict | None, +) -> None: + guard = AsyncMock(side_effect=ValueError('LangBot Models is managed by Cloud and cannot be modified')) + monkeypatch.setattr(model_service_module, '_assert_cloud_managed_provider_mutable', guard) + application = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='cloud_runtime'))) + service = service_type(application) + monkeypatch.setattr( + service, + get_method, + AsyncMock(return_value={'uuid': MODEL, 'provider_uuid': PROVIDER, 'extra_args': {}}), + ) + + args = (WORKSPACE, MODEL) if payload is None else (WORKSPACE, MODEL, payload) + with pytest.raises(ValueError, match='managed by Cloud'): + await getattr(service, write_method)(*args) + + guard.assert_awaited_once() diff --git a/tests/unit_tests/api/service/test_provider_service.py b/tests/unit_tests/api/service/test_provider_service.py index 15b995895..fdf8b01bc 100644 --- a/tests/unit_tests/api/service/test_provider_service.py +++ b/tests/unit_tests/api/service/test_provider_service.py @@ -25,6 +25,7 @@ from langbot.pkg.workspace.errors import WorkspaceNotFoundError pytestmark = pytest.mark.asyncio WORKSPACE_UUID = 'workspace-a' +SYSTEM_REQUESTER = 'space-chat-completions' def _create_mock_provider( @@ -1005,3 +1006,56 @@ class TestProviderSecretRoundtrip: ) ap.persistence_mgr.execute_async.assert_not_awaited() + + +class TestCloudManagedProviderProtection: + @staticmethod + def _service() -> ModelProviderService: + ap = SimpleNamespace( + persistence_mgr=SimpleNamespace( + mode=SimpleNamespace(value='cloud_runtime'), + execute_async=AsyncMock(), + ), + model_mgr=SimpleNamespace(), + ) + return ModelProviderService(ap) + + async def test_cloud_rejects_user_created_system_requester(self): + service = self._service() + + with pytest.raises(ValueError, match='reserved'): + await service.create_provider( + WORKSPACE_UUID, + { + 'name': 'Fake LangBot Models', + 'requester': SYSTEM_REQUESTER, + 'base_url': 'https://example.invalid/v1', + 'api_keys': ['fake'], + }, + ) + + with pytest.raises(ValueError, match='reserved'): + await service.find_or_create_provider( + WORKSPACE_UUID, + SYSTEM_REQUESTER, + 'https://api.langbot.cloud/v1', + ['fake'], + ) + service.ap.persistence_mgr.execute_async.assert_not_awaited() + + async def test_cloud_rejects_update_and_delete_of_managed_provider(self): + service = self._service() + service.get_provider = AsyncMock( + return_value={'uuid': 'system-provider', 'requester': SYSTEM_REQUESTER} + ) + + with pytest.raises(ValueError, match='managed by Cloud'): + await service.update_provider(WORKSPACE_UUID, 'system-provider', {'name': 'Renamed'}) + with pytest.raises(ValueError, match='managed by Cloud'): + await service.delete_provider(WORKSPACE_UUID, 'system-provider') + service.ap.persistence_mgr.execute_async.assert_not_awaited() + + async def test_oss_does_not_reserve_space_requester(self): + ap = SimpleNamespace(persistence_mgr=SimpleNamespace(mode=SimpleNamespace(value='oss_compat'))) + service = ModelProviderService(ap) + assert service._system_requester_is_reserved(SYSTEM_REQUESTER) is False diff --git a/tests/unit_tests/cloud/test_bootstrap.py b/tests/unit_tests/cloud/test_bootstrap.py index e27367d8d..6225ba7b7 100644 --- a/tests/unit_tests/cloud/test_bootstrap.py +++ b/tests/unit_tests/cloud/test_bootstrap.py @@ -66,6 +66,10 @@ class _Provider: def __init__(self): self.manifest_provider = _Manifest() + async def fetch_model_catalog(self, instance_uuid: str): + del instance_uuid + raise AssertionError('not used by bootstrap contract tests') + def bootstrap(self, *, instance_uuid: str, instance_config: dict): del instance_config return VerifiedCloudDeployment( @@ -79,6 +83,7 @@ class _Provider: entitlement_provider=_Entitlements(), directory_provider=_Directory(), manifest_provider=self.manifest_provider, + model_catalog_provider=self, verification_key_id='root-2026', ) diff --git a/tests/unit_tests/cloud/test_model_catalog.py b/tests/unit_tests/cloud/test_model_catalog.py new file mode 100644 index 000000000..fc6833d28 --- /dev/null +++ b/tests/unit_tests/cloud/test_model_catalog.py @@ -0,0 +1,271 @@ +from __future__ import annotations + +import logging +from datetime import UTC, datetime +from types import SimpleNamespace + +import pytest +import sqlalchemy +from sqlalchemy.ext.asyncio import create_async_engine + +from langbot.pkg.cloud.model_catalog import ( + CloudModelCatalogSnapshot, + CloudModelCatalogSyncService, + system_model_uuid, + system_provider_uuid, +) +from langbot.pkg.entity.persistence.base import Base +from langbot.pkg.entity.persistence.model import EmbeddingModel, LLMModel, ModelProvider +from langbot.pkg.entity.persistence.workspace import Workspace +from langbot.pkg.persistence.mgr import PersistenceManager, PersistenceMode + + +pytestmark = pytest.mark.asyncio +INSTANCE_UUID = 'instance-model-catalog' +WORKSPACE_A = '00000000-0000-4000-8000-000000000001' +WORKSPACE_B = '00000000-0000-4000-8000-000000000002' +OWNER_A = '10000000-0000-4000-8000-000000000001' +OWNER_B = '10000000-0000-4000-8000-000000000002' + + +class _CatalogProvider: + def __init__(self, snapshot: CloudModelCatalogSnapshot) -> None: + self.snapshot = snapshot + + async def fetch_model_catalog(self, instance_uuid: str) -> CloudModelCatalogSnapshot: + assert instance_uuid == INSTANCE_UUID + return self.snapshot + + +def _snapshot( + *, + key_a: str | None = 'owner-a-key', + model_id: str = 'gpt-test', + include_embedding: bool = True, +) -> CloudModelCatalogSnapshot: + models = [ + { + 'uuid': 'upstream-chat', + 'model_id': model_id, + 'category': 'chat', + 'llm_abilities': ['chat', 'vision'], + 'is_featured': True, + 'featured_order': 7, + } + ] + if include_embedding: + models.append( + { + 'uuid': 'upstream-embedding', + 'model_id': 'embedding-test', + 'category': 'embedding', + } + ) + return CloudModelCatalogSnapshot.model_validate( + { + 'instance_uuid': INSTANCE_UUID, + 'generated_at': datetime.now(UTC), + 'base_url': 'https://api.langbot.cloud/v1/', + 'models': models, + 'workspaces': [ + { + 'workspace_uuid': WORKSPACE_A, + 'owner_account_uuid': OWNER_A, + 'api_key': key_a, + }, + { + 'workspace_uuid': WORKSPACE_B, + 'owner_account_uuid': OWNER_B, + 'api_key': 'owner-b-key', + }, + ], + } + ) + + +async def test_catalog_reconciles_every_workspace_idempotently_and_tracks_owner_and_downlisting(tmp_path) -> None: + engine = create_async_engine(f'sqlite+aiosqlite:///{tmp_path / "model-catalog.db"}') + manager = PersistenceManager(object(), mode=PersistenceMode.CLOUD_RUNTIME) + manager.db = SimpleNamespace(get_engine=lambda: engine) + bindings = [ + SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_A, placement_generation=1), + SimpleNamespace(instance_uuid=INSTANCE_UUID, workspace_uuid=WORKSPACE_B, placement_generation=1), + ] + workspace_service = SimpleNamespace(list_active_execution_bindings=lambda: _async_value(bindings)) + reload_counter = _AsyncCounter() + runtime_reload = SimpleNamespace(load_models_from_db=reload_counter) + app = SimpleNamespace( + persistence_mgr=manager, + workspace_service=workspace_service, + model_mgr=runtime_reload, + logger=logging.getLogger(__name__), + ) + provider = _CatalogProvider(_snapshot()) + service = CloudModelCatalogSyncService(app, provider, INSTANCE_UUID) + + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + await connection.execute( + sqlalchemy.insert(Workspace), + [ + { + 'uuid': WORKSPACE_A, + 'instance_uuid': INSTANCE_UUID, + 'name': 'A', + 'slug': 'a', + 'source': 'cloud_projection', + }, + { + 'uuid': WORKSPACE_B, + 'instance_uuid': INSTANCE_UUID, + 'name': 'B', + 'slug': 'b', + 'source': 'cloud_projection', + }, + ], + ) + await connection.execute( + sqlalchemy.insert(ModelProvider).values( + uuid='custom-provider', + workspace_uuid=WORKSPACE_A, + name='Custom', + requester='openai-chat-completions', + base_url='https://custom.example/v1', + api_keys=['custom-key'], + ) + ) + await connection.execute( + sqlalchemy.insert(LLMModel).values( + uuid='custom-model', + workspace_uuid=WORKSPACE_A, + name='custom-model', + provider_uuid='custom-provider', + abilities=['chat'], + extra_args={}, + prefered_ranking=0, + ) + ) + + first = await service.sync_once() + assert first == {'workspaces': 2, 'created': 6, 'updated': 0, 'deleted': 0} + assert reload_counter.calls == 1 + + async with engine.connect() as connection: + providers = ( + await connection.execute( + sqlalchemy.select( + ModelProvider.uuid, + ModelProvider.workspace_uuid, + ModelProvider.api_keys, + ).where(ModelProvider.requester == 'space-chat-completions') + ) + ).all() + assert {item.workspace_uuid for item in providers} == {WORKSPACE_A, WORKSPACE_B} + assert {item.uuid for item in providers} == { + system_provider_uuid(WORKSPACE_A), + system_provider_uuid(WORKSPACE_B), + } + assert {item.workspace_uuid: item.api_keys for item in providers} == { + WORKSPACE_A: ['owner-a-key'], + WORKSPACE_B: ['owner-b-key'], + } + assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(LLMModel)) == 3 + assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(EmbeddingModel)) == 2 + + second = await service.sync_once() + assert second == {'workspaces': 2, 'created': 0, 'updated': 0, 'deleted': 0} + assert reload_counter.calls == 1 + + provider.snapshot = _snapshot( + key_a='new-owner-key', + model_id='gpt-renamed', + include_embedding=False, + ) + third = await service.sync_once() + assert third == {'workspaces': 2, 'created': 0, 'updated': 3, 'deleted': 2} + assert reload_counter.calls == 2 + + async with engine.connect() as connection: + provider_a_keys = await connection.scalar( + sqlalchemy.select(ModelProvider.api_keys).where(ModelProvider.uuid == system_provider_uuid(WORKSPACE_A)) + ) + assert provider_a_keys == ['new-owner-key'] + system_model_names = ( + ( + await connection.execute( + sqlalchemy.select(LLMModel.name).where( + LLMModel.provider_uuid.in_( + [system_provider_uuid(WORKSPACE_A), system_provider_uuid(WORKSPACE_B)] + ) + ) + ) + ) + .scalars() + .all() + ) + assert set(system_model_names) == {'gpt-renamed'} + assert await connection.scalar(sqlalchemy.select(sqlalchemy.func.count()).select_from(EmbeddingModel)) == 0 + assert ( + await connection.scalar( + sqlalchemy.select(sqlalchemy.func.count()) + .select_from(ModelProvider) + .where(ModelProvider.uuid == 'custom-provider') + ) + == 1 + ) + assert ( + await connection.scalar( + sqlalchemy.select(sqlalchemy.func.count()) + .select_from(LLMModel) + .where(LLMModel.uuid == 'custom-model') + ) + == 1 + ) + + provider.snapshot = _snapshot(key_a=None, model_id='gpt-renamed', include_embedding=False) + fourth = await service.sync_once() + assert fourth == {'workspaces': 2, 'created': 0, 'updated': 1, 'deleted': 0} + assert reload_counter.calls == 3 + async with engine.connect() as connection: + provider_a_keys = await connection.scalar( + sqlalchemy.select(ModelProvider.api_keys).where(ModelProvider.uuid == system_provider_uuid(WORKSPACE_A)) + ) + assert provider_a_keys == [] + finally: + await engine.dispose() + + +def test_workspace_scoped_ids_are_stable_and_secrets_are_redacted() -> None: + assert system_provider_uuid(WORKSPACE_A) == system_provider_uuid(WORKSPACE_A) + assert system_provider_uuid(WORKSPACE_A) != system_provider_uuid(WORKSPACE_B) + assert system_model_uuid(WORKSPACE_A, 'chat', 'upstream') != system_model_uuid(WORKSPACE_B, 'chat', 'upstream') + snapshot = _snapshot() + assert 'owner-a-key' not in repr(snapshot) + + +async def test_snapshot_must_cover_every_active_workspace() -> None: + snapshot = _snapshot().model_copy(update={'workspaces': _snapshot().workspaces[:1]}) + app = SimpleNamespace( + workspace_service=SimpleNamespace( + list_active_execution_bindings=lambda: _async_value( + [SimpleNamespace(workspace_uuid=WORKSPACE_A), SimpleNamespace(workspace_uuid=WORKSPACE_B)] + ) + ), + logger=logging.getLogger(__name__), + ) + service = CloudModelCatalogSyncService(app, _CatalogProvider(snapshot), INSTANCE_UUID) + with pytest.raises(ValueError, match='missing billing projections for 1 active Workspaces'): + await service.sync_once() + + +async def _async_value(value): + return value + + +class _AsyncCounter: + def __init__(self) -> None: + self.calls = 0 + + async def __call__(self) -> None: + self.calls += 1