unique_toolkit 1.1.9__tar.gz → 1.2.0__tar.gz

{unique_toolkit-1.1.9 → unique_toolkit-1.2.0}/CHANGELOG.md

@@ -6,6 +6,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [1.2.0] - 2025-09-24
+- Add ability to display sub agent responses in the chat.
+
 ## [1.1.9] - 2025-09-24
 - Fix bug in `LanguageModelFunction` to extend support mistral tool calling.
 

{unique_toolkit-1.1.9 → unique_toolkit-1.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 1.1.9
+Version: 1.2.0
 Summary: 
 License: Proprietary
 Author: Cedric Klinkert
@@ -119,6 +119,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [1.2.0] - 2025-09-24
+- Add ability to display sub agent responses in the chat.
+
 ## [1.1.9] - 2025-09-24
 - Fix bug in `LanguageModelFunction` to extend support mistral tool calling.
 

{unique_toolkit-1.1.9 → unique_toolkit-1.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "unique_toolkit"
-version = "1.1.9"
+version = "1.2.0"
 description = ""
 authors = [
     "Cedric Klinkert <cedric.klinkert@unique.ch>",

{unique_toolkit-1.1.9 → unique_toolkit-1.2.0}/unique_toolkit/agentic/postprocessor/postprocessor_manager.py

@@ -16,10 +16,10 @@ class Postprocessor(ABC):
     def get_name(self) -> str:
         return self.name
 
-    async def run(self, loop_response: LanguageModelStreamResponse) -> str:
+    async def run(self, loop_response: LanguageModelStreamResponse) -> None:
         raise NotImplementedError("Subclasses must implement this method.")
 
-    async def apply_postprocessing_to_response(
+    def apply_postprocessing_to_response(
         self, loop_response: LanguageModelStreamResponse
     ) -> bool:
         raise NotImplementedError(
@@ -102,6 +102,7 @@ class PostprocessorManager:
             self._chat_service.modify_assistant_message(
                 content=loop_response.message.text,
                 message_id=loop_response.message.id,
+                references=loop_response.message.references,
             )
 
     async def execute_postprocessors(

{unique_toolkit-1.1.9 → unique_toolkit-1.2.0}/unique_toolkit/agentic/tools/a2a/config.py

@@ -1,3 +1,7 @@
+from enum import StrEnum
+
+from pydantic import BaseModel
+
 from unique_toolkit.agentic.tools.config import get_configuration_dict
 from unique_toolkit.agentic.tools.schemas import BaseToolConfig
 
@@ -6,22 +10,36 @@ This is the message that will be sent to the sub-agent.
 """.strip()
 
 
+class ResponseDisplayMode(StrEnum):
+    HIDDEN = "hidden"
+    DETAILS_OPEN = "details_open"
+    DETAILS_CLOSED = "details_closed"
+
+
+class SubAgentToolDisplayConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    mode: ResponseDisplayMode = ResponseDisplayMode.HIDDEN
+    remove_from_history: bool = True
+
+
 class SubAgentToolConfig(BaseToolConfig):
     model_config = get_configuration_dict()
 
-    name: str = "default_name"
     assistant_id: str = ""
     chat_id: str | None = None
     reuse_chat: bool = True
+
     tool_description_for_system_prompt: str = ""
     tool_description: str = ""
     param_description_sub_agent_user_message: str = (
         DEFAULT_PARAM_DESCRIPTION_SUB_AGENT_USER_MESSAGE
     )
     tool_format_information_for_system_prompt: str = ""
-
     tool_description_for_user_prompt: str = ""
     tool_format_information_for_user_prompt: str = ""
 
     poll_interval: float = 1.0
     max_wait: float = 120.0
+
+    response_display_config: SubAgentToolDisplayConfig = SubAgentToolDisplayConfig()

{unique_toolkit-1.1.9 → unique_toolkit-1.2.0}/unique_toolkit/agentic/tools/a2a/manager.py

@@ -3,8 +3,6 @@ from logging import Logger
 from unique_toolkit.agentic.tools.a2a.config import SubAgentToolConfig
 from unique_toolkit.agentic.tools.a2a.service import SubAgentTool, ToolProgressReporter
 from unique_toolkit.agentic.tools.config import ToolBuildConfig
-from unique_toolkit.agentic.tools.schemas import BaseToolConfig
-from unique_toolkit.agentic.tools.tool import Tool
 from unique_toolkit.app.schemas import ChatEvent
 
 
@@ -19,7 +17,7 @@ class A2AManager:
 
     def get_all_sub_agents(
         self, tool_configs: list[ToolBuildConfig], event: ChatEvent
-    ) -> tuple[list[ToolBuildConfig], list[Tool[BaseToolConfig]]]:
+    ) -> tuple[list[ToolBuildConfig], list[SubAgentTool]]:
         sub_agents = []
 
         for tool_config in tool_configs:
@@ -32,13 +30,15 @@ class A2AManager:
                 )
                 continue
 
-            sub_agent_tool_config: SubAgentToolConfig = tool_config.configuration
+            sub_agent_tool_config = tool_config.configuration
 
             sub_agents.append(
                 SubAgentTool(
                     configuration=sub_agent_tool_config,
                     event=event,
                     tool_progress_reporter=self._tool_progress_reporter,
+                    name=tool_config.name,
+                    display_name=tool_config.display_name,
                 )
             )
 

unique_toolkit-1.2.0/unique_toolkit/agentic/tools/a2a/postprocessing/__init__.py

@@ -0,0 +1,5 @@
+from unique_toolkit.agentic.tools.a2a.postprocessing.postprocessor import (
+    SubAgentResponsesPostprocessor,
+)
+
+__all__ = ["SubAgentResponsesPostprocessor"]

unique_toolkit-1.2.0/unique_toolkit/agentic/tools/a2a/postprocessing/display.py

@@ -0,0 +1,113 @@
+import re
+from abc import ABC, abstractmethod
+from typing import Literal, override
+
+from unique_toolkit.agentic.tools.a2a.config import ResponseDisplayMode
+
+
+class _ResponseDisplayHandler(ABC):
+    @abstractmethod
+    def build_response_display(
+        self, display_name: str, assistant_id: str, answer: str
+    ) -> str:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def remove_response_display(self, assistant_id: str, text: str) -> str:
+        raise NotImplementedError()
+
+
+class _DetailsResponseDisplayHandler(_ResponseDisplayHandler):
+    def __init__(self, mode: Literal["open", "closed"]) -> None:
+        self._mode = mode
+
+    DETAILS_CLOSED_TEMPLATE = (
+        "<details><summary>{display_name}</summary>\n"
+        "\n"
+        '<div style="display: none;">{assistant_id}</div>\n'
+        "\n"
+        "{answer}\n"
+        "</details>\n"
+        "<br>\n"
+        "\n"
+    )
+
+    DETAILS_OPEN_TEMPLATE = (
+        "<details open><summary>{display_name}</summary>\n"
+        "\n"
+        '<div style="display: none;">{assistant_id}</div>\n'
+        "\n"
+        "{answer}\n"
+        "\n"
+        "</details>\n"
+        "<br>\n"
+        "\n"
+    )
+
+    def _get_detect_re(self, assistant_id: str) -> str:
+        if self._mode == "open":
+            return (
+                r"(?s)<details open>\s*"
+                r"<summary>(.*?)</summary>\s*"
+                rf"<div style=\"display: none;\">{re.escape(assistant_id)}</div>\s*"
+                r"(.*?)\s*"
+                r"</details>\s*"
+                r"<br>\s*"
+            )
+        else:
+            return (
+                r"(?s)<details>\s*"
+                r"<summary>(.*?)</summary>\s*"
+                rf"<div style=\"display: none;\">{re.escape(assistant_id)}</div>\s*"
+                r"(.*?)\s*"
+                r"</details>\s*"
+                r"<br>\s*"
+            )
+
+    def _get_template(self) -> str:
+        if self._mode == "open":
+            return self.DETAILS_OPEN_TEMPLATE
+        else:
+            return self.DETAILS_CLOSED_TEMPLATE
+
+    @override
+    def build_response_display(
+        self, display_name: str, assistant_id: str, answer: str
+    ) -> str:
+        return self._get_template().format(
+            assistant_id=assistant_id, display_name=display_name, answer=answer
+        )
+
+    @override
+    def remove_response_display(self, assistant_id: str, text: str) -> str:
+        return re.sub(self._get_detect_re(assistant_id=assistant_id), "", text)
+
+
+_DISPLAY_HANDLERS = {
+    ResponseDisplayMode.DETAILS_OPEN: _DetailsResponseDisplayHandler(mode="open"),
+    ResponseDisplayMode.DETAILS_CLOSED: _DetailsResponseDisplayHandler(mode="closed"),
+}
+
+
+def build_sub_agent_answer_display(
+    display_name: str, display_mode: ResponseDisplayMode, answer: str, assistant_id: str
+) -> str:
+    if display_mode not in _DISPLAY_HANDLERS:
+        return ""
+
+    display_f = _DISPLAY_HANDLERS[display_mode]
+
+    return display_f.build_response_display(
+        display_name=display_name, answer=answer, assistant_id=assistant_id
+    )
+
+
+def remove_sub_agent_answer_from_text(
+    display_mode: ResponseDisplayMode, text: str, assistant_id: str
+) -> str:
+    if display_mode not in _DISPLAY_HANDLERS:
+        return text
+
+    display_f = _DISPLAY_HANDLERS[display_mode]
+
+    return display_f.remove_response_display(assistant_id=assistant_id, text=text)

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

@@ -0,0 +1,204 @@
+import logging
+import re
+from typing import NotRequired, 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.service import SubAgentTool
+from unique_toolkit.content.schemas import ContentReference
+from unique_toolkit.language_model.schemas import LanguageModelStreamResponse
+
+logger = logging.getLogger(__name__)
+
+SpaceMessage = unique_sdk.Space.Message
+
+
+class _SubAgentMessageInfo(TypedDict):
+    text: str | None
+    references: list[unique_sdk.Space.Reference]
+
+
+class _SubAgentToolInfo(TypedDict):
+    display_name: str
+    display_config: SubAgentToolDisplayConfig
+    response: NotRequired[_SubAgentMessageInfo]
+
+
+class SubAgentResponsesPostprocessor(Postprocessor):
+    def __init__(
+        self,
+        user_id: str,
+        company_id: str,
+        agent_chat_id: str,
+        sub_agent_tools: list[SubAgentTool],
+    ):
+        super().__init__(name=self.__class__.__name__)
+
+        self._user_id = user_id
+        self._company_id = company_id
+
+        self._agent_chat_id = 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(),
+                )
+            )
+
+        self._sub_agent_message = None
+
+    @override
+    async def run(self, loop_response: LanguageModelStreamResponse) -> None:
+        self._sub_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,
+        )
+
+    @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 "response" not in tool_info:
+                logger.warning(
+                    "No response from assistant %s",
+                    assistant_id,
+                )
+                continue
+
+            if display_mode != ResponseDisplayMode.HIDDEN:
+                displayed[assistant_id] = tool_info["response"]
+
+        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, message in reversed(displayed.items()):
+            tool_info = self._assistant_id_to_tool_info[assistant_id]
+            display_mode = tool_info["display_config"].mode
+            display_name = tool_info["display_name"]
+            loop_response.message.text = (
+                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"],
+                    source_id=ref["sourceId"],
+                    url=ref["url"],
+                    source=ref["source"],
+                    name=ref["name"],
+                    sequence_number=ref["sequenceNumber"],
+                )
+                for ref in message["references"]
+            )
+
+        return modified
+
+    @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(
+                    display_mode=display_config.mode,
+                    text=text,
+                    assistant_id=assistant_id,
+                )
+        return text
+
+    def notify_sub_agent_response(
+        self, sub_agent_assistant_id: str, response: SpaceMessage
+    ) -> None:
+        if sub_agent_assistant_id not in self._assistant_id_to_tool_info:
+            logger.warning(
+                "Unknown assistant id %s received, message will be ignored.",
+                sub_agent_assistant_id,
+            )
+            return
+
+        self._assistant_id_to_tool_info[sub_agent_assistant_id]["response"] = {
+            "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, _SubAgentMessageInfo], existing_refs: dict[str, int]
+) -> None:
+    start_index = max(existing_refs.values(), default=0) + 1
+
+    for assistant_id, message in messages.items():
+        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,
+            )
+            continue
+
+        references = list(sorted(references, key=lambda ref: ref["sequenceNumber"]))
+
+        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
+                start_index += 1
+
+            reference_num = existing_refs[source_id]
+
+            seq_num = reference["sequenceNumber"]
+            message["text"] = re.sub(
+                rf"<sup>{seq_num}</sup>",
+                f"<sup>{reference_num}</sup>",
+                message["text"],
+            )
+            reference["sequenceNumber"] = reference_num
+
+        message["references"] = message_new_refs