Merge branch 'main' into issue-875

Author: yrobla

api/openapi.json

@@ -1581,7 +1581,9 @@
       "MuxMatcherType": {
         "type": "string",
         "enum": [
-          "catch_all"
+          "catch_all",
+          "filename_match",
+          "request_type_match"
         ],
         "title": "MuxMatcherType",
         "description": "Represents the different types of matchers we support."

src/codegate/clients/clients.py

@@ -12,3 +12,4 @@ class ClientType(Enum):
     COPILOT = "copilot"  # Copilot client
     OPEN_INTERPRETER = "open_interpreter"  # Open Interpreter client
     AIDER = "aider"  # Aider client
+    CONTINUE = "continue"  # Continue client

src/codegate/clients/detector.py

@@ -160,6 +160,24 @@ def client_name(self) -> ClientType:
         return ClientType.OPEN_INTERPRETER
 
 
+class ContinueDetector(BaseClientDetector):
+    """
+    Detector for Continue client based on message content
+    """
+
+    def __init__(self):
+        super().__init__()
+        # This is a hack that really only detects Continue with DeepSeek
+        # we should get a header or user agent for this (upstream PR pending)
+        self.content_detector = ContentDetector(
+            "You are an AI programming assistant, utilizing the DeepSeek Coder model"
+        )
+
+    @property
+    def client_name(self) -> ClientType:
+        return ClientType.CONTINUE
+
+
 class CopilotDetector(BaseClientDetector):
     """
     Detector for Copilot client based on user agent
@@ -191,6 +209,7 @@ def __init__(self):
             KoduDetector(),
             OpenInterpreter(),
             CopilotDetector(),
+            ContinueDetector(),
         ]
 
     def __call__(self, func):

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,45 @@ 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]
+            logger.info(
+                "Muxing request routed to destination provider",
+                model=model_route.model.name,
+                provider_type=model_route.endpoint.provider_type,
+                provider_name=model_route.endpoint.name,
+            )
 
-            # 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,82 @@ 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:
+        logger.info("Catch all rule matched")
         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
+        )
+        is_filename_match = any(self._matcher_blob in filename for filename in filenames_to_match)
+        if is_filename_match:
+            logger.info(
+                "Filename rule matched", filenames=filenames_to_match, matcher=self._matcher_blob
+            )
+        return is_filename_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"
+        is_request_type_match = self._matcher_blob == incoming_request_type
+        if is_request_type_match:
+            logger.info(
+                "Request type rule matched",
+                matcher=self._matcher_blob,
+                request_type=incoming_request_type,
+            )
+        return is_request_type_match
+
+
 class MuxingRulesinWorkspaces:
     """A thread safe dictionary to store the muxing rules in workspaces."""
 
@@ -111,7 +179,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/cli/cli.py

@@ -13,6 +13,8 @@
 )
 from codegate.pipeline.cli.commands import CustomInstructions, Version, Workspace
 
+codegate_regex = re.compile(r"^codegate(?:\s+(.*))?", re.IGNORECASE)
+
 HELP_TEXT = """
 ## CodeGate CLI\n
 **Usage**: `codegate [-h] <command> [args]`\n
@@ -77,6 +79,22 @@ def _get_cli_from_open_interpreter(last_user_message_str: str) -> Optional[re.Ma
     return re.match(r"^codegate\s*(.*?)\s*$", last_user_block, re.IGNORECASE)
 
 
+def _get_cli_from_continue(last_user_message_str: str) -> Optional[re.Match[str]]:
+    """
+    Continue sends a differently formatted message to the CLI if DeepSeek is used
+    """
+    deepseek_match = re.search(
+        r"utilizing the DeepSeek Coder model.*?### Instruction:\s*codegate\s+(.*?)\s*### Response:",
+        last_user_message_str,
+        re.DOTALL | re.IGNORECASE,
+    )
+    if deepseek_match:
+        command = deepseek_match.group(1).strip()
+        return re.match(r"^(.*?)$", command)  # This creates a match object with the command
+
+    return codegate_regex.match(last_user_message_str)
+
+
 class CodegateCli(PipelineStep):
     """Pipeline step that handles codegate cli."""
 
@@ -110,12 +128,14 @@ async def process(
         if last_user_message is not None:
             last_user_message_str, _ = last_user_message
             last_user_message_str = last_user_message_str.strip()
-            codegate_regex = re.compile(r"^codegate(?:\s+(.*))?", re.IGNORECASE)
 
+            # Check client-specific matchers first
             if context.client in [ClientType.CLINE, ClientType.KODU]:
                 match = _get_cli_from_cline(codegate_regex, last_user_message_str)
             elif context.client in [ClientType.OPEN_INTERPRETER]:
                 match = _get_cli_from_open_interpreter(last_user_message_str)
+            elif context.client in [ClientType.CONTINUE]:
+                match = _get_cli_from_continue(last_user_message_str)
             else:
                 # Check if "codegate" is the first word in the message
                 match = codegate_regex.match(last_user_message_str)