Add mux rules for match according to file and request type

Closes: #885 and #944

This PR introduces the necessary changes for:
1. Muxing to match on a specific filename (`main.py`) or file type (`.py`)
2. Muxing on a specific request type, i.e. chat or FIM

The matching is tightly coupled with the priority of the muxing rules.
A match will occur on the first possible rule.

Example:
Muxing rule 1 -> match all `.py` files -> go to `chat-gpt`
Muxing rule 2 -> match all `.md` files -> go to `ollama`

If in a request we recieve a `README.md` and a `main.py` file the
request will be routed to `chat-gpt` since it's the match with
the highest priority in the list.

To introduce above changes there were some minor changes:
1. Separate FIM detection into its own class. Before it used to be part of `BaseProvider`
2. Detect if a request is FIM or not before calling `process_request` method.

Both of them were necessary in order to re-use the logic to create the
matcher for request type

Author: aponcedeleonch

src/codegate/muxing/models.py

@@ -3,6 +3,8 @@
 
 import pydantic
 
+from codegate.clients.clients import ClientType
+
 
 class MuxMatcherType(str, Enum):
     """
@@ -11,6 +13,12 @@ class MuxMatcherType(str, Enum):
 
     # Always match this prompt
     catch_all = "catch_all"
+    # Match based on the filename. It will match if there is a filename
+    # in the request that matches the matcher either extension or full name (*.py or main.py)
+    filename_match = "filename_match"
+    # Match based on the request type. It will match if the request type
+    # matches the matcher (e.g. FIM or chat)
+    request_type_match = "request_type_match"
 
 
 class MuxRule(pydantic.BaseModel):
@@ -25,3 +33,14 @@ class MuxRule(pydantic.BaseModel):
     # The actual matcher to use. Note that
     # this depends on the matcher type.
     matcher: Optional[str] = None
+
+
+class ThingToMatchMux(pydantic.BaseModel):
+    """
+    Represents the fields we can use to match a mux rule.
+    """
+
+    body: dict
+    url_request_path: str
+    is_fim_request: bool
+    client_type: ClientType

src/codegate/muxing/router.py

@@ -1,14 +1,14 @@
 import json
+from typing import Optional
 
 import structlog
 from fastapi import APIRouter, HTTPException, Request
 
-from codegate.clients.clients import ClientType
 from codegate.clients.detector import DetectClient
-from codegate.extract_snippets.body_extractor import BodyCodeSnippetExtractorError
-from codegate.extract_snippets.factory import BodyCodeExtractorFactory
+from codegate.muxing import models as mux_models
 from codegate.muxing import rulematcher
 from codegate.muxing.adapter import BodyAdapter, ResponseAdapter
+from codegate.providers.fim_analyzer import FIMAnalyzer
 from codegate.providers.registry import ProviderRegistry
 from codegate.workspaces.crud import WorkspaceCrud
 
@@ -39,40 +39,20 @@ def get_routes(self) -> APIRouter:
     def _ensure_path_starts_with_slash(self, path: str) -> str:
         return path if path.startswith("/") else f"/{path}"
 
-    def _extract_request_filenames(self, detected_client: ClientType, data: dict) -> set[str]:
+    async def _get_model_route(
+        self, thing_to_match: mux_models.ThingToMatchMux
+    ) -> Optional[rulematcher.ModelRoute]:
         """
-        Extract filenames from the request data.
+        Get the model route for the given things_to_match.
         """
-        try:
-            body_extractor = BodyCodeExtractorFactory.create_snippet_extractor(detected_client)
-            return body_extractor.extract_unique_filenames(data)
-        except BodyCodeSnippetExtractorError as e:
-            logger.error(f"Error extracting filenames from request: {e}")
-            return set()
-
-    async def _get_model_routes(self, filenames: set[str]) -> list[rulematcher.ModelRoute]:
-        """
-        Get the model routes for the given filenames.
-        """
-        model_routes = []
         mux_registry = await rulematcher.get_muxing_rules_registry()
         try:
-            # Try to get a catch_all route
-            single_model_route = await mux_registry.get_match_for_active_workspace(
-                thing_to_match=None
-            )
-            model_routes.append(single_model_route)
-
-            # Get the model routes for each filename
-            for filename in filenames:
-                model_route = await mux_registry.get_match_for_active_workspace(
-                    thing_to_match=filename
-                )
-                model_routes.append(model_route)
+            # Try to get a model route for the active workspace
+            model_route = await mux_registry.get_match_for_active_workspace(thing_to_match)
+            return model_route
         except Exception as e:
             logger.error(f"Error getting active workspace muxes: {e}")
             raise HTTPException(str(e), status_code=404)
-        return model_routes
 
     def _setup_routes(self):
 
@@ -88,34 +68,38 @@ async def route_to_dest_provider(
             1. Get destination provider from DB and active workspace.
             2. Map the request body to the destination provider format.
             3. Run pipeline. Selecting the correct destination provider.
-            4. Transmit the response back to the client in the correct format.
+            4. Transmit the response back to the client in OpenAI format.
             """
 
             body = await request.body()
             data = json.loads(body)
+            is_fim_request = FIMAnalyzer.is_fim_request(rest_of_path, data)
+
+            # 1. Get destination provider from DB and active workspace.
+            thing_to_match = mux_models.ThingToMatchMux(
+                body=data,
+                url_request_path=rest_of_path,
+                is_fim_request=is_fim_request,
+                client_type=request.state.detected_client,
+            )
+            model_route = await self._get_model_route(thing_to_match)
+            if not model_route:
+                raise HTTPException(
+                    "No matching rule found for the active workspace", status_code=404
+                )
 
-            filenames_in_data = self._extract_request_filenames(request.state.detected_client, data)
-            logger.info(f"Extracted filenames from request: {filenames_in_data}")
-
-            model_routes = await self._get_model_routes(filenames_in_data)
-            if not model_routes:
-                raise HTTPException("No rule found for the active workspace", status_code=404)
-
-            # We still need some logic here to handle the case where we have multiple model routes.
-            # For the moment since we match all only pick the first.
-            model_route = model_routes[0]
-
-            # Parse the input data and map it to the destination provider format
+            # 2. Map the request body 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(model_route, data)
+
+            # 3. Run pipeline. Selecting the correct destination provider.
             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, request.state.detected_client
+                new_data, api_key, is_fim_request, request.state.detected_client
             )
-            # Format the response to the client always using the OpenAI format
+
+            # 4. Transmit the response back to the client in OpenAI format.
             return self._response_adapter.format_response_to_client(
                 response, model_route.endpoint.provider_type
             )

src/codegate/muxing/rulematcher.py

@@ -1,9 +1,17 @@
 import copy
 from abc import ABC, abstractmethod
 from asyncio import Lock
-from typing import List, Optional
+from typing import Dict, List, Optional
 
+import structlog
+
+from codegate.clients.clients import ClientType
 from codegate.db import models as db_models
+from codegate.extract_snippets.body_extractor import BodyCodeSnippetExtractorError
+from codegate.extract_snippets.factory import BodyCodeExtractorFactory
+from codegate.muxing import models as mux_models
+
+logger = structlog.get_logger("codegate")
 
 _muxrules_sgtn = None
 
@@ -40,11 +48,12 @@ def __init__(
 class MuxingRuleMatcher(ABC):
     """Base class for matching muxing rules."""
 
-    def __init__(self, route: ModelRoute):
+    def __init__(self, route: ModelRoute, matcher_blob: str):
         self._route = route
+        self._matcher_blob = matcher_blob
 
     @abstractmethod
-    def match(self, thing_to_match) -> bool:
+    def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
         """Return True if the rule matches the thing_to_match."""
         pass
 
@@ -61,23 +70,69 @@ class MuxingMatcherFactory:
     def create(mux_rule: db_models.MuxRule, route: ModelRoute) -> MuxingRuleMatcher:
         """Create a muxing matcher for the given endpoint and model."""
 
-        factory = {
-            "catch_all": CatchAllMuxingRuleMatcher,
+        factory: Dict[mux_models.MuxMatcherType, MuxingRuleMatcher] = {
+            mux_models.MuxMatcherType.catch_all: CatchAllMuxingRuleMatcher,
+            mux_models.MuxMatcherType.filename_match: FileMuxingRuleMatcher,
+            mux_models.MuxMatcherType.request_type_match: RequestTypeMuxingRuleMatcher,
         }
 
         try:
-            return factory[mux_rule.matcher_type](route)
+            # Initialize the MuxingRuleMatcher
+            return factory[mux_rule.matcher_type](route, mux_rule.matcher_blob)
         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:
+    def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
         return True
 
 
+class FileMuxingRuleMatcher(MuxingRuleMatcher):
+    """A file muxing rule matcher."""
+
+    def _extract_request_filenames(self, detected_client: ClientType, data: dict) -> set[str]:
+        """
+        Extract filenames from the request data.
+        """
+        try:
+            body_extractor = BodyCodeExtractorFactory.create_snippet_extractor(detected_client)
+            return body_extractor.extract_unique_filenames(data)
+        except BodyCodeSnippetExtractorError as e:
+            logger.error(f"Error extracting filenames from request: {e}")
+            return set()
+
+    def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
+        """
+        Retun True if there is a filename in the request that matches the matcher_blob.
+        The matcher_blob is either an extension (e.g. .py) or a filename (e.g. main.py).
+        """
+        # If there is no matcher_blob, we don't match
+        if not self._matcher_blob:
+            return False
+        filenames_to_match = self._extract_request_filenames(
+            thing_to_match.client_type, thing_to_match.body
+        )
+        return any(self._matcher_blob in filename for filename in filenames_to_match)
+
+
+class RequestTypeMuxingRuleMatcher(MuxingRuleMatcher):
+    """A catch all muxing rule matcher."""
+
+    def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
+        """
+        Return True if the request type matches the matcher_blob.
+        The matcher_blob is either "fim" or "chat".
+        """
+        # If there is no matcher_blob, we don't match
+        if not self._matcher_blob:
+            return False
+        incoming_request_type = "fim" if thing_to_match.is_fim_request else "chat"
+        return self._matcher_blob == incoming_request_type
+
+
 class MuxingRulesinWorkspaces:
     """A thread safe dictionary to store the muxing rules in workspaces."""
 
@@ -111,7 +166,9 @@ async def get_registries(self) -> List[str]:
         async with self._lock:
             return list(self._ws_rules.keys())
 
-    async def get_match_for_active_workspace(self, thing_to_match) -> Optional[ModelRoute]:
+    async def get_match_for_active_workspace(
+        self, thing_to_match: mux_models.ThingToMatchMux
+    ) -> Optional[ModelRoute]:
         """Get the first match for the given thing_to_match."""
 
         # We iterate over all the rules and return the first match

src/codegate/pipeline/secrets/secrets.py

@@ -2,12 +2,12 @@
 from abc import abstractmethod
 from typing import List, Optional, Tuple
 
-from codegate.extract_snippets.factory import MessageCodeExtractorFactory
 import structlog
 from litellm import ChatCompletionRequest, ChatCompletionSystemMessage, ModelResponse
 from litellm.types.utils import Delta, StreamingChoices
 
 from codegate.config import Config
+from codegate.extract_snippets.factory import MessageCodeExtractorFactory
 from codegate.pipeline.base import (
     AlertSeverity,
     CodeSnippet,

src/codegate/providers/anthropic/provider.py

@@ -11,6 +11,7 @@
 from codegate.providers.anthropic.adapter import AnthropicInputNormalizer, AnthropicOutputNormalizer
 from codegate.providers.anthropic.completion_handler import AnthropicCompletion
 from codegate.providers.base import BaseProvider, ModelFetchError
+from codegate.providers.fim_analyzer import FIMAnalyzer
 from codegate.providers.litellmshim import anthropic_stream_generator
 
 
@@ -57,10 +58,9 @@ async def process_request(
         self,
         data: dict,
         api_key: str,
-        request_url_path: str,
+        is_fim_request: bool,
         client_type: ClientType,
     ):
-        is_fim_request = self._is_fim_request(request_url_path, data)
         try:
             stream = await self.complete(data, api_key, is_fim_request, client_type)
         except Exception as e:
@@ -98,10 +98,11 @@ async def create_message(
 
             body = await request.body()
             data = json.loads(body)
+            is_fim_request = FIMAnalyzer.is_fim_request(request.url.path, data)
 
             return await self.process_request(
                 data,
                 x_api_key,
-                request.url.path,
+                is_fim_request,
                 request.state.detected_client,
             )

src/codegate/providers/base.py

@@ -79,7 +79,7 @@ async def process_request(
         self,
         data: dict,
         api_key: str,
-        request_url_path: str,
+        is_fim_request: bool,
         client_type: ClientType,
     ):
         pass
@@ -173,61 +173,6 @@ async def _run_input_pipeline(
 
         return result
 
-    def _is_fim_request_url(self, request_url_path: str) -> bool:
-        """
-        Checks the request URL to determine if a request is FIM or chat completion.
-        Used by: llama.cpp
-        """
-        # Evaluate first a larger substring.
-        if request_url_path.endswith("/chat/completions"):
-            return False
-
-        # /completions is for OpenAI standard. /api/generate is for ollama.
-        if request_url_path.endswith("/completions") or request_url_path.endswith("/api/generate"):
-            return True
-
-        return False
-
-    def _is_fim_request_body(self, data: Dict) -> bool:
-        """
-        Determine from the raw incoming data if it's a FIM request.
-        Used by: OpenAI and Anthropic
-        """
-        messages = data.get("messages", [])
-        if not messages:
-            return False
-
-        first_message_content = messages[0].get("content")
-        if first_message_content is None:
-            return False
-
-        fim_stop_sequences = ["</COMPLETION>", "<COMPLETION>", "</QUERY>", "<QUERY>"]
-        if isinstance(first_message_content, str):
-            msg_prompt = first_message_content
-        elif isinstance(first_message_content, list):
-            msg_prompt = first_message_content[0].get("text", "")
-        else:
-            logger.warning(f"Could not determine if message was FIM from data: {data}")
-            return False
-        return all([stop_sequence in msg_prompt for stop_sequence in fim_stop_sequences])
-
-    def _is_fim_request(self, request_url_path: str, data: Dict) -> bool:
-        """
-        Determine if the request is FIM by the URL or the data of the request.
-        """
-        # first check if we are in specific tools to discard FIM
-        prompt = data.get("prompt", "")
-        tools = ["cline", "kodu", "open interpreter"]
-        for tool in tools:
-            if tool in prompt.lower():
-                #  those tools can never be FIM
-                return False
-        # Avoid more expensive inspection of body by just checking the URL.
-        if self._is_fim_request_url(request_url_path):
-            return True
-
-        return self._is_fim_request_body(data)
-
     async def _cleanup_after_streaming(
         self, stream: AsyncIterator[ModelResponse], context: PipelineContext
     ) -> AsyncIterator[ModelResponse]: