convert muxing rule match to coroutines

Author: aponcedeleonch

src/codegate/muxing/rulematcher.py

@@ -61,7 +61,7 @@ def __init__(self, route: ModelRoute, mux_rule: mux_models.MuxRule):
         self._mux_rule = mux_rule
 
     @abstractmethod
-    def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
+    async def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
         """Return True if the rule matches the thing_to_match."""
         pass
 
@@ -97,7 +97,7 @@ def create(db_mux_rule: db_models.MuxRule, route: ModelRoute) -> MuxingRuleMatch
 class CatchAllMuxingRuleMatcher(MuxingRuleMatcher):
     """A catch all muxing rule matcher."""
 
-    def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
+    async def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
         logger.info("Catch all rule matched")
         return True
 
@@ -132,7 +132,7 @@ 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:
+    async def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
         """
         Return True if the matcher is in one of the request filenames.
         """
@@ -156,7 +156,7 @@ def _is_request_type_match(self, is_fim_request: bool) -> bool:
             return True
         return False
 
-    def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
+    async def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
         """
         Return True if the matcher is in one of the request filenames and
         if the request type matches the MuxMatcherType.
@@ -194,7 +194,7 @@ def _get_user_messages_from_body(self, body: Dict) -> List[str]:
                     user_messages.append(msgs_content)
         return user_messages
 
-    def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
+    async def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
         """
         Return True if the matcher is the persona description matched with the
         user messages.
@@ -204,14 +204,11 @@ def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
             return False
 
         persona_manager = PersonaManager()
-        is_persona_matched = persona_manager.check_persona_match(
+        is_persona_matched = await persona_manager.check_persona_match(
             persona_name=self._mux_rule.matcher, queries=user_messages
         )
-        logger.info(
-            "Persona rule matched",
-            matcher=self._mux_rule.matcher,
-            is_persona_matched=is_persona_matched,
-        )
+        if is_persona_matched:
+            logger.info("Persona rule matched", persona=self._mux_rule.matcher)
         return is_persona_matched
 
 
@@ -258,7 +255,7 @@ async def get_match_for_active_workspace(
         try:
             rules = await self.get_ws_rules(self._active_workspace)
             for rule in rules:
-                if rule.match(thing_to_match):
+                if await rule.match(thing_to_match):
                     return rule.destination()
             return None
         except KeyError:

tests/muxing/test_rulematcher.py

@@ -24,6 +24,7 @@
 )
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize(
     "matcher_blob, thing_to_match",
     [
@@ -40,12 +41,13 @@
         ),
     ],
 )
-def test_catch_all(matcher_blob, thing_to_match):
+async def test_catch_all(matcher_blob, thing_to_match):
     muxing_rule_matcher = rulematcher.CatchAllMuxingRuleMatcher(mocked_route_openai, matcher_blob)
     # It should always match
-    assert muxing_rule_matcher.match(thing_to_match) is True
+    assert await muxing_rule_matcher.match(thing_to_match) is True
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize(
     "matcher, filenames_to_match, expected_bool",
     [
@@ -60,7 +62,7 @@ def test_catch_all(matcher_blob, thing_to_match):
         ("*.ts", ["main.tsx", "test.tsx"], False),  # Extension no match
     ],
 )
-def test_file_matcher(
+async def test_file_matcher(
     matcher,
     filenames_to_match,
     expected_bool,
@@ -81,9 +83,10 @@ def test_file_matcher(
         is_fim_request=False,
         client_type="generic",
     )
-    assert muxing_rule_matcher.match(mocked_thing_to_match) is expected_bool
+    assert await muxing_rule_matcher.match(mocked_thing_to_match) is expected_bool
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize(
     "matcher, filenames_to_match, expected_bool_filenames",
     [
@@ -107,7 +110,7 @@ def test_file_matcher(
         (True, "chat_filename", False),  # No match
     ],
 )
-def test_request_file_matcher(
+async def test_request_file_matcher(
     matcher,
     filenames_to_match,
     expected_bool_filenames,
@@ -143,7 +146,7 @@ def test_request_file_matcher(
         )
         is expected_bool_filenames
     )
-    assert muxing_rule_matcher.match(mocked_thing_to_match) is (
+    assert await muxing_rule_matcher.match(mocked_thing_to_match) is (
         expected_bool_request and expected_bool_filenames
     )
 
@@ -155,6 +158,10 @@ def test_request_file_matcher(
         (mux_models.MuxMatcherType.filename_match, rulematcher.FileMuxingRuleMatcher),
         (mux_models.MuxMatcherType.fim_filename, rulematcher.RequestTypeAndFileMuxingRuleMatcher),
         (mux_models.MuxMatcherType.chat_filename, rulematcher.RequestTypeAndFileMuxingRuleMatcher),
+        (
+            mux_models.MuxMatcherType.persona_description,
+            rulematcher.PersonaDescriptionMuxingRuleMatcher,
+        ),
         ("invalid_matcher", None),
     ],
 )