Implement muxes CRUD

Signed-off-by: Juan Antonio Osorio <[email protected]>

Author: JAORMX

src/codegate/api/v1.py

@@ -481,22 +481,15 @@ async def get_workspace_muxes(
 
     The list is ordered in order of priority. That is, the first rule in the list
     has the highest priority."""
-    # TODO: This is a dummy implementation. In the future, we should have a proper
-    # implementation that fetches the mux rules from the database.
-    return [
-        v1_models.MuxRule(
-            # Hardcode some UUID just for mocking purposes
-            provider_id="00000000-0000-0000-0000-000000000001",
-            model="gpt-3.5-turbo",
-            matcher_type=v1_models.MuxMatcherType.file_regex,
-            matcher=".*\\.txt",
-        ),
-        v1_models.MuxRule(
-            provider_id="00000000-0000-0000-0000-000000000002",
-            model="davinci",
-            matcher_type=v1_models.MuxMatcherType.catch_all,
-        ),
-    ]
+    try:
+        muxes = await wscrud.get_muxes(workspace_name)
+    except crud.WorkspaceDoesNotExistError:
+        raise HTTPException(status_code=404, detail="Workspace does not exist")
+    except Exception:
+        logger.exception("Error while getting workspace")
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+    return muxes
 
 
 @v1.put(
@@ -510,8 +503,16 @@ async def set_workspace_muxes(
     request: List[v1_models.MuxRule],
 ):
     """Set the mux rules of a workspace."""
-    # TODO: This is a dummy implementation. In the future, we should have a proper
-    # implementation that sets the mux rules in the database.
+    try:
+        await wscrud.set_muxes(workspace_name, request)
+    except crud.WorkspaceDoesNotExistError:
+        raise HTTPException(status_code=404, detail="Workspace does not exist")
+    except crud.WorkspaceCrudError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+    except Exception:
+        logger.exception("Error while setting muxes")
+        raise HTTPException(status_code=500, detail="Internal server error")
+
     return Response(status_code=204)
 
 

src/codegate/api/v1_models.py

@@ -287,11 +287,6 @@ class MuxMatcherType(str, Enum):
     Represents the different types of matchers we support.
     """
 
-    # Match a regular expression for a file path
-    # in the prompt. Note that if no file is found,
-    # the prompt will be passed through.
-    file_regex = "file_regex"
-
     # Always match this prompt
     catch_all = "catch_all"
 

src/codegate/db/connection.py

@@ -19,6 +19,7 @@
     Alert,
     GetPromptWithOutputsRow,
     GetWorkspaceByNameConditions,
+    MuxRule,
     Output,
     Prompt,
     ProviderAuthMaterial,
@@ -112,6 +113,15 @@ async def _execute_update_pydantic_model(
                 raise e
             return None
 
+    async def _execute_with_no_return(self, sql_command: TextClause, conditions: dict):
+        """Execute a command that doesn't return anything."""
+        try:
+            async with self._async_db_engine.begin() as conn:
+                await conn.execute(sql_command, conditions)
+        except Exception as e:
+            logger.error(f"Failed to execute command: {sql_command}.", error=str(e))
+            raise e
+
     async def record_request(self, prompt_params: Optional[Prompt] = None) -> Optional[Prompt]:
         if prompt_params is None:
             return None
@@ -459,22 +469,44 @@ async def add_provider_model(self, model: ProviderModel) -> ProviderModel:
         added_model = await self._execute_update_pydantic_model(model, sql, should_raise=True)
         return added_model
 
-    async def delete_provider_models(self, provider_id: str) -> Optional[ProviderModel]:
+    async def delete_provider_models(self, provider_id: str):
         sql = text(
             """
             DELETE FROM provider_models
             WHERE provider_endpoint_id = :provider_endpoint_id
+            """
+        )
+        conditions = {"provider_endpoint_id": provider_id}
+        await self._execute_with_no_return(sql, conditions)
+
+    async def delete_muxes_by_workspace(self, workspace_id: str):
+        sql = text(
+            """
+            DELETE FROM muxes
+            WHERE workspace_id = :workspace_id
             RETURNING *
             """
         )
-        await self._execute_update_pydantic_model(
-            ProviderModel(
-                provider_endpoint_id=provider_id,
-                name="Fake name to respect the signature of the function",
-            ),
-            sql,
-            should_raise=True,
+
+        conditions = {"workspace_id": workspace_id}
+        await self._execute_with_no_return(sql, conditions)
+
+    async def add_mux(self, mux: MuxRule) -> MuxRule:
+        sql = text(
+            """
+            INSERT INTO muxes (
+                id, provider_endpoint_id, provider_model_name, workspace_id, matcher_type,
+                matcher_blob, priority, created_at, updated_at
+            )
+            VALUES (
+                :id, :provider_endpoint_id, :provider_model_name, :workspace_id,
+                :matcher_type, :matcher_blob, :priority, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP
+            )
+            RETURNING *
+            """
         )
+        added_mux = await self._execute_update_pydantic_model(mux, sql, should_raise=True)
+        return added_mux
 
 
 class DbReader(DbCodeGate):
@@ -701,6 +733,22 @@ async def get_provider_models_by_provider_id(self, provider_id: str) -> List[Pro
         )
         return models
 
+    async def get_provider_model_by_provider_id_and_name(
+        self, provider_id: str, model_name: str
+    ) -> Optional[ProviderModel]:
+        sql = text(
+            """
+            SELECT provider_endpoint_id, name
+            FROM provider_models
+            WHERE provider_endpoint_id = :provider_endpoint_id AND name = :name
+            """
+        )
+        conditions = {"provider_endpoint_id": provider_id, "name": model_name}
+        models = await self._exec_select_conditions_to_pydantic(
+            ProviderModel, sql, conditions, should_raise=True
+        )
+        return models[0] if models else None
+
     async def get_all_provider_models(self) -> List[ProviderModel]:
         sql = text(
             """
@@ -712,6 +760,22 @@ async def get_all_provider_models(self) -> List[ProviderModel]:
         models = await self._execute_select_pydantic_model(ProviderModel, sql)
         return models
 
+    async def get_muxes_by_workspace(self, workspace_id: str) -> List[MuxRule]:
+        sql = text(
+            """
+            SELECT id, provider_endpoint_id, provider_model_name, workspace_id, matcher_type,
+            matcher_blob, priority, created_at, updated_at
+            FROM muxes
+            WHERE workspace_id = :workspace_id
+            ORDER BY priority ASC
+            """
+        )
+        conditions = {"workspace_id": workspace_id}
+        muxes = await self._exec_select_conditions_to_pydantic(
+            MuxRule, sql, conditions, should_raise=True
+        )
+        return muxes
+
 
 def init_db_sync(db_path: Optional[str] = None):
     """DB will be initialized in the constructor in case it doesn't exist."""

src/codegate/db/models.py

@@ -173,3 +173,15 @@ class ProviderModel(BaseModel):
     provider_endpoint_id: str
     provider_endpoint_name: Optional[str] = None
     name: str
+
+
+class MuxRule(BaseModel):
+    id: str
+    provider_endpoint_id: str
+    provider_model_name: str
+    workspace_id: str
+    matcher_type: str
+    matcher_blob: str
+    priority: int
+    created_at: Optional[datetime.datetime] = None
+    updated_at: Optional[datetime.datetime] = None

src/codegate/workspaces/crud.py

@@ -1,8 +1,16 @@
+import asyncio
 import datetime
 from typing import List, Optional, Tuple
+from uuid import uuid4 as uuid
 
 from codegate.db.connection import DbReader, DbRecorder
-from codegate.db.models import ActiveWorkspace, Session, WorkspaceRow, WorkspaceWithSessionInfo
+from codegate.db.models import (
+    ActiveWorkspace,
+    MuxRule,
+    Session,
+    WorkspaceRow,
+    WorkspaceWithSessionInfo,
+)
 
 
 class WorkspaceCrudError(Exception):
@@ -17,6 +25,10 @@ class WorkspaceAlreadyActiveError(WorkspaceCrudError):
     pass
 
 
+class WorkspaceMuxRuleDoesNotExistError(WorkspaceCrudError):
+    pass
+
+
 DEFAULT_WORKSPACE_NAME = "default"
 
 # These are reserved keywords that cannot be used for workspaces
@@ -202,3 +214,72 @@ async def get_workspace_by_name(self, workspace_name: str) -> WorkspaceRow:
         if not workspace:
             raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.")
         return workspace
+
+    # Can't use type hints since the models are not yet defined
+    # Note that I'm explicitly importing the models here to avoid circular imports.
+    async def get_muxes(self, workspace_name: str):
+        from codegate.api import v1_models
+
+        # Verify if workspace exists
+        workspace = await self._db_reader.get_workspace_by_name(workspace_name)
+        if not workspace:
+            raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.")
+
+        dbmuxes = await self._db_reader.get_muxes_by_workspace(workspace.id)
+
+        muxes = []
+        # These are already sorted by priority
+        for dbmux in dbmuxes:
+            muxes.append(
+                v1_models.MuxRule(
+                    provider_id=dbmux.provider_endpoint_id,
+                    model=dbmux.provider_model_name,
+                    matcher_type=dbmux.matcher_type,
+                    matcher=dbmux.matcher_blob,
+                )
+            )
+
+        return muxes
+
+    # Can't use type hints since the models are not yet defined
+    async def set_muxes(self, workspace_name: str, muxes):
+        # Verify if workspace exists
+        workspace = await self._db_reader.get_workspace_by_name(workspace_name)
+        if not workspace:
+            raise WorkspaceDoesNotExistError(f"Workspace {workspace_name} does not exist.")
+
+        # Delete all muxes for the workspace
+        db_recorder = DbRecorder()
+        await db_recorder.delete_muxes_by_workspace(workspace.id)
+
+        tasks = set()
+
+        # Add the new muxes
+        priority = 0
+
+        # Verify all models are valid
+        for mux in muxes:
+            dbm = await self._db_reader.get_provider_model_by_provider_id_and_name(
+                mux.provider_id,
+                mux.model,
+            )
+            if not dbm:
+                raise WorkspaceCrudError(
+                    f"Model {mux.model} does not exist for provider {mux.provider_id}"
+                )
+
+        for mux in muxes:
+            new_mux = MuxRule(
+                id=str(uuid()),
+                provider_endpoint_id=mux.provider_id,
+                provider_model_name=mux.model,
+                workspace_id=workspace.id,
+                matcher_type=mux.matcher_type,
+                matcher_blob=mux.matcher if mux.matcher else "",
+                priority=priority,
+            )
+            tasks.add(db_recorder.add_mux(new_mux))
+
+            priority += 1
+
+        await asyncio.gather(*tasks)