unique_toolkit 1.7.0__py3-none-any.whl → 1.8.1__py3-none-any.whl

unique_toolkit/agentic/tools/a2a/postprocessing/postprocessor.py

@@ -1,19 +1,21 @@
 import logging
-import re
 from typing import TypedDict, override
 
 import unique_sdk
 
 from unique_toolkit.agentic.postprocessor.postprocessor_manager import Postprocessor
-from unique_toolkit.agentic.tools.a2a.config import (
-    ResponseDisplayMode,
-    SubAgentToolDisplayConfig,
+from unique_toolkit.agentic.tools.a2a.postprocessing._display import (
+    _build_sub_agent_answer_display,
+    _remove_sub_agent_answer_from_text,
 )
-from unique_toolkit.agentic.tools.a2a.postprocessing.display import (
-    build_sub_agent_answer_display,
-    remove_sub_agent_answer_from_text,
+from unique_toolkit.agentic.tools.a2a.postprocessing._utils import (
+    _replace_references_in_text,
 )
-from unique_toolkit.agentic.tools.a2a.service import SubAgentTool
+from unique_toolkit.agentic.tools.a2a.postprocessing.config import (
+    SubAgentDisplayConfig,
+    SubAgentResponseDisplayMode,
+)
+from unique_toolkit.agentic.tools.a2a.tool import SubAgentTool
 from unique_toolkit.content.schemas import ContentReference
 from unique_toolkit.language_model.schemas import LanguageModelStreamResponse
 
@@ -23,14 +25,14 @@ SpaceMessage = unique_sdk.Space.Message
 
 
 class _SubAgentMessageInfo(TypedDict):
-    text: str | None
+    text: str
     references: list[unique_sdk.Space.Reference]
 
 
 class _SubAgentToolInfo(TypedDict):
     display_name: str
-    display_config: SubAgentToolDisplayConfig
-    responses: list[_SubAgentMessageInfo]
+    display_config: SubAgentDisplayConfig
+    responses: dict[int, _SubAgentMessageInfo]
 
 
 class SubAgentResponsesPostprocessor(Postprocessor):
@@ -38,93 +40,79 @@ class SubAgentResponsesPostprocessor(Postprocessor):
         self,
         user_id: str,
         company_id: str,
-        agent_chat_id: str,
-        sub_agent_tools: list[SubAgentTool],
-    ):
+        main_agent_chat_id: str,
+    ) -> None:
         super().__init__(name=self.__class__.__name__)
 
         self._user_id = user_id
         self._company_id = company_id
-
-        self._agent_chat_id = agent_chat_id
+        self._main_agent_chat_id = main_agent_chat_id
 
         self._assistant_id_to_tool_info: dict[str, _SubAgentToolInfo] = {}
-
-        for sub_agent_tool in sub_agent_tools:
-            sub_agent_tool.subscribe(self)
-
-            self._assistant_id_to_tool_info[sub_agent_tool.config.assistant_id] = (
-                _SubAgentToolInfo(
-                    display_config=sub_agent_tool.config.response_display_config,
-                    display_name=sub_agent_tool.display_name(),
-                    responses=[],
-                )
-            )
-
-        self._sub_agent_message = None
+        self._main_agent_message: SpaceMessage | None = None
 
     @override
     async def run(self, loop_response: LanguageModelStreamResponse) -> None:
-        self._sub_agent_message = await unique_sdk.Space.get_latest_message_async(
+        self._main_agent_message = await unique_sdk.Space.get_latest_message_async(
             user_id=self._user_id,
             company_id=self._company_id,
-            chat_id=self._agent_chat_id,
+            chat_id=self._main_agent_chat_id,
         )
 
     @override
     def apply_postprocessing_to_response(
         self, loop_response: LanguageModelStreamResponse
     ) -> bool:
-        logger.info("Adding sub agent responses to the response")
-
-        # Get responses to display
-        displayed = {}
-        for assistant_id, tool_info in self._assistant_id_to_tool_info.items():
-            display_mode = tool_info["display_config"].mode
-
-            if len(tool_info["responses"]) == 0:
-                logger.warning(
-                    "No response from assistant %s",
-                    assistant_id,
-                )
-                continue
-
-            if display_mode != ResponseDisplayMode.HIDDEN:
-                displayed[assistant_id] = tool_info["responses"]
+        logger.info("Prepending sub agent responses to the main agent response")
+
+        if len(self._assistant_id_to_tool_info) == 0 or all(
+            len(tool_info["responses"]) == 0
+            for tool_info in self._assistant_id_to_tool_info.values()
+        ):
+            logger.info("No sub agent responses to prepend")
+            return False
+
+        if self._main_agent_message is None:
+            raise ValueError(
+                "Main agent message is not set, the `run` method must be called first"
+            )
 
         existing_refs = {
             ref.source_id: ref.sequence_number
             for ref in loop_response.message.references
         }
-        _consolidate_references_in_place(displayed, existing_refs)
 
-        modified = len(displayed) > 0
-        for assistant_id, messages in reversed(displayed.items()):
-            for i in reversed(range(len(messages))):
-                message = messages[i]
+        _consolidate_references_in_place(
+            list(self._assistant_id_to_tool_info.values()), existing_refs
+        )
+
+        answers = []
+        for assistant_id in self._assistant_id_to_tool_info.keys():
+            messages = self._assistant_id_to_tool_info[assistant_id]["responses"]
+
+            for sequence_number in sorted(messages):
+                message = messages[sequence_number]
                 tool_info = self._assistant_id_to_tool_info[assistant_id]
+
                 display_mode = tool_info["display_config"].mode
                 display_name = tool_info["display_name"]
                 if len(messages) > 1:
-                    display_name += f" {i + 1}"
+                    display_name += f" {sequence_number}"
 
-                loop_response.message.text = (
-                    build_sub_agent_answer_display(
+                answers.append(
+                    _build_sub_agent_answer_display(
                         display_name=display_name,
                         assistant_id=assistant_id,
                         display_mode=display_mode,
                         answer=message["text"],
                     )
-                    + loop_response.message.text
                 )
 
-                assert self._sub_agent_message is not None
-
                 loop_response.message.references.extend(
                     ContentReference(
-                        message_id=self._sub_agent_message["id"],
+                        message_id=self._main_agent_message["id"],
                         source_id=ref["sourceId"],
-                        url=ref["url"],
+                        url=ref["url"] or "",
                         source=ref["source"],
                         name=ref["name"],
                         sequence_number=ref["sequenceNumber"],
@@ -132,22 +120,44 @@ class SubAgentResponsesPostprocessor(Postprocessor):
                     for ref in message["references"]
                 )
 
-        return modified
+        loop_response.message.text = "\n".join(answers) + loop_response.message.text
+
+        return True
 
     @override
     async def remove_from_text(self, text) -> str:
         for assistant_id, tool_info in self._assistant_id_to_tool_info.items():
             display_config = tool_info["display_config"]
             if display_config.remove_from_history:
-                text = remove_sub_agent_answer_from_text(
+                text = _remove_sub_agent_answer_from_text(
                     display_mode=display_config.mode,
                     text=text,
                     assistant_id=assistant_id,
                 )
         return text
 
+    def register_sub_agent_tool(
+        self, tool: SubAgentTool, display_config: SubAgentDisplayConfig
+    ) -> None:
+        if display_config.mode == SubAgentResponseDisplayMode.HIDDEN:
+            logger.info(
+                "Sub agent tool %s has display mode `hidden`, responses will be ignored.",
+                tool.config.assistant_id,
+            )
+            return
+
+        if tool.config.assistant_id not in self._assistant_id_to_tool_info:
+            tool.subscribe(self)
+            self._assistant_id_to_tool_info[tool.config.assistant_id] = (
+                _SubAgentToolInfo(
+                    display_config=display_config,
+                    display_name=tool.display_name(),
+                    responses={},
+                )
+            )
+
     def notify_sub_agent_response(
-        self, sub_agent_assistant_id: str, response: SpaceMessage
+        self, response: SpaceMessage, sub_agent_assistant_id: str, sequence_number: int
     ) -> None:
         if sub_agent_assistant_id not in self._assistant_id_to_tool_info:
             logger.warning(
@@ -156,47 +166,57 @@ class SubAgentResponsesPostprocessor(Postprocessor):
             )
             return
 
-        self._assistant_id_to_tool_info[sub_agent_assistant_id]["responses"].append(
-            {
-                "text": response["text"],
-                "references": [
-                    {
-                        "name": ref["name"],
-                        "url": ref["url"],
-                        "sequenceNumber": ref["sequenceNumber"],
-                        "originalIndex": [],
-                        "sourceId": ref["sourceId"],
-                        "source": ref["source"],
-                    }
-                    for ref in response["references"] or []
-                ],
-            }
-        )
+        if response["text"] is None:
+            logger.warning(
+                "Sub agent response %s has no text, message will be ignored.",
+                sequence_number,
+            )
+            return
+
+        self._assistant_id_to_tool_info[sub_agent_assistant_id]["responses"][
+            sequence_number
+        ] = {
+            "text": response["text"],
+            "references": [
+                {
+                    "name": ref["name"],
+                    "url": ref["url"],
+                    "sequenceNumber": ref["sequenceNumber"],
+                    "originalIndex": [],
+                    "sourceId": ref["sourceId"],
+                    "source": ref["source"],
+                }
+                for ref in response["references"] or []
+            ],
+        }
 
 
 def _consolidate_references_in_place(
-    messages: dict[str, list[_SubAgentMessageInfo]], existing_refs: dict[str, int]
+    messages: list[_SubAgentToolInfo], existing_refs: dict[str, int]
 ) -> None:
     start_index = max(existing_refs.values(), default=0) + 1
 
-    for assistant_id, assistant_messages in messages.items():
-        for message in assistant_messages:
+    for assistant_tool_info in messages:
+        assistant_messages = assistant_tool_info["responses"]
+
+        for sequence_number in sorted(assistant_messages):
+            message = assistant_messages[sequence_number]
+
             references = message["references"]
-            if len(references) == 0 or message["text"] is None:
-                logger.info(
-                    "Message from assistant %s does not contain any references",
-                    assistant_id,
+            if len(references) == 0:
+                logger.debug(
+                    "Message from assistant %s with sequence number %s does not contain any references",
+                    assistant_tool_info["display_name"],
+                    sequence_number,
                 )
                 continue
 
             references = list(sorted(references, key=lambda ref: ref["sequenceNumber"]))
-
             ref_map = {}
-
             message_new_refs = []
+
             for reference in references:
                 source_id = reference["sourceId"]
-
                 if source_id not in existing_refs:
                     message_new_refs.append(reference)
                     existing_refs[source_id] = start_index
@@ -208,24 +228,3 @@ def _consolidate_references_in_place(
 
             message["text"] = _replace_references_in_text(message["text"], ref_map)
             message["references"] = message_new_refs
-
-
-def _replace_references_in_text_non_overlapping(
-    text: str, ref_map: dict[int, int]
-) -> str:
-    for orig, repl in ref_map.items():
-        text = re.sub(rf"<sup>{orig}</sup>", f"<sup>{repl}</sup>", text)
-    return text
-
-
-def _replace_references_in_text(text: str, ref_map: dict[int, int]) -> str:
-    # 2 phase replacement, since the map keys and values can overlap
-    max_ref = max(max(ref_map.keys(), default=0), max(ref_map.values(), default=0)) + 1
-    unique_refs = range(max_ref, max_ref + len(ref_map))
-
-    text = _replace_references_in_text_non_overlapping(
-        text, dict(zip(ref_map.keys(), unique_refs))
-    )
-    return _replace_references_in_text_non_overlapping(
-        text, dict(zip(unique_refs, ref_map.values()))
-    )