Implement muxing routing cache (#851)

* Implement muxing routing cache

This allows us to have a parsed and in-memory representation of the
routing rule engine. The intent is to reduce calls to the database.

This is more expensive to maintain since we need to refresh the cache on
every operation pertaining to models, endpoints, workspaces, and muxes themselves.

Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com>

* Add logic to repopulate mux cache in case a model changes.

Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com>

* Add TODO

Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com>

* Changes to integrate mux cache with mux router

* Changed Lock primites to asyncio

---------

Signed-off-by: Juan Antonio Osorio <ozz@stacklok.com>
Co-authored-by: Alejandro Ponce <aponcedeleonch@stacklok.com>

Author: JAORMX

src/codegate/cli.py

@@ -21,6 +21,7 @@
 from codegate.providers.copilot.provider import CopilotProvider
 from codegate.server import init_app
 from codegate.storage.utils import restore_storage_backup
+from codegate.workspaces import crud as wscrud
 
 
 class UvicornServer:
@@ -341,6 +342,8 @@ def serve(  # noqa: C901
 
         registry = app.provider_registry
         loop.run_until_complete(provendcrud.initialize_provider_endpoints(registry))
+        wsc = wscrud.WorkspaceCrud()
+        loop.run_until_complete(wsc.initialize_mux_registry())
 
         # Run the server
         try:

src/codegate/db/connection.py

@@ -20,7 +20,6 @@
     GetPromptWithOutputsRow,
     GetWorkspaceByNameConditions,
     MuxRule,
-    MuxRuleProviderEndpoint,
     Output,
     Prompt,
     ProviderAuthMaterial,
@@ -711,6 +710,22 @@ async def get_provider_endpoint_by_id(self, provider_id: str) -> Optional[Provid
         )
         return provider[0] if provider else None
 
+    async def get_auth_material_by_provider_id(
+        self, provider_id: str
+    ) -> Optional[ProviderAuthMaterial]:
+        sql = text(
+            """
+            SELECT id as provider_endpoint_id, auth_type, auth_blob
+            FROM provider_endpoints
+            WHERE id = :provider_endpoint_id
+            """
+        )
+        conditions = {"provider_endpoint_id": provider_id}
+        auth_material = await self._exec_select_conditions_to_pydantic(
+            ProviderAuthMaterial, sql, conditions, should_raise=True
+        )
+        return auth_material[0] if auth_material else None
+
     async def get_provider_endpoints(self) -> List[ProviderEndpoint]:
         sql = text(
             """
@@ -778,26 +793,6 @@ async def get_muxes_by_workspace(self, workspace_id: str) -> List[MuxRule]:
         )
         return muxes
 
-    async def get_muxes_with_provider_by_workspace(
-        self, workspace_id: str
-    ) -> List[MuxRuleProviderEndpoint]:
-        sql = text(
-            """
-            SELECT m.id, m.provider_endpoint_id, m.provider_model_name, m.workspace_id,
-            m.matcher_type, m.matcher_blob, m.priority, m.created_at, m.updated_at,
-            pe.provider_type, pe.endpoint, pe.auth_type, pe.auth_blob
-            FROM muxes m
-            INNER JOIN provider_endpoints pe ON pe.id = m.provider_endpoint_id
-            WHERE m.workspace_id = :workspace_id
-            ORDER BY priority ASC
-            """
-        )
-        conditions = {"workspace_id": workspace_id}
-        muxes = await self._exec_select_conditions_to_pydantic(
-            MuxRuleProviderEndpoint, 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

@@ -199,19 +199,3 @@ class MuxRule(BaseModel):
     priority: int
     created_at: Optional[datetime.datetime] = None
     updated_at: Optional[datetime.datetime] = None
-
-
-class MuxRuleProviderEndpoint(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
-    provider_type: ProviderType
-    endpoint: str
-    auth_type: str
-    auth_blob: str

src/codegate/muxing/__init__.py

@@ -10,6 +10,7 @@
 from ollama import ChatResponse
 
 from codegate.db import models as db_models
+from codegate.muxing import rulematcher
 from codegate.providers.ollama.adapter import OLlamaToModel
 
 logger = structlog.get_logger("codegate")
@@ -55,21 +56,17 @@ def _from_anthropic_to_openai(self, anthropic_body: dict) -> dict:
             del new_body["system"]
         return new_body
 
-    def _get_provider_formatted_url(
-        self, mux_and_provider: db_models.MuxRuleProviderEndpoint
-    ) -> str:
+    def _get_provider_formatted_url(self, model_route: rulematcher.ModelRoute) -> str:
         """Get the provider formatted URL to use in base_url. Note this value comes from DB"""
-        if mux_and_provider.provider_type == db_models.ProviderType.openai:
-            return f"{mux_and_provider.endpoint}/v1"
-        return mux_and_provider.endpoint
+        if model_route.endpoint.provider_type == db_models.ProviderType.openai:
+            return f"{model_route.endpoint.endpoint}/v1"
+        return model_route.endpoint.endpoint
 
-    def _set_destination_info(
-        self, data: dict, mux_and_provider: db_models.MuxRuleProviderEndpoint
-    ) -> dict:
+    def _set_destination_info(self, data: dict, model_route: rulematcher.ModelRoute) -> dict:
         """Set the destination provider info."""
         new_data = copy.deepcopy(data)
-        new_data["model"] = mux_and_provider.provider_model_name
-        new_data["base_url"] = self._get_provider_formatted_url(mux_and_provider)
+        new_data["model"] = model_route.model.name
+        new_data["base_url"] = self._get_provider_formatted_url(model_route)
         return new_data
 
     def _identify_provider(self, data: dict) -> db_models.ProviderType:
@@ -79,22 +76,20 @@ def _identify_provider(self, data: dict) -> db_models.ProviderType:
         else:
             return db_models.ProviderType.openai
 
-    def map_body_to_dest(
-        self, mux_and_provider: db_models.MuxRuleProviderEndpoint, data: dict
-    ) -> dict:
+    def map_body_to_dest(self, model_route: rulematcher.ModelRoute, data: dict) -> dict:
         """
         Map the body to the destination provider.
 
         We only need to transform the body if the destination or origin provider is Anthropic.
         """
         origin_prov = self._identify_provider(data)
-        if mux_and_provider.provider_type == db_models.ProviderType.anthropic:
+        if model_route.endpoint.provider_type == db_models.ProviderType.anthropic:
             if origin_prov != db_models.ProviderType.anthropic:
                 data = self._from_openai_to_antrhopic(data)
         else:
             if origin_prov == db_models.ProviderType.anthropic:
                 data = self._from_anthropic_to_openai(data)
-        return self._set_destination_info(data, mux_and_provider)
+        return self._set_destination_info(data, model_route)
 
 
 class StreamChunkFormatter:

src/codegate/mux/adapter.py renamed to src/codegate/muxing/adapter.py

@@ -3,7 +3,8 @@
 import structlog
 from fastapi import APIRouter, HTTPException, Request
 
-from codegate.mux.adapter import BodyAdapter, ResponseAdapter
+from codegate.muxing import rulematcher
+from codegate.muxing.adapter import BodyAdapter, ResponseAdapter
 from codegate.providers.registry import ProviderRegistry
 from codegate.workspaces.crud import WorkspaceCrud
 
@@ -53,27 +54,27 @@ async def route_to_dest_provider(
             body = await request.body()
             data = json.loads(body)
 
+            mux_registry = await rulematcher.get_muxing_rules_registry()
             try:
-                active_ws_muxes = await self._ws_crud.get_active_workspace_muxes()
+                # TODO: For future releases we will have to idenify a thing_to_match
+                # and use our registry to get the correct muxes for the active workspace
+                model_route = await mux_registry.get_match_for_active_workspace(thing_to_match=None)
             except Exception as e:
                 logger.error(f"Error getting active workspace muxes: {e}")
                 raise HTTPException(str(e))
 
-            # TODO: Implement the muxing logic here. For the moment we will assume
-            # that we have a single mux, i.e. a single destination provider.
-            if len(active_ws_muxes) == 0:
-                raise HTTPException(status_code=404, detail="No active workspace muxes found")
-            mux_and_provider = active_ws_muxes[0]
+            if not model_route:
+                raise HTTPException("No rule found for the active workspace", status_code=404)
 
             # Parse the input data and map it to the destination provider format
             rest_of_path = self._ensure_path_starts_with_slash(rest_of_path)
-            new_data = self._body_adapter.map_body_to_dest(mux_and_provider, data)
-            provider = self._provider_registry.get_provider(mux_and_provider.provider_type)
-            api_key = mux_and_provider.auth_blob
+            new_data = self._body_adapter.map_body_to_dest(model_route, data)
+            provider = self._provider_registry.get_provider(model_route.endpoint.provider_type)
+            api_key = model_route.auth_material.auth_blob
 
             # Send the request to the destination provider. It will run the pipeline
             response = await provider.process_request(new_data, api_key, rest_of_path)
             # Format the response to the client always using the OpenAI format
             return self._response_adapter.format_response_to_client(
-                response, mux_and_provider.provider_type
+                response, model_route.endpoint.provider_type
             )

src/codegate/mux/router.py renamed to src/codegate/muxing/router.py

@@ -0,0 +1,126 @@
+import copy
+from abc import ABC, abstractmethod
+from asyncio import Lock
+from typing import List, Optional
+
+from codegate.db import models as db_models
+
+_muxrules_sgtn = None
+
+_singleton_lock = Lock()
+
+
+async def get_muxing_rules_registry():
+    """Returns a singleton instance of the muxing rules registry."""
+
+    global _muxrules_sgtn
+
+    if _muxrules_sgtn is None:
+        async with _singleton_lock:
+            if _muxrules_sgtn is None:
+                _muxrules_sgtn = MuxingRulesinWorkspaces()
+
+    return _muxrules_sgtn
+
+
+class ModelRoute:
+    """A route for a model."""
+
+    def __init__(
+        self,
+        model: db_models.ProviderModel,
+        endpoint: db_models.ProviderEndpoint,
+        auth_material: db_models.ProviderAuthMaterial,
+    ):
+        self.model = model
+        self.endpoint = endpoint
+        self.auth_material = auth_material
+
+
+class MuxingRuleMatcher(ABC):
+    """Base class for matching muxing rules."""
+
+    def __init__(self, route: ModelRoute):
+        self._route = route
+
+    @abstractmethod
+    def match(self, thing_to_match) -> bool:
+        """Return True if the rule matches the thing_to_match."""
+        pass
+
+    def destination(self) -> ModelRoute:
+        """Return the destination of the rule."""
+
+        return self._route
+
+
+class MuxingMatcherFactory:
+    """Factory for creating muxing matchers."""
+
+    @staticmethod
+    def create(mux_rule: db_models.MuxRule, route: ModelRoute) -> MuxingRuleMatcher:
+        """Create a muxing matcher for the given endpoint and model."""
+
+        factory = {
+            "catch_all": CatchAllMuxingRuleMatcher,
+        }
+
+        try:
+            return factory[mux_rule.matcher_type](route)
+        except KeyError:
+            raise ValueError(f"Unknown matcher type: {mux_rule.matcher_type}")
+
+
+class CatchAllMuxingRuleMatcher(MuxingRuleMatcher):
+    """A catch all muxing rule matcher."""
+
+    def match(self, thing_to_match) -> bool:
+        return True
+
+
+class MuxingRulesinWorkspaces:
+    """A thread safe dictionary to store the muxing rules in workspaces."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._lock = Lock()
+        self._active_workspace = ""
+        self._ws_rules = {}
+
+    async def get_ws_rules(self, workspace_name: str) -> List[MuxingRuleMatcher]:
+        """Get the rules for the given workspace."""
+        async with self._lock:
+            return copy.deepcopy(self._ws_rules.get(workspace_name, []))
+
+    async def set_ws_rules(self, workspace_name: str, rules: List[MuxingRuleMatcher]) -> None:
+        """Set the rules for the given workspace."""
+        async with self._lock:
+            self._ws_rules[workspace_name] = rules
+
+    async def delete_ws_rules(self, workspace_name: str) -> None:
+        """Delete the rules for the given workspace."""
+        async with self._lock:
+            del self._ws_rules[workspace_name]
+
+    async def set_active_workspace(self, workspace_name: str) -> None:
+        """Set the active workspace."""
+        self._active_workspace = workspace_name
+
+    async def get_registries(self) -> List[str]:
+        """Get the list of workspaces."""
+        async with self._lock:
+            return list(self._ws_rules.keys())
+
+    async def get_match_for_active_workspace(self, thing_to_match) -> Optional[ModelRoute]:
+        """Get the first match for the given thing_to_match."""
+
+        # We iterate over all the rules and return the first match
+        # Since we already do a deepcopy in __getitem__, we don't need to lock here
+        try:
+            rules = await self.get_ws_rules(self._active_workspace)
+            for rule in rules:
+                if rule.match(thing_to_match):
+                    return rule.destination()
+            return None
+        except KeyError:
+            raise RuntimeError("No rules found for the active workspace")

src/codegate/muxing/rulematcher.py

@@ -12,6 +12,7 @@
 from codegate.db.connection import DbReader, DbRecorder
 from codegate.providers.base import BaseProvider
 from codegate.providers.registry import ProviderRegistry, get_provider_registry
+from codegate.workspaces import crud as workspace_crud
 
 logger = structlog.get_logger("codegate")
 
@@ -32,6 +33,7 @@ class ProviderCrud:
     def __init__(self):
         self._db_reader = DbReader()
         self._db_writer = DbRecorder()
+        self._ws_crud = workspace_crud.WorkspaceCrud()
 
     async def list_endpoints(self) -> List[apimodelsv1.ProviderEndpoint]:
         """List all the endpoints."""
@@ -176,6 +178,9 @@ async def update_endpoint(
                 )
             )
 
+        # a model might have been deleted, let's repopulate the cache
+        await self._ws_crud.repopulate_mux_cache()
+
         return apimodelsv1.ProviderEndpoint.from_db_model(dbendpoint)
 
     async def configure_auth_material(
@@ -208,6 +213,8 @@ async def delete_endpoint(self, provider_id: UUID):
 
         await self._db_writer.delete_provider_endpoint(dbendpoint)
 
+        await self._ws_crud.repopulate_mux_cache()
+
     async def models_by_provider(self, provider_id: UUID) -> List[apimodelsv1.ModelByProvider]:
         """Get the models by provider."""
 

src/codegate/providers/crud/crud.py

@@ -11,7 +11,7 @@
 from codegate import __description__, __version__
 from codegate.api.v1 import v1
 from codegate.db.models import ProviderType
-from codegate.mux.router import MuxRouter
+from codegate.muxing.router import MuxRouter
 from codegate.pipeline.factory import PipelineFactory
 from codegate.providers.anthropic.provider import AnthropicProvider
 from codegate.providers.llamacpp.provider import LlamaCppProvider