add a persona matcher to the system prompt and unit tests

Author: aponcedeleonch

src/codegate/muxing/models.py

@@ -32,8 +32,10 @@ class MuxMatcherType(str, Enum):
     fim_filename = "fim_filename"
     # Match based on chat request type. It will match if the request type is chat
     chat_filename = "chat_filename"
-    # Match the request content to the persona description
+    # Match the user messages to the persona description
     persona_description = "persona_description"
+    # Match the system prompt to the persona description
+    sys_prompt_persona_desc = "sys_prompt_persona_desc"
 
 
 class MuxRule(pydantic.BaseModel):

src/codegate/muxing/rulematcher.py

@@ -83,7 +83,8 @@ def create(db_mux_rule: db_models.MuxRule, route: ModelRoute) -> MuxingRuleMatch
             mux_models.MuxMatcherType.filename_match: FileMuxingRuleMatcher,
             mux_models.MuxMatcherType.fim_filename: RequestTypeAndFileMuxingRuleMatcher,
             mux_models.MuxMatcherType.chat_filename: RequestTypeAndFileMuxingRuleMatcher,
-            mux_models.MuxMatcherType.persona_description: PersonaDescriptionMuxingRuleMatcher,
+            mux_models.MuxMatcherType.persona_description: UserMsgsPersonaDescMuxMatcher,
+            mux_models.MuxMatcherType.sys_prompt_persona_desc: SysPromptPersonaDescMuxMatcher,
         }
 
         try:
@@ -173,12 +174,42 @@ async def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
         return is_rule_matched
 
 
-class PersonaDescriptionMuxingRuleMatcher(MuxingRuleMatcher):
+class PersonaDescMuxMatcher(MuxingRuleMatcher):
     """Muxing rule to match the request content to a persona description."""
 
-    def _get_user_messages_from_body(self, body: Dict) -> List[str]:
+    @abstractmethod
+    def _get_queries_for_persona_match(self, body: Dict) -> List[str]:
+        """
+        Get the queries to use for persona matching.
+        """
+        pass
+
+    async def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
+        """
+        Return True if the matcher is the persona description matched with the queries.
+
+        The queries are extracted from the body and will depend on the type of matcher.
+        1. UserMessagesPersonaDescMuxMatcher: Extracts queries from the user messages in the body.
+        2. SysPromptPersonaDescMuxMatcher: Extracts queries from the system messages in the body.
+        """
+        queries = self._get_queries_for_persona_match(thing_to_match.body)
+        if not queries:
+            return False
+
+        persona_manager = PersonaManager()
+        is_persona_matched = await persona_manager.check_persona_match(
+            persona_name=self._mux_rule.matcher, queries=queries
+        )
+        if is_persona_matched:
+            logger.info("Persona rule matched", persona=self._mux_rule.matcher)
+        return is_persona_matched
+
+
+class UserMsgsPersonaDescMuxMatcher(PersonaDescMuxMatcher):
+
+    def _get_queries_for_persona_match(self, body: Dict) -> List[str]:
         """
-        Get the user messages from the body to use as queries.
+        Get the queries from the user messages in the body.
         """
         user_messages = []
         for msg in body.get("messages", []):
@@ -194,22 +225,34 @@ def _get_user_messages_from_body(self, body: Dict) -> List[str]:
                     user_messages.append(msgs_content)
         return user_messages
 
-    async def match(self, thing_to_match: mux_models.ThingToMatchMux) -> bool:
+
+class SysPromptPersonaDescMuxMatcher(PersonaDescMuxMatcher):
+
+    def _get_queries_for_persona_match(self, body: Dict) -> List[str]:
         """
-        Return True if the matcher is the persona description matched with the
-        user messages.
+        Get the queries from the system messages in the body.
         """
-        user_messages = self._get_user_messages_from_body(thing_to_match.body)
-        if not user_messages:
-            return False
+        system_messages = []
+        for msg in body.get("messages", []):
+            if msg.get("role", "") in ["system", "developer"]:
+                msgs_content = msg.get("content")
+                if not msgs_content:
+                    continue
+                if isinstance(msgs_content, list):
+                    for msg_content in msgs_content:
+                        if msg_content.get("type", "") == "text":
+                            system_messages.append(msg_content.get("text", ""))
+                elif isinstance(msgs_content, str):
+                    system_messages.append(msgs_content)
 
-        persona_manager = PersonaManager()
-        is_persona_matched = await persona_manager.check_persona_match(
-            persona_name=self._mux_rule.matcher, queries=user_messages
-        )
-        if is_persona_matched:
-            logger.info("Persona rule matched", persona=self._mux_rule.matcher)
-        return is_persona_matched
+        # Handling the anthropic system prompt
+        anthropic_sys_prompt = body.get("system")
+        if anthropic_sys_prompt:
+            system_messages.append(anthropic_sys_prompt)
+
+        # In an ideal world, the length of system_messages should be 1. Returnin the list
+        # to handle any edge cases and to not break parent function's signature.
+        return system_messages
 
 
 class MuxingRulesinWorkspaces:

tests/muxing/test_rulematcher.py

@@ -1,4 +1,5 @@
-from unittest.mock import MagicMock
+from typing import Dict, List
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -151,6 +152,100 @@ async def test_request_file_matcher(
     )
 
 
+# We mock PersonaManager because it's tested in /tests/persona/test_manager.py
+MOCK_PERSONA_MANAGER = AsyncMock()
+MOCK_PERSONA_MANAGER.check_persona_match.return_value = True
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "body, expected_queries",
+    [
+        ({"messages": [{"role": "system", "content": "Youre helpful"}]}, []),
+        ({"messages": [{"role": "user", "content": "hello"}]}, ["hello"]),
+        (
+            {"messages": [{"role": "user", "content": [{"type": "text", "text": "hello_dict"}]}]},
+            ["hello_dict"],
+        ),
+    ],
+)
+async def test_user_msgs_persona_desc_matcher(body: Dict, expected_queries: List[str]):
+    mux_rule = mux_models.MuxRule(
+        provider_id="1",
+        model="fake-gpt",
+        matcher_type="user_messages_persona_desc",
+        matcher="foo_persona",
+    )
+    muxing_rule_matcher = rulematcher.UserMsgsPersonaDescMuxMatcher(mocked_route_openai, mux_rule)
+
+    mocked_thing_to_match = mux_models.ThingToMatchMux(
+        body=body,
+        url_request_path="/chat/completions",
+        is_fim_request=False,
+        client_type="generic",
+    )
+
+    resulting_queries = muxing_rule_matcher._get_queries_for_persona_match(body)
+    assert set(resulting_queries) == set(expected_queries)
+
+    with patch("codegate.muxing.rulematcher.PersonaManager", return_value=MOCK_PERSONA_MANAGER):
+        result = await muxing_rule_matcher.match(mocked_thing_to_match)
+
+    if expected_queries:
+        assert result is True
+    else:
+        assert result is False
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "body, expected_queries",
+    [
+        ({"messages": [{"role": "system", "content": "Youre helpful"}]}, ["Youre helpful"]),
+        ({"messages": [{"role": "user", "content": "hello"}]}, []),
+        (
+            {
+                "messages": [
+                    {"role": "system", "content": "Youre helpful"},
+                    {"role": "user", "content": "hello"},
+                ]
+            },
+            ["Youre helpful"],
+        ),
+        (
+            {"messages": [{"role": "user", "content": "hello"}], "system": "Anthropic system"},
+            ["Anthropic system"],
+        ),
+    ],
+)
+async def test_sys_prompt_persona_desc_matcher(body: Dict, expected_queries: List[str]):
+    mux_rule = mux_models.MuxRule(
+        provider_id="1",
+        model="fake-gpt",
+        matcher_type="sys_prompt_persona_desc",
+        matcher="foo_persona",
+    )
+    muxing_rule_matcher = rulematcher.SysPromptPersonaDescMuxMatcher(mocked_route_openai, mux_rule)
+
+    mocked_thing_to_match = mux_models.ThingToMatchMux(
+        body=body,
+        url_request_path="/chat/completions",
+        is_fim_request=False,
+        client_type="generic",
+    )
+
+    resulting_queries = muxing_rule_matcher._get_queries_for_persona_match(body)
+    assert set(resulting_queries) == set(expected_queries)
+
+    with patch("codegate.muxing.rulematcher.PersonaManager", return_value=MOCK_PERSONA_MANAGER):
+        result = await muxing_rule_matcher.match(mocked_thing_to_match)
+
+    if expected_queries:
+        assert result is True
+    else:
+        assert result is False
+
+
 @pytest.mark.parametrize(
     "matcher_type, expected_class",
     [
@@ -159,8 +254,12 @@ async def test_request_file_matcher(
         (mux_models.MuxMatcherType.fim_filename, rulematcher.RequestTypeAndFileMuxingRuleMatcher),
         (mux_models.MuxMatcherType.chat_filename, rulematcher.RequestTypeAndFileMuxingRuleMatcher),
         (
-            mux_models.MuxMatcherType.persona_description,
-            rulematcher.PersonaDescriptionMuxingRuleMatcher,
+            mux_models.MuxMatcherType.user_messages_persona_desc,
+            rulematcher.UserMsgsPersonaDescMuxMatcher,
+        ),
+        (
+            mux_models.MuxMatcherType.sys_prompt_persona_desc,
+            rulematcher.SysPromptPersonaDescMuxMatcher,
         ),
         ("invalid_matcher", None),
     ],