[owl] Add model replacement route (#922)

added an admin model management route to replace the model used in a gentable
currently doesn't support replacing embedding model
hardened redis async client creation with a loop-aware design

Author: noobHappylife

clients/python/src/jamaibase/client.py

@@ -42,6 +42,8 @@
     EmbeddingResponse,
     FileUploadResponse,
     GenConfigUpdateRequest,
+    GenTableModelReplaceProgressKeys,
+    GenTableModelReplaceRequest,
     GetURLRequest,
     GetURLResponse,
     KnowledgeTableSchemaCreate,
@@ -1229,6 +1231,37 @@ async def delete_deployment(self, deployment_id: str, **kwargs) -> OkResponse:
             **kwargs,
         )
 
+    async def replace_model_ids(
+        self,
+        request: GenTableModelReplaceRequest,
+        **kwargs,
+    ) -> OkResponse:
+        """
+        Replace model IDs in GenTable generation configs.
+
+        Args:
+            request (GenTableModelReplaceRequest): The model replacement request.
+
+        Returns:
+            response (OkResponse): Response containing the progress key.
+        """
+        return await self._post(
+            "/v2/models/replace",
+            body=request,
+            response_model=OkResponse,
+            **kwargs,
+        )
+
+    async def list_model_replace_progress_keys(
+        self,
+        **kwargs,
+    ) -> GenTableModelReplaceProgressKeys:
+        return await self._get(
+            "/v2/models/replace/progress_keys",
+            response_model=GenTableModelReplaceProgressKeys,
+            **kwargs,
+        )
+
 
 class _OrganizationsAsync(_ClientAsync):
     """Organization methods."""
@@ -4731,6 +4764,19 @@ def update_deployment(
     def delete_deployment(self, deployment_id: str, **kwargs) -> OkResponse:
         return LOOP.run(super().delete_deployment(deployment_id, **kwargs))
 
+    def replace_model_ids(
+        self,
+        request: GenTableModelReplaceRequest,
+        **kwargs,
+    ) -> OkResponse:
+        return LOOP.run(super().replace_model_ids(request, **kwargs))
+
+    def list_model_replace_progress_keys(
+        self,
+        **kwargs,
+    ) -> GenTableModelReplaceProgressKeys:
+        return LOOP.run(super().list_model_replace_progress_keys(**kwargs))
+
 
 class _Organizations(_OrganizationsAsync):
     """Organization methods."""

clients/python/src/jamaibase/types/__init__.py

@@ -164,6 +164,8 @@
     DiscriminatedGenConfig,
     EmbedGenConfig,
     GenConfigUpdateRequest,
+    GenTableModelReplaceProgressKeys,
+    GenTableModelReplaceRequest,
     ImageGenConfig,
     KnowledgeTableSchemaCreate,
     LLMGenConfig,

clients/python/src/jamaibase/types/gen_table.py

@@ -392,6 +392,31 @@ def check_column_map(self) -> Self:
         return self
 
 
+class GenTableModelReplaceRequest(BaseModel):
+    mapping: dict[str, str] = Field(
+        min_length=1,
+        description="Mapping of old model IDs to replacement model IDs.",
+    )
+    organization_ids: list[str] | None = Field(
+        None,
+        description="Optional organization IDs to scan. If omitted, all organizations are scanned.",
+    )
+
+    @model_validator(mode="after")
+    def check_mapping(self) -> Self:
+        for old_id, new_id in self.mapping.items():
+            if old_id == new_id:
+                raise ValueError(f'Model replacement maps "{old_id}" to itself.')
+        return self
+
+
+class GenTableModelReplaceProgressKeys(BaseModel):
+    items: list[str] = Field(
+        default_factory=list,
+        description="Recent GenTable model replacement progress keys, newest first.",
+    )
+
+
 class ColumnRenameRequest(BaseModel):
     table_id: str = Field(
         description="Table name or ID.",

docker/override.dev.yml

@@ -56,6 +56,10 @@ services:
       - ${PWD}/docker_data/owl/db:/usr/src/app/db
       - ${PWD}/services/api/src:/usr/src/app/api/src
 
+  starling:
+    volumes:
+      - ${PWD}/services/api/src:/usr/src/app/api/src
+
   kopi:
     ports:
       - 5569:3000

services/api/src/owl/docparse.py

@@ -245,9 +245,7 @@ async def load_document(
                 md = await CACHE.get(cache_key)
             if md is not None:
                 # Extend cache TTL
-                await CACHE._redis_async.expire(
-                    cache_key, ENV_CONFIG.document_loader_cache_ttl_sec
-                )
+                await CACHE.expire(cache_key, ENV_CONFIG.document_loader_cache_ttl_sec)
                 logger.debug(f'File "{file_name}" loaded from cache (cache key="{cache_key}").')
                 return md
             try:
@@ -361,7 +359,7 @@ async def load_document_chunks(
                 chunk_json_str = await CACHE.get(cache_key)
             if chunk_json_str is not None:
                 # Extend cache TTL
-                await CACHE._redis_async.expire(cache_key, cache_ttl)
+                await CACHE.expire(cache_key, cache_ttl)
                 logger.info(
                     f'File chunks "{file_name}" loaded from cache (cache key="{cache_key}").'
                 )

services/api/src/owl/routers/models.py

@@ -5,20 +5,25 @@
 from sqlmodel import func, select
 
 from owl.db import AsyncSession, yield_async_session
-from owl.db.models import Deployment, ModelConfig
+from owl.db.models import Deployment, ModelConfig, Organization
+from owl.tasks.gen_table import replace_gen_table_models
 from owl.types import (
     CloudProvider,
     DeploymentCreate,
     DeploymentRead,
     DeploymentUpdate,
+    GenTableModelReplaceProgressKeys,
+    GenTableModelReplaceRequest,
     ListQuery,
     ModelConfigCreate,
     ModelConfigRead,
     ModelConfigUpdate,
+    ModelType,
     OkResponse,
     Page,
     UserAuth,
 )
+from owl.utils import uuid7_str
 from owl.utils.auth import auth_user_service_key, has_permissions
 from owl.utils.dates import now
 from owl.utils.exceptions import (
@@ -27,6 +32,14 @@
     ResourceNotFoundError,
     handle_exception,
 )
+from owl.utils.gen_table_model_replace import (
+    acquire_model_replace_lock,
+    delete_model_replace_progress,
+    get_model_replace_lock,
+    initialize_model_replace_progress,
+    list_recent_model_replace_progress_keys,
+    release_model_replace_lock,
+)
 
 router = APIRouter()
 
@@ -316,3 +329,76 @@ async def delete_deployment(
     await session.delete(deployment)
     await session.commit()
     return OkResponse()
+
+
+@router.post(
+    "/v2/models/replace",
+    summary="Replace model IDs in GenTable generation configs.",
+    description="Permissions: `system.MEMBER`.",
+)
+@handle_exception
+async def replace_gen_table_model_ids(
+    request: Request,
+    user: Annotated[UserAuth, Depends(auth_user_service_key)],
+    session: Annotated[AsyncSession, Depends(yield_async_session)],
+    body: GenTableModelReplaceRequest,
+) -> OkResponse:
+    logger.info(f"{request.state.id} - Replacing GenTable model IDs: {body}")
+    has_permissions(user, ["system.MEMBER"])
+    if body.organization_ids is not None:
+        for organization_id in body.organization_ids:
+            await Organization.get(session, organization_id, name="Organization")
+    model_configs = {
+        model_id: await ModelConfig.get(session, model_id, name="Model config")
+        for model_id in sorted(set(body.mapping) | set(body.mapping.values()))
+    }
+    embedding_model_ids = [
+        model_id
+        for model_id, model_config in model_configs.items()
+        if model_config is not None and model_config.type == ModelType.EMBED
+    ]
+    if embedding_model_ids:
+        raise BadInputError(f"Replacing embedding model is not supported: {embedding_model_ids}.")
+    progress_key = f"gen_table_model_replace:{uuid7_str()}"
+    if not await acquire_model_replace_lock(progress_key):
+        active_progress_key = await get_model_replace_lock()
+        raise ResourceExistsError(
+            (
+                "Another GenTable model replace task is already running"
+                + (f': progress_key="{active_progress_key}".' if active_progress_key else ".")
+            )
+        )
+    try:
+        await initialize_model_replace_progress(
+            progress_key=progress_key,
+            mapping=body.mapping,
+            organization_ids=body.organization_ids,
+            requested_by=user.id,
+        )
+        replace_gen_table_models.delay(
+            mapping=body.mapping,
+            organization_ids=body.organization_ids,
+            progress_key=progress_key,
+            requested_by=user.id,
+        )
+    except Exception:
+        await release_model_replace_lock(progress_key)
+        await delete_model_replace_progress(progress_key)
+        raise
+    logger.bind(user_id=user.id).info(
+        f'Enqueued GenTable model replace task progress_key="{progress_key}".'
+    )
+    return OkResponse(progress_key=progress_key)
+
+
+@router.get(
+    "/v2/models/replace/progress_keys",
+    summary="List recent GenTable model replacement progress keys.",
+    description="Permissions: `system.MEMBER`.",
+)
+@handle_exception
+async def list_gen_table_model_replace_progress_keys(
+    user: Annotated[UserAuth, Depends(auth_user_service_key)],
+) -> GenTableModelReplaceProgressKeys:
+    has_permissions(user, ["system.MEMBER"])
+    return GenTableModelReplaceProgressKeys(items=await list_recent_model_replace_progress_keys())

services/api/src/owl/tasks/gen_table.py

@@ -7,6 +7,7 @@
 from owl.db.gen_table import ActionTable, ChatTable, KnowledgeTable
 from owl.types import TableType
 from owl.utils.exceptions import JamaiException, ResourceExistsError
+from owl.utils.gen_table_model_replace import GenTableModelReplacer, release_model_replace_lock
 from owl.utils.io import open_uri_async
 
 TABLE_CLS: dict[TableType, ActionTable | KnowledgeTable | ChatTable] = {
@@ -60,3 +61,29 @@ async def _task():
     table = asyncio.get_event_loop().run_until_complete(_task())
     logger.info("Generative Table import task completed.")
     return table.v1_meta_response.model_dump_json()
+
+
+@celery_app.task
+def replace_gen_table_models(
+    mapping: dict[str, str],
+    *,
+    organization_ids: list[str] | None = None,
+    progress_key: str = "",
+    requested_by: str = "",
+) -> dict[str, bool | int]:
+    async def _task():
+        try:
+            replacer = GenTableModelReplacer(
+                mapping=mapping,
+                organization_ids=organization_ids,
+                progress_key=progress_key,
+                requested_by=requested_by,
+            )
+            return await replacer.run()
+        finally:
+            await release_model_replace_lock(progress_key)
+
+    logger.info("GenTable model replace task started.")
+    result = asyncio.get_event_loop().run_until_complete(_task())
+    logger.info("GenTable model replace task completed.")
+    return result.to_dict()

services/api/src/owl/types/__init__.py

@@ -75,6 +75,8 @@
     FunctionCall,
     FunctionParameters,
     GenConfigUpdateRequest,
+    GenTableModelReplaceProgressKeys,
+    GenTableModelReplaceRequest,
     GetURLRequest,
     GetURLResponse,
     Host,