unique_toolkit 1.1.8__py3-none-any.whl → 1.2.0__py3-none-any.whl

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/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/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/agentic/tools/a2a/postprocessing/__init__.py

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

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/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

unique_toolkit/agentic/tools/a2a/postprocessing/test/test_display.py

@@ -0,0 +1,412 @@
+import re
+
+import pytest
+
+from unique_toolkit.agentic.tools.a2a.config import ResponseDisplayMode
+from unique_toolkit.agentic.tools.a2a.postprocessing.display import (
+    _DetailsResponseDisplayHandler,
+    build_sub_agent_answer_display,
+    remove_sub_agent_answer_from_text,
+)
+
+
+class TestDetailsResponseDisplayHandler:
+    """Test suite for DetailsResponseDisplayHandler class."""
+
+    @pytest.fixture
+    def open_handler(self):
+        """Create a handler with open mode."""
+        return _DetailsResponseDisplayHandler(mode="open")
+
+    @pytest.fixture
+    def closed_handler(self):
+        """Create a handler with closed mode."""
+        return _DetailsResponseDisplayHandler(mode="closed")
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for testing."""
+        return {
+            "display_name": "Test Assistant",
+            "assistant_id": "test_assistant_123",
+            "answer": "This is a test answer with multiple lines.\nSecond line here.",
+        }
+
+    def test_build_response_display_open_mode(self, open_handler, sample_data):
+        """Test building response display in open mode."""
+        result = open_handler.build_response_display(
+            display_name=sample_data["display_name"],
+            assistant_id=sample_data["assistant_id"],
+            answer=sample_data["answer"],
+        )
+
+        assert "<details open>" in result
+        assert (
+            f'<div style="display: none;">{sample_data["assistant_id"]}</div>' in result
+        )
+        assert f"<summary>{sample_data['display_name']}</summary>" in result
+        assert sample_data["answer"] in result
+        assert "</details>" in result
+
+    def test_build_response_display_closed_mode(self, closed_handler, sample_data):
+        """Test building response display in closed mode."""
+        result = closed_handler.build_response_display(
+            display_name=sample_data["display_name"],
+            assistant_id=sample_data["assistant_id"],
+            answer=sample_data["answer"],
+        )
+
+        assert "<details>" in result
+        assert "<details open>" not in result
+        assert (
+            f'<div style="display: none;">{sample_data["assistant_id"]}</div>' in result
+        )
+        assert f"<summary>{sample_data['display_name']}</summary>" in result
+        assert sample_data["answer"] in result
+        assert "</details>" in result
+
+    def test_build_response_display_with_special_characters(self, open_handler):
+        """Test building response display with special characters in content."""
+        result = open_handler.build_response_display(
+            display_name="Test & Co.",
+            assistant_id="test<>123",
+            answer="Answer with <tags> & symbols",
+        )
+
+        assert "Test & Co." in result
+        assert "test<>123" in result
+        assert "Answer with <tags> & symbols" in result
+
+    def test_remove_response_display_open_mode(self, open_handler, sample_data):
+        """Test removing response display from text in open mode."""
+        # First build the display
+        display_html = open_handler.build_response_display(
+            display_name=sample_data["display_name"],
+            assistant_id=sample_data["assistant_id"],
+            answer=sample_data["answer"],
+        )
+
+        # Create text with the display embedded
+        text_with_display = f"Some text before\n{display_html}\nSome text after"
+
+        # Remove the display
+        result = open_handler.remove_response_display(
+            assistant_id=sample_data["assistant_id"], text=text_with_display
+        )
+
+        assert "Some text before" in result
+        assert "Some text after" in result
+        assert sample_data["display_name"] not in result
+        assert sample_data["answer"] not in result
+
+    def test_remove_response_display_closed_mode(self, closed_handler, sample_data):
+        """Test removing response display from text in closed mode."""
+        # First build the display
+        display_html = closed_handler.build_response_display(
+            display_name=sample_data["display_name"],
+            assistant_id=sample_data["assistant_id"],
+            answer=sample_data["answer"],
+        )
+
+        # Create text with the display embedded
+        text_with_display = f"Some text before\n{display_html}\nSome text after"
+
+        # Remove the display
+        result = closed_handler.remove_response_display(
+            assistant_id=sample_data["assistant_id"], text=text_with_display
+        )
+
+        assert "Some text before" in result
+        assert "Some text after" in result
+        assert sample_data["display_name"] not in result
+        assert sample_data["answer"] not in result
+
+    def test_remove_response_display_multiple_instances(self, open_handler):
+        """Test removing multiple instances of response display."""
+        assistant_id = "test_123"
+
+        display1 = open_handler.build_response_display(
+            display_name="First", assistant_id=assistant_id, answer="First answer"
+        )
+
+        display2 = open_handler.build_response_display(
+            display_name="Second", assistant_id=assistant_id, answer="Second answer"
+        )
+
+        text_with_displays = f"Start\n{display1}\nMiddle\n{display2}\nEnd"
+
+        result = open_handler.remove_response_display(
+            assistant_id=assistant_id, text=text_with_displays
+        )
+
+        assert "Start" in result
+        assert "Middle" in result
+        assert "End" in result
+        assert "First answer" not in result
+        assert "Second answer" not in result
+
+    def test_remove_response_display_no_match(self, open_handler):
+        """Test removing response display when no match exists."""
+        text = "This is some text without any displays"
+        result = open_handler.remove_response_display(
+            assistant_id="nonexistent", text=text
+        )
+        assert result == text
+
+    def test_remove_response_display_with_regex_special_chars(self, open_handler):
+        """Test removing response display with regex special characters in assistant_id."""
+        assistant_id = "test.+*?[]{}()^$|"
+
+        display_html = open_handler.build_response_display(
+            display_name="Test", assistant_id=assistant_id, answer="Test answer"
+        )
+
+        text_with_display = f"Before\n{display_html}\nAfter"
+
+        result = open_handler.remove_response_display(
+            assistant_id=assistant_id, text=text_with_display
+        )
+
+        assert "Before" in result
+        assert "After" in result
+        assert "Test answer" not in result
+
+    def test_get_detect_re_pattern_validity(self, open_handler, closed_handler):
+        """Test that the regex patterns are valid and compilable."""
+        assistant_id = "test_123"
+
+        open_pattern = open_handler._get_detect_re(assistant_id)
+        closed_pattern = closed_handler._get_detect_re(assistant_id)
+
+        # Should not raise exceptions
+        re.compile(open_pattern)
+        re.compile(closed_pattern)
+
+        assert "(?s)" in open_pattern  # multiline flag
+        assert "(?s)" in closed_pattern
+        assert "details open" in open_pattern
+        assert "details>" in closed_pattern
+        assert "details open" not in closed_pattern
+
+
+class TestDisplayFunctions:
+    """Test suite for module-level display functions."""
+
+    @pytest.fixture
+    def sample_data(self):
+        """Sample data for testing."""
+        return {
+            "display_name": "Test Assistant",
+            "assistant_id": "test_assistant_123",
+            "answer": "This is a test answer.",
+        }
+
+    def test_build_sub_agent_answer_display_details_open(self, sample_data):
+        """Test building sub-agent answer display with DETAILS_OPEN mode."""
+        result = build_sub_agent_answer_display(
+            display_name=sample_data["display_name"],
+            display_mode=ResponseDisplayMode.DETAILS_OPEN,
+            answer=sample_data["answer"],
+            assistant_id=sample_data["assistant_id"],
+        )
+
+        assert "<details open>" in result
+        assert sample_data["display_name"] in result
+        assert sample_data["answer"] in result
+        assert sample_data["assistant_id"] in result
+
+    def test_build_sub_agent_answer_display_details_closed(self, sample_data):
+        """Test building sub-agent answer display with DETAILS_CLOSED mode."""
+        result = build_sub_agent_answer_display(
+            display_name=sample_data["display_name"],
+            display_mode=ResponseDisplayMode.DETAILS_CLOSED,
+            answer=sample_data["answer"],
+            assistant_id=sample_data["assistant_id"],
+        )
+
+        assert "<details>" in result
+        assert "<details open>" not in result
+        assert sample_data["display_name"] in result
+        assert sample_data["answer"] in result
+        assert sample_data["assistant_id"] in result
+
+    def test_build_sub_agent_answer_display_hidden_mode(self, sample_data):
+        """Test building sub-agent answer display with HIDDEN mode."""
+        result = build_sub_agent_answer_display(
+            display_name=sample_data["display_name"],
+            display_mode=ResponseDisplayMode.HIDDEN,
+            answer=sample_data["answer"],
+            assistant_id=sample_data["assistant_id"],
+        )
+
+        assert result == ""
+
+    def test_remove_sub_agent_answer_from_text_details_open(self, sample_data):
+        """Test removing sub-agent answer from text with DETAILS_OPEN mode."""
+        # First build the display
+        display_html = build_sub_agent_answer_display(
+            display_name=sample_data["display_name"],
+            display_mode=ResponseDisplayMode.DETAILS_OPEN,
+            answer=sample_data["answer"],
+            assistant_id=sample_data["assistant_id"],
+        )
+
+        text_with_display = f"Before\n{display_html}\nAfter"
+
+        result = remove_sub_agent_answer_from_text(
+            display_mode=ResponseDisplayMode.DETAILS_OPEN,
+            text=text_with_display,
+            assistant_id=sample_data["assistant_id"],
+        )
+
+        assert "Before" in result
+        assert "After" in result
+        assert sample_data["answer"] not in result
+
+    def test_remove_sub_agent_answer_from_text_details_closed(self, sample_data):
+        """Test removing sub-agent answer from text with DETAILS_CLOSED mode."""
+        # First build the display
+        display_html = build_sub_agent_answer_display(
+            display_name=sample_data["display_name"],
+            display_mode=ResponseDisplayMode.DETAILS_CLOSED,
+            answer=sample_data["answer"],
+            assistant_id=sample_data["assistant_id"],
+        )
+
+        text_with_display = f"Before\n{display_html}\nAfter"
+
+        result = remove_sub_agent_answer_from_text(
+            display_mode=ResponseDisplayMode.DETAILS_CLOSED,
+            text=text_with_display,
+            assistant_id=sample_data["assistant_id"],
+        )
+
+        assert "Before" in result
+        assert "After" in result
+        assert sample_data["answer"] not in result
+
+    def test_remove_sub_agent_answer_from_text_hidden_mode(self, sample_data):
+        """Test removing sub-agent answer from text with HIDDEN mode."""
+        text = "Some text here"
+        result = remove_sub_agent_answer_from_text(
+            display_mode=ResponseDisplayMode.HIDDEN,
+            text=text,
+            assistant_id=sample_data["assistant_id"],
+        )
+
+        assert result == text
+
+    def test_roundtrip_build_and_remove(self, sample_data):
+        """Test that building and then removing display results in clean text."""
+        original_text = "This is the original text."
+
+        # Build display
+        display_html = build_sub_agent_answer_display(
+            display_name=sample_data["display_name"],
+            display_mode=ResponseDisplayMode.DETAILS_OPEN,
+            answer=sample_data["answer"],
+            assistant_id=sample_data["assistant_id"],
+        )
+
+        # Insert into text
+        text_with_display = f"{original_text}\n{display_html}"
+
+        # Remove display
+        result = remove_sub_agent_answer_from_text(
+            display_mode=ResponseDisplayMode.DETAILS_OPEN,
+            text=text_with_display,
+            assistant_id=sample_data["assistant_id"],
+        )
+
+        # Should be back to original (with some whitespace differences)
+        assert original_text in result.strip()
+        assert sample_data["answer"] not in result
+
+
+class TestEdgeCases:
+    """Test suite for edge cases and error conditions."""
+
+    def test_empty_strings(self):
+        """Test handling of empty strings."""
+        handler = _DetailsResponseDisplayHandler(mode="open")
+
+        result = handler.build_response_display(
+            display_name="", assistant_id="test", answer=""
+        )
+
+        assert "<details open>" in result
+        assert "<summary></summary>" in result
+
+    def test_multiline_content(self):
+        """Test handling of multiline content."""
+        handler = _DetailsResponseDisplayHandler(mode="open")
+
+        multiline_answer = """Line 1
+        Line 2
+        Line 3 with    spaces
+        
+        Line 5 after blank line"""
+
+        result = handler.build_response_display(
+            display_name="Multi-line Test",
+            assistant_id="test_ml",
+            answer=multiline_answer,
+        )
+
+        assert multiline_answer in result
+
+        # Test removal
+        text_with_display = f"Before\n{result}\nAfter"
+        clean_result = handler.remove_response_display(
+            assistant_id="test_ml", text=text_with_display
+        )
+
+        assert "Before" in clean_result
+        assert "After" in clean_result
+        assert multiline_answer not in clean_result
+
+    def test_html_content_in_answer(self):
+        """Test handling of HTML content within the answer."""
+        handler = _DetailsResponseDisplayHandler(mode="open")
+
+        html_answer = "<p>This is <strong>bold</strong> text with <em>emphasis</em></p>"
+
+        result = handler.build_response_display(
+            display_name="HTML Test", assistant_id="test_html", answer=html_answer
+        )
+
+        assert html_answer in result
+
+        # Test removal
+        text_with_display = f"Before\n{result}\nAfter"
+        clean_result = handler.remove_response_display(
+            assistant_id="test_html", text=text_with_display
+        )
+
+        assert "Before" in clean_result
+        assert "After" in clean_result
+        assert html_answer not in clean_result
+
+    def test_unicode_content(self):
+        """Test handling of Unicode content."""
+        handler = _DetailsResponseDisplayHandler(mode="open")
+
+        unicode_content = "Testing Unicode: 你好 🌟 café naïve résumé"
+
+        result = handler.build_response_display(
+            display_name="Unicode Test",
+            assistant_id="test_unicode",
+            answer=unicode_content,
+        )
+
+        assert unicode_content in result
+
+        # Test removal
+        text_with_display = f"Before\n{result}\nAfter"
+        clean_result = handler.remove_response_display(
+            assistant_id="test_unicode", text=text_with_display
+        )
+
+        assert "Before" in clean_result
+        assert "After" in clean_result
+        assert unicode_content not in clean_result

unique_toolkit/agentic/tools/a2a/service.py

@@ -1,8 +1,13 @@
+from typing import Protocol, override
+
+import unique_sdk
 from pydantic import Field, create_model
 from unique_sdk.utils.chat_in_space import send_message_and_wait_for_completion
 
 from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
-from unique_toolkit.agentic.tools.a2a.config import SubAgentToolConfig
+from unique_toolkit.agentic.tools.a2a.config import (
+    SubAgentToolConfig,
+)
 from unique_toolkit.agentic.tools.a2a.memory import (
     get_sub_agent_short_term_memory_manager,
 )
@@ -10,6 +15,7 @@ from unique_toolkit.agentic.tools.a2a.schema import (
     SubAgentShortTermMemorySchema,
     SubAgentToolInput,
 )
+from unique_toolkit.agentic.tools.agent_chunks_hanlder import AgentChunksHandler
 from unique_toolkit.agentic.tools.factory import ToolFactory
 from unique_toolkit.agentic.tools.schemas import ToolCallResponse
 from unique_toolkit.agentic.tools.tool import Tool
@@ -20,9 +26,17 @@ from unique_toolkit.agentic.tools.tool_progress_reporter import (
 from unique_toolkit.app import ChatEvent
 from unique_toolkit.language_model import (
     LanguageModelFunction,
-    LanguageModelMessage,
     LanguageModelToolDescription,
 )
+from unique_toolkit.language_model.schemas import LanguageModelMessage
+
+
+class SubAgentResponseSubscriber(Protocol):
+    def notify_sub_agent_response(
+        self,
+        sub_agent_assistant_id: str,
+        response: unique_sdk.Space.Message,
+    ) -> None: ...
 
 
 class SubAgentTool(Tool[SubAgentToolConfig]):
@@ -33,11 +47,15 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
         configuration: SubAgentToolConfig,
         event: ChatEvent,
         tool_progress_reporter: ToolProgressReporter | None = None,
+        name: str = "SubAgentTool",
+        display_name: str = "SubAgentTool",
     ):
         super().__init__(configuration, event, tool_progress_reporter)
         self._user_id = event.user_id
         self._company_id = event.company_id
-        self.name = configuration.name
+
+        self.name = name
+        self._display_name = display_name
 
         self._short_term_memory_manager = get_sub_agent_short_term_memory_manager(
             company_id=self._company_id,
@@ -45,6 +63,13 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             chat_id=event.payload.chat_id,
             assistant_id=self.config.assistant_id,
         )
+        self._subscribers: list[SubAgentResponseSubscriber] = []
+
+    def display_name(self) -> str:
+        return self._display_name
+
+    def subscribe(self, subscriber: SubAgentResponseSubscriber) -> None:
+        self._subscribers.append(subscriber)
 
     def tool_description(self) -> LanguageModelToolDescription:
         tool_input_model_with_description = create_model(
@@ -105,52 +130,96 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             SubAgentShortTermMemorySchema(chat_id=chat_id)
         )
 
-    async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
-        tool_input = SubAgentToolInput.model_validate(tool_call.arguments)
-
-        if self.tool_progress_reporter:
+    async def _notify_progress(
+        self,
+        tool_call: LanguageModelFunction,
+        message: str,
+        state: ProgressState,
+    ) -> None:
+        if self.tool_progress_reporter is not None:
             await self.tool_progress_reporter.notify_from_tool_call(
                 tool_call=tool_call,
-                name=f"{self.name}",
-                message=f"{tool_input.user_message}",
-                state=ProgressState.RUNNING,
+                name=self._display_name,
+                message=message,
+                state=state,
+            )
+
+    async def _execute_and_handle_timeout(
+        self, tool_user_message: str, chat_id: str, tool_call: LanguageModelFunction
+    ) -> unique_sdk.Space.Message:
+        try:
+            return await send_message_and_wait_for_completion(
+                user_id=self._user_id,
+                assistant_id=self.config.assistant_id,
+                company_id=self._company_id,
+                text=tool_user_message,
+                chat_id=chat_id,
+                poll_interval=self.config.poll_interval,
+                max_wait=self.config.max_wait,
+                stop_condition="completedAt",
+            )
+        except TimeoutError as e:
+            await self._notify_progress(
+                tool_call=tool_call,
+                message="Timeout while waiting for response from sub agent.",
+                state=ProgressState.FAILED,
             )
 
+            raise TimeoutError(
+                "Timeout while waiting for response from sub agent. The user should consider increasing the max wait time.",
+            ) from e
+
+    def _notify_subscribers(self, response: unique_sdk.Space.Message) -> None:
+        for subsciber in self._subscribers:
+            subsciber.notify_sub_agent_response(
+                sub_agent_assistant_id=self.config.assistant_id,
+                response=response,
+            )
+
+    async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
+        tool_input = SubAgentToolInput.model_validate(tool_call.arguments)
+
+        await self._notify_progress(
+            tool_call=tool_call,
+            message=tool_input.user_message,
+            state=ProgressState.RUNNING,
+        )
+
         # Check if there is a saved chat id in short term memory
         chat_id = await self._get_chat_id()
 
-        response = await send_message_and_wait_for_completion(
-            user_id=self._user_id,
-            assistant_id=self.config.assistant_id,
-            company_id=self._company_id,
-            text=tool_input.user_message,  # type: ignore
+        response = await self._execute_and_handle_timeout(
+            tool_user_message=tool_input.user_message,  # type: ignore
             chat_id=chat_id,  # type: ignore
-            poll_interval=self.config.poll_interval,
-            max_wait=self.config.max_wait,
+            tool_call=tool_call,
         )
 
-        if chat_id is None:
+        self._notify_subscribers(response)
+
+        if chat_id is None and self.config.reuse_chat:
             await self._save_chat_id(response["chatId"])
 
         if response["text"] is None:
             raise ValueError("No response returned from sub agent")
 
-        self._text = response["text"]
+        await self._notify_progress(
+            tool_call=tool_call,
+            message=tool_input.user_message,
+            state=ProgressState.FINISHED,
+        )
+
         return ToolCallResponse(
             id=tool_call.id,  # type: ignore
             name=tool_call.name,
             content=response["text"],
         )
 
+    @override
     def get_tool_call_result_for_loop_history(
         self,
         tool_response: ToolCallResponse,
-    ) -> LanguageModelMessage:
-        return ToolCallResponse(
-            id=tool_response.id,  # type: ignore
-            name=tool_response.name,
-            content=tool_response.content,
-        )
+        agent_chunks_handler: AgentChunksHandler,
+    ) -> LanguageModelMessage: ...  # Empty as method is deprecated
 
 
 ToolFactory.register_tool(SubAgentTool, SubAgentToolConfig)

unique_toolkit/agentic/tools/config.py

@@ -78,8 +78,17 @@ class ToolBuildConfig(BaseModel):
             # Configuration can remain as a dict
             return value
 
+        is_sub_agent_tool = (
+            value.get("is_sub_agent") or value.get("isSubAgent") or False
+        )
+
         configuration = value.get("configuration", {})
-        if isinstance(configuration, dict):
+
+        if is_sub_agent_tool:
+            from unique_toolkit.agentic.tools.a2a.config import SubAgentToolConfig
+
+            config = SubAgentToolConfig.model_validate(configuration)
+        elif isinstance(configuration, dict):
             # Local import to avoid circular import at module import time
             from unique_toolkit.agentic.tools.factory import ToolFactory
 

unique_toolkit/agentic/tools/tool_manager.py

@@ -6,6 +6,7 @@ from pydantic import BaseModel, Field
 
 from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
 from unique_toolkit.agentic.tools.a2a.manager import A2AManager
+from unique_toolkit.agentic.tools.a2a.service import SubAgentTool
 from unique_toolkit.agentic.tools.config import ToolBuildConfig
 from unique_toolkit.agentic.tools.factory import ToolFactory
 from unique_toolkit.agentic.tools.mcp.manager import MCPManager
@@ -113,6 +114,7 @@ class ToolManager:
         mcp_tools = self._mcp_manager.get_all_mcp_tools()
         # Combine both types of tools
         self.available_tools = internal_tools + mcp_tools + sub_agents
+        self._sub_agents = sub_agents
 
         for t in self.available_tools:
             if not t.is_enabled():
@@ -136,6 +138,10 @@ class ToolManager:
 
             self._tools.append(t)
 
+    @property
+    def sub_agents(self) -> list[SubAgentTool]:
+        return self._sub_agents
+
     def get_evaluation_check_list(self) -> list[EvaluationMetricName]:
         return list(self._tool_evaluation_check_list)
 
@@ -143,7 +149,7 @@ class ToolManager:
         self._logger.info(f"Loaded tools: {[tool.name for tool in self._tools]}")
 
     def get_tools(self) -> list[Tool]:
-        return self._tools
+        return self._tools  # type: ignore
 
     def get_tool_by_name(self, name: str) -> Tool | None:
         for tool in self._tools:

unique_toolkit/language_model/schemas.py

@@ -80,13 +80,14 @@ class LanguageModelFunction(BaseModel):
 
     @field_validator("id", mode="before")
     def randomize_id(cls, value):
-        return uuid4().hex
+        if not value:
+            return uuid4().hex
+        else:
+            return value
 
     @model_serializer()
     def serialize_model(self):
         seralization = {}
-        if self.id:
-            seralization["id"] = self.id
         seralization["name"] = self.name
         if self.arguments:
             seralization["arguments"] = json.dumps(self.arguments)

{unique_toolkit-1.1.8.dist-info → unique_toolkit-1.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 1.1.8
+Version: 1.2.0
 Summary: 
 License: Proprietary
 Author: Cedric Klinkert
@@ -118,6 +118,13 @@ All notable changes to this project will be documented in this file.
 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.
+
 ## [1.1.8] - 2025-09-23
 - Revert last to version 1.1.6
 

{unique_toolkit-1.1.8.dist-info → unique_toolkit-1.2.0.dist-info}/RECORD

@@ -44,19 +44,23 @@ unique_toolkit/agentic/history_manager/history_construction_with_contents.py,sha
 unique_toolkit/agentic/history_manager/history_manager.py,sha256=qo6vjEXueCXUxHSrawYeMmFn7tuTEXGTWnBVb0H8bDY,8423
 unique_toolkit/agentic/history_manager/loop_token_reducer.py,sha256=9ZWh3dfs7SGsU33bxzKo9x94haEhc7Rerk9PtfFzcxg,18474
 unique_toolkit/agentic/history_manager/utils.py,sha256=NDSSz0Jp3oVJU3iKlVScmM1AOe-6hTiVjLr16DUPsV0,5656
-unique_toolkit/agentic/postprocessor/postprocessor_manager.py,sha256=lUCPQsYIGjpJdw7ZqEG3JcGc3B2bFBFqH9Rs1KxWzaM,4221
+unique_toolkit/agentic/postprocessor/postprocessor_manager.py,sha256=GDzJhaoOUwxZ37IINkQ7au4CHmAOFS5miP2lqv8ZwZA,4277
 unique_toolkit/agentic/reference_manager/reference_manager.py,sha256=1GeoFX1-RLdTcns1358GJADDSAcTAM2J0jJJpln08qo,4005
 unique_toolkit/agentic/short_term_memory_manager/persistent_short_term_memory_manager.py,sha256=uF3HSoZF0hBfuNhIE9N8KRtuwDfpoeXUFVrv_cyZ3Sw,5839
 unique_toolkit/agentic/thinking_manager/thinking_manager.py,sha256=41QWFsdRrbWlQHBfYCFv726UDom4WbcvaRfjCmoUOQI,4183
 unique_toolkit/agentic/tools/__init__.py,sha256=-ToY9-Xiz0K7qCUydH1h1yG6n4h1hQS8sBuSVPNEq2Y,43
 unique_toolkit/agentic/tools/a2a/__init__.py,sha256=NdY0J33b1G4sbx6UWwNS74JVSAeEtu8u_iEXOT64Uq0,187
-unique_toolkit/agentic/tools/a2a/config.py,sha256=6A75FdiCOh35D4lJbmcUZNiYQ6v32T-kWYMA4v3NTl8,880
-unique_toolkit/agentic/tools/a2a/manager.py,sha256=JYX8Sh4zWdMG9altCkAIwUxwqC0kz_W_F3pSI89qKRw,1706
+unique_toolkit/agentic/tools/a2a/config.py,sha256=m6INkff6jHEMB_DPm8MmtXfrGaMy1TIwXd1vFwHBFAU,1312
+unique_toolkit/agentic/tools/a2a/manager.py,sha256=yuuQuBrAcsT3gAWEdxf6EvRnL_iWtvaK14lRs21w5PA,1665
 unique_toolkit/agentic/tools/a2a/memory.py,sha256=4VFBzITCv5E_8YCc4iF4Y6FhzplS2C-FZaZHdeC7DyA,1028
+unique_toolkit/agentic/tools/a2a/postprocessing/__init__.py,sha256=9M5RRSO4gmQVI1gfve6MGEfVktT9WcELQFZ8Sv6xu4Y,160
+unique_toolkit/agentic/tools/a2a/postprocessing/display.py,sha256=rlBWO2M8Lr5Kx-vmvwlV-vEu33BZE4votP-TMr3-3Dw,3366
+unique_toolkit/agentic/tools/a2a/postprocessing/postprocessor.py,sha256=BWXFsv5akPhkNxRQ3JLktrBkoT1_66joTel_uM4BgRY,6963
+unique_toolkit/agentic/tools/a2a/postprocessing/test/test_display.py,sha256=UVOd5tVTltngVfsdOn6cuyvG7QmRBkfgUj_SruTgmHI,15279
 unique_toolkit/agentic/tools/a2a/schema.py,sha256=T1l5z6trtPE5nhqPzt5tvfRNDhqL_ST1Wj7_lBWJ58g,304
-unique_toolkit/agentic/tools/a2a/service.py,sha256=7n7BkJxTR-hpoQczfHLK8eS0rtATL7-YARH-4netbbI,5339
+unique_toolkit/agentic/tools/a2a/service.py,sha256=cACUPwWXzi_XfPlFqK0x6nawugx17Vq18-XRUbhgL4k,7669
 unique_toolkit/agentic/tools/agent_chunks_hanlder.py,sha256=x32Dp1Z8cVW5i-XzXbaMwX2KHPcNGmqEU-FB4AV9ZGo,1909
-unique_toolkit/agentic/tools/config.py,sha256=B9Tl0uJ81XB4SjSke4kSydame4glGT2s69ZXfr0ViZU,4634
+unique_toolkit/agentic/tools/config.py,sha256=91Gw92YoTC6WeWa4lfOpXSvIYekCwELcVNkHZZEkW2o,4936
 unique_toolkit/agentic/tools/factory.py,sha256=Wt0IGSbLg8ZTq5PU9p_JTW0LtNATWLpc3336irJKXlM,1277
 unique_toolkit/agentic/tools/mcp/__init__.py,sha256=RLF_p-LDRC7GhiB3fdCi4u3bh6V9PY_w26fg61BLyco,122
 unique_toolkit/agentic/tools/mcp/manager.py,sha256=DPYwwDe6RSZyuPaxn-je49fP_qOOs0ZV46EM6GZcV4c,2748
@@ -66,7 +70,7 @@ unique_toolkit/agentic/tools/schemas.py,sha256=0ZR8xCdGj1sEdPE0lfTIG2uSR5zqWoprU
 unique_toolkit/agentic/tools/test/test_mcp_manager.py,sha256=9F7FjpYKeOkg2Z4bt2H1WGaxu9fzB-1iiE-b7g3KzQk,15724
 unique_toolkit/agentic/tools/test/test_tool_progress_reporter.py,sha256=dod5QPqgGUInVAGXAbsAKNTEypIi6pUEWhDbJr9YfUU,6307
 unique_toolkit/agentic/tools/tool.py,sha256=m56VLxiHuKU2_J5foZp00xhm5lTxWEW7zRLGbIE9ssU,6744
-unique_toolkit/agentic/tools/tool_manager.py,sha256=E1eKZuc7eHPeGFqrGfdMGPNLqu3NYhrfKDjSJxMSuvU,10925
+unique_toolkit/agentic/tools/tool_manager.py,sha256=l8OGQiSeMWqesnFQ4vmgnrSU4e6ipAbB2pq0HmQM4AE,11140
 unique_toolkit/agentic/tools/tool_progress_reporter.py,sha256=ixud9VoHey1vlU1t86cW0-WTvyTwMxNSWBon8I11SUk,7955
 unique_toolkit/agentic/tools/utils/__init__.py,sha256=iD1YYzf9LcJFv95Z8BqCAFSewNBabybZRZyvPKGfvro,27
 unique_toolkit/agentic/tools/utils/execution/__init__.py,sha256=OHiKpqBnfhBiEQagKVWJsZlHv8smPp5OI4dFIexzibw,37
@@ -117,7 +121,7 @@ unique_toolkit/language_model/functions.py,sha256=PNCmbYovhgMSkY89p7-3DunG6jIeka
 unique_toolkit/language_model/infos.py,sha256=eHln--Y5f6znFxknV6A8m-fRaEpH5-kmRh9m-ZWqco4,57188
 unique_toolkit/language_model/prompt.py,sha256=JSawaLjQg3VR-E2fK8engFyJnNdk21zaO8pPIodzN4Q,3991
 unique_toolkit/language_model/reference.py,sha256=nkX2VFz-IrUz8yqyc3G5jUMNwrNpxITBrMEKkbqqYoI,8583
-unique_toolkit/language_model/schemas.py,sha256=EOgy-p1GRcS46Sq0qEsN8MfOMl-KCcvEd9aCmqm9d08,16497
+unique_toolkit/language_model/schemas.py,sha256=w23zH2OAYkTsS-wAqelUdhO9TCgis0TbFa8PszmhZYY,16501
 unique_toolkit/language_model/service.py,sha256=JkYGtCug8POQskTv_aoYkzTMOaPCWRM94y73o3bUttQ,12019
 unique_toolkit/language_model/utils.py,sha256=bPQ4l6_YO71w-zaIPanUUmtbXC1_hCvLK0tAFc3VCRc,1902
 unique_toolkit/protocols/support.py,sha256=V15WEIFKVMyF1QCnR8vIi4GrJy4dfTCB6d6JlqPZ58o,2341
@@ -128,7 +132,7 @@ unique_toolkit/short_term_memory/schemas.py,sha256=OhfcXyF6ACdwIXW45sKzjtZX_gkcJ
 unique_toolkit/short_term_memory/service.py,sha256=5PeVBu1ZCAfyDb2HLVvlmqSbyzBBuE9sI2o9Aajqjxg,8884
 unique_toolkit/smart_rules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unique_toolkit/smart_rules/compile.py,sha256=cxWjb2dxEI2HGsakKdVCkSNi7VK9mr08w5sDcFCQyWI,9553
-unique_toolkit-1.1.8.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
-unique_toolkit-1.1.8.dist-info/METADATA,sha256=PZk9DqoRKtH35hnqcqRyNkcy20mLwytAiuqB6EAJp28,33104
-unique_toolkit-1.1.8.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-unique_toolkit-1.1.8.dist-info/RECORD,,
+unique_toolkit-1.2.0.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
+unique_toolkit-1.2.0.dist-info/METADATA,sha256=Fp7cz1tbhox_Vcohuj8hL8ahu14rw0GJ6gUbL0t-T1c,33290
+unique_toolkit-1.2.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+unique_toolkit-1.2.0.dist-info/RECORD,,