updated matcher types to preserve old types

Author: aponcedeleonch

migrations/versions/2025_02_19_1452-5e5cd2288147_update_matcher_types.py

@@ -26,24 +26,17 @@ def upgrade() -> None:
     op.execute(
         """
         UPDATE muxes
-        SET matcher_type = 'fim', matcher_blob = ''
+        SET matcher_type = 'fim_filename', matcher_blob = ''
         WHERE matcher_type = 'request_type_match' AND matcher_blob = 'fim';
         """
     )
     op.execute(
         """
         UPDATE muxes
-        SET matcher_type = 'chat', matcher_blob = ''
+        SET matcher_type = 'chat_filename', matcher_blob = ''
         WHERE matcher_type = 'request_type_match' AND matcher_blob = 'chat';
         """
     )
-    op.execute(
-        """
-        UPDATE muxes
-        SET matcher_type = 'catch_all'
-        WHERE matcher_type = 'filename_match' AND matcher_blob != '';
-        """
-    )
 
     # Finish transaction
     op.execute("COMMIT;")
@@ -67,13 +60,6 @@ def downgrade() -> None:
         WHERE matcher_type = 'chat';
         """
     )
-    op.execute(
-        """
-        UPDATE muxes
-        SET matcher_type = 'filename_match', matcher_blob = 'catch_all'
-        WHERE matcher_type = 'catch_all';
-        """
-    )
 
     # Finish transaction
     op.execute("COMMIT;")

src/codegate/muxing/models.py

@@ -24,10 +24,13 @@ 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 fim request type. It will match if the request type is fim
-    fim = "fim"
+    fim_filename = "fim_filename"
     # Match based on chat request type. It will match if the request type is chat
-    chat = "chat"
+    chat_filename = "chat_filename"
 
 
 class MuxRule(pydantic.BaseModel):

src/codegate/muxing/rulematcher.py

@@ -78,8 +78,9 @@ def create(db_mux_rule: db_models.MuxRule, route: ModelRoute) -> MuxingRuleMatch
 
         factory: Dict[mux_models.MuxMatcherType, MuxingRuleMatcher] = {
             mux_models.MuxMatcherType.catch_all: RequestTypeAndFileMuxingRuleMatcher,
-            mux_models.MuxMatcherType.fim: RequestTypeAndFileMuxingRuleMatcher,
-            mux_models.MuxMatcherType.chat: RequestTypeAndFileMuxingRuleMatcher,
+            mux_models.MuxMatcherType.filename_match: FileMuxingRuleMatcher,
+            mux_models.MuxMatcherType.fim_filename: RequestTypeAndFileMuxingRuleMatcher,
+            mux_models.MuxMatcherType.chat_filename: RequestTypeAndFileMuxingRuleMatcher,
         }
 
         try:
@@ -90,7 +91,16 @@ def create(db_mux_rule: db_models.MuxRule, route: ModelRoute) -> MuxingRuleMatch
             raise ValueError(f"Unknown matcher type: {mux_rule.matcher_type}")
 
 
-class RequestTypeAndFileMuxingRuleMatcher(MuxingRuleMatcher):
+class CatchAllMuxingRuleMatcher(MuxingRuleMatcher):
+    """A catch all muxing rule matcher."""
+
+    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]:
         """
@@ -119,14 +129,26 @@ def _is_matcher_in_filenames(self, detected_client: ClientType, data: dict) -> b
         )
         return is_filename_match
 
+    def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
+        """
+        Return True if the matcher is in one of the request filenames.
+        """
+        is_rule_matched = self._is_matcher_in_filenames(
+            thing_to_match.client_type, thing_to_match.body
+        )
+        if is_rule_matched:
+            logger.info("Filename rule matched", matcher=self._mux_rule.matcher)
+        return is_rule_matched
+
+
+class RequestTypeAndFileMuxingRuleMatcher(FileMuxingRuleMatcher):
+    """A request type and file muxing rule matcher."""
+
     def _is_request_type_match(self, is_fim_request: bool) -> bool:
         """
         Check if the request type matches the MuxMatcherType.
         """
-        # Catch all rule matches both chat and FIM requests
-        if self._mux_rule.matcher_type == mux_models.MuxMatcherType.catch_all:
-            return True
-        incoming_request_type = "fim" if is_fim_request else "chat"
+        incoming_request_type = "fim_filename" if is_fim_request else "chat_filename"
         if incoming_request_type == self._mux_rule.matcher_type:
             return True
         return False

tests/muxing/test_rulematcher.py

@@ -24,6 +24,55 @@
 )
 
 
+def test_catch_all():
+    mux_rule = mux_models.MuxRule(
+        provider_id="1",
+        model="fake-gpt",
+        matcher_type="catch_all",
+        matcher="",
+    )
+    muxing_rule_matcher = rulematcher.CatchAllMuxingRuleMatcher(mocked_route_openai, mux_rule)
+    assert muxing_rule_matcher.match(None) is True
+
+
+@pytest.mark.parametrize(
+    "matcher, filenames_to_match, expected_bool",
+    [
+        (None, [], True),  # Empty filenames and no blob
+        (None, ["main.py"], True),  # Empty blob should match
+        (".py", ["main.py"], True),  # Extension match
+        ("main.py", ["main.py"], True),  # Full name match
+        (".py", ["main.py", "test.py"], True),  # Extension match
+        ("main.py", ["main.py", "test.py"], True),  # Full name match
+        ("main.py", ["test.py"], False),  # Full name no match
+        (".js", ["main.py", "test.py"], False),  # Extension no match
+        (".ts", ["main.tsx", "test.tsx"], False),  # Extension no match
+    ],
+)
+def test_file_matcher(
+    matcher,
+    filenames_to_match,
+    expected_bool,
+):
+    mux_rule = mux_models.MuxRule(
+        provider_id="1",
+        model="fake-gpt",
+        matcher_type="filename_match",
+        matcher=matcher,
+    )
+    muxing_rule_matcher = rulematcher.FileMuxingRuleMatcher(mocked_route_openai, mux_rule)
+    # We mock the _extract_request_filenames method to return a list of filenames
+    # The logic to get the correct filenames from snippets is tested in /tests/extract_snippets
+    muxing_rule_matcher._extract_request_filenames = MagicMock(return_value=filenames_to_match)
+    mocked_thing_to_match = mux_models.ThingToMatchMux(
+        body={},
+        url_request_path="/chat/completions",
+        is_fim_request=False,
+        client_type="generic",
+    )
+    assert muxing_rule_matcher.match(mocked_thing_to_match) is expected_bool
+
+
 @pytest.mark.parametrize(
     "matcher, filenames_to_match, expected_bool_filenames",
     [
@@ -41,15 +90,13 @@
 @pytest.mark.parametrize(
     "is_fim_request, matcher_type, expected_bool_request",
     [
-        (False, "fim", False),  # No match
-        (True, "fim", True),  # Match
-        (False, "chat", True),  # Match
-        (True, "chat", False),  # No match
-        (True, "catch_all", True),  # Match
-        (False, "catch_all", True),  # Match
+        (False, "fim_filename", False),  # No match
+        (True, "fim_filename", True),  # Match
+        (False, "chat_filename", True),  # Match
+        (True, "chat_filename", False),  # No match
     ],
 )
-def test_file_matcher(
+def test_request_file_matcher(
     matcher,
     filenames_to_match,
     expected_bool_filenames,