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

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

@@ -2,11 +2,13 @@ import re
 
 import pytest
 
-from unique_toolkit.agentic.tools.a2a.config import ResponseDisplayMode
-from unique_toolkit.agentic.tools.a2a.postprocessing.display import (
+from unique_toolkit.agentic.tools.a2a.postprocessing._display import (
+    _build_sub_agent_answer_display,
     _DetailsResponseDisplayHandler,
-    build_sub_agent_answer_display,
-    remove_sub_agent_answer_from_text,
+    _remove_sub_agent_answer_from_text,
+)
+from unique_toolkit.agentic.tools.a2a.postprocessing.config import (
+    SubAgentResponseDisplayMode,
 )
 
 
@@ -14,12 +16,12 @@ class TestDetailsResponseDisplayHandler:
     """Test suite for DetailsResponseDisplayHandler class."""
 
     @pytest.fixture
-    def open_handler(self):
+    def open_handler(self) -> _DetailsResponseDisplayHandler:
         """Create a handler with open mode."""
         return _DetailsResponseDisplayHandler(mode="open")
 
     @pytest.fixture
-    def closed_handler(self):
+    def closed_handler(self) -> _DetailsResponseDisplayHandler:
         """Create a handler with closed mode."""
         return _DetailsResponseDisplayHandler(mode="closed")
 
@@ -201,82 +203,59 @@ class TestDisplayFunctions:
             "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(
+    @pytest.mark.parametrize(
+        "display_mode,expected_content,not_expected_content",
+        [
+            (SubAgentResponseDisplayMode.DETAILS_OPEN, "<details open>", None),
+            (SubAgentResponseDisplayMode.DETAILS_CLOSED, "<details>", "<details open>"),
+            (SubAgentResponseDisplayMode.HIDDEN, "", None),
+        ],
+    )
+    def test_build_sub_agent_answer_display(
+        self, sample_data, display_mode, expected_content, not_expected_content
+    ):
+        """Test building sub-agent answer display with different modes."""
+        result = _build_sub_agent_answer_display(
             display_name=sample_data["display_name"],
-            display_mode=ResponseDisplayMode.HIDDEN,
+            display_mode=display_mode,
             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."""
+        if display_mode == SubAgentResponseDisplayMode.HIDDEN:
+            assert result == ""
+        else:
+            assert expected_content in result
+            assert sample_data["display_name"] in result
+            assert sample_data["answer"] in result
+            assert sample_data["assistant_id"] in result
+
+            if not_expected_content:
+                assert not_expected_content not in result
+
+    @pytest.mark.parametrize(
+        "display_mode",
+        [
+            SubAgentResponseDisplayMode.DETAILS_OPEN,
+            SubAgentResponseDisplayMode.DETAILS_CLOSED,
+        ],
+    )
+    def test_remove_sub_agent_answer_from_text_details_modes(
+        self, sample_data, display_mode
+    ):
+        """Test removing sub-agent answer from text with DETAILS_OPEN and DETAILS_CLOSED modes."""
         # First build the display
-        display_html = build_sub_agent_answer_display(
+        display_html = _build_sub_agent_answer_display(
             display_name=sample_data["display_name"],
-            display_mode=ResponseDisplayMode.DETAILS_CLOSED,
+            display_mode=display_mode,
             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,
+        result = _remove_sub_agent_answer_from_text(
+            display_mode=display_mode,
             text=text_with_display,
             assistant_id=sample_data["assistant_id"],
         )
@@ -288,8 +267,8 @@ class TestDisplayFunctions:
     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,
+        result = _remove_sub_agent_answer_from_text(
+            display_mode=SubAgentResponseDisplayMode.HIDDEN,
             text=text,
             assistant_id=sample_data["assistant_id"],
         )
@@ -301,9 +280,9 @@ class TestDisplayFunctions:
         original_text = "This is the original text."
 
         # Build display
-        display_html = build_sub_agent_answer_display(
+        display_html = _build_sub_agent_answer_display(
             display_name=sample_data["display_name"],
-            display_mode=ResponseDisplayMode.DETAILS_OPEN,
+            display_mode=SubAgentResponseDisplayMode.DETAILS_OPEN,
             answer=sample_data["answer"],
             assistant_id=sample_data["assistant_id"],
         )
@@ -312,8 +291,8 @@ class TestDisplayFunctions:
         text_with_display = f"{original_text}\n{display_html}"
 
         # Remove display
-        result = remove_sub_agent_answer_from_text(
-            display_mode=ResponseDisplayMode.DETAILS_OPEN,
+        result = _remove_sub_agent_answer_from_text(
+            display_mode=SubAgentResponseDisplayMode.DETAILS_OPEN,
             text=text_with_display,
             assistant_id=sample_data["assistant_id"],
         )

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

@@ -1,4 +1,6 @@
-from unique_toolkit.agentic.tools.a2a.postprocessing.postprocessor import (
+import pytest
+
+from unique_toolkit.agentic.tools.a2a.postprocessing._utils import (
     _replace_references_in_text,
     _replace_references_in_text_non_overlapping,
 )
@@ -30,19 +32,21 @@ class TestReplaceReferencesInTextNonOverlapping:
         result = _replace_references_in_text_non_overlapping(text, ref_map)
         assert result == text
 
-    def test_empty_ref_map(self):
-        """Test with empty reference map."""
-        text = "This text has<sup>1</sup> references but empty map."
-        ref_map = {}
-        result = _replace_references_in_text_non_overlapping(text, ref_map)
-        assert result == text
-
-    def test_empty_text(self):
-        """Test with empty text."""
-        text = ""
-        ref_map = {1: 5}
+    @pytest.mark.parametrize(
+        "text,ref_map,expected",
+        [
+            (
+                "This text has<sup>1</sup> references but empty map.",
+                {},
+                "This text has<sup>1</sup> references but empty map.",
+            ),
+            ("", {1: 5}, ""),
+        ],
+    )
+    def test_empty_inputs(self, text, ref_map, expected):
+        """Test with empty reference map or empty text."""
         result = _replace_references_in_text_non_overlapping(text, ref_map)
-        assert result == ""
+        assert result == expected
 
     def test_reference_not_in_map(self):
         """Test with references in text that are not in the map."""
@@ -92,20 +96,24 @@ class TestReplaceReferencesInTextNonOverlapping:
         expected = "Good<sup>10</sup> and bad<sup>abc</sup> and<sup></sup>."
         assert result == expected
 
-    def test_zero_reference_number(self):
-        """Test with zero as reference number."""
-        text = "Zero reference<sup>0</sup> here."
-        ref_map = {0: 100}
-        result = _replace_references_in_text_non_overlapping(text, ref_map)
-        expected = "Zero reference<sup>100</sup> here."
-        assert result == expected
-
-    def test_negative_reference_numbers(self):
-        """Test with negative reference numbers (edge case)."""
-        text = "Negative<sup>-1</sup> reference."
-        ref_map = {-1: 5}
+    @pytest.mark.parametrize(
+        "text,ref_map,expected",
+        [
+            (
+                "Zero reference<sup>0</sup> here.",
+                {0: 100},
+                "Zero reference<sup>100</sup> here.",
+            ),
+            (
+                "Negative<sup>-1</sup> reference.",
+                {-1: 5},
+                "Negative<sup>5</sup> reference.",
+            ),
+        ],
+    )
+    def test_special_reference_numbers(self, text, ref_map, expected):
+        """Test with zero and negative reference numbers."""
         result = _replace_references_in_text_non_overlapping(text, ref_map)
-        expected = "Negative<sup>5</sup> reference."
         assert result == expected
 
 
@@ -152,26 +160,26 @@ class TestReplaceReferencesInText:
         expected = "A<sup>4</sup>B<sup>1</sup>C<sup>2</sup>D<sup>3</sup>."
         assert result == expected
 
-    def test_empty_ref_map(self):
-        """Test with empty reference map."""
-        text = "Text with<sup>1</sup> references."
-        ref_map = {}
-        result = _replace_references_in_text(text, ref_map)
-        assert result == text
-
-    def test_empty_text(self):
-        """Test with empty text."""
-        text = ""
-        ref_map = {1: 2}
-        result = _replace_references_in_text(text, ref_map)
-        assert result == ""
-
-    def test_no_references_in_text(self):
-        """Test with text containing no references."""
-        text = "This text has no references."
-        ref_map = {1: 10, 2: 20}
+    @pytest.mark.parametrize(
+        "text,ref_map,expected",
+        [
+            (
+                "Text with<sup>1</sup> references.",
+                {},
+                "Text with<sup>1</sup> references.",
+            ),
+            ("", {1: 2}, ""),
+            (
+                "This text has no references.",
+                {1: 10, 2: 20},
+                "This text has no references.",
+            ),
+        ],
+    )
+    def test_edge_cases(self, text, ref_map, expected):
+        """Test edge cases: empty reference map, empty text, and text with no references."""
         result = _replace_references_in_text(text, ref_map)
-        assert result == text
+        assert result == expected
 
     def test_single_reference_no_overlap(self):
         """Test single reference with no overlap potential."""

unique_toolkit/agentic/tools/a2a/tool/__init__.py

@@ -0,0 +1,4 @@
+from unique_toolkit.agentic.tools.a2a.tool.config import SubAgentToolConfig
+from unique_toolkit.agentic.tools.a2a.tool.service import SubAgentTool
+
+__all__ = ["SubAgentTool", "SubAgentToolConfig"]

unique_toolkit/agentic/tools/a2a/{memory.py → tool/_memory.py}

@@ -2,7 +2,7 @@ from unique_toolkit import ShortTermMemoryService
 from unique_toolkit.agentic.short_term_memory_manager.persistent_short_term_memory_manager import (
     PersistentShortMemoryManager,
 )
-from unique_toolkit.agentic.tools.a2a.schema import SubAgentShortTermMemorySchema
+from unique_toolkit.agentic.tools.a2a.tool._schema import SubAgentShortTermMemorySchema
 
 
 def _get_short_term_memory_name(assistant_id: str) -> str:

unique_toolkit/agentic/tools/a2a/{schema.py → tool/_schema.py}

@@ -1,15 +1,9 @@
 from pydantic import BaseModel
 
-from unique_toolkit.agentic.tools.schemas import ToolCallResponse
-
 
 class SubAgentToolInput(BaseModel):
     user_message: str
 
 
-class SubAgentToolCallResponse(ToolCallResponse):
-    assistant_message: str
-
-
 class SubAgentShortTermMemorySchema(BaseModel):
     chat_id: str

unique_toolkit/agentic/tools/a2a/tool/config.py

@@ -0,0 +1,63 @@
+from pydantic import Field
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit.agentic.tools.schemas import BaseToolConfig
+
+DEFAULT_PARAM_DESCRIPTION_SUB_AGENT_USER_MESSAGE = """
+This is the message that will be sent to the sub-agent.
+""".strip()
+
+DEFAULT_FORMAT_INFORMATION_SUB_AGENT_SYSTEM_MESSAGE = """
+NEVER mention any references from sub-agent answers in your response.
+""".strip()
+
+
+class SubAgentToolConfig(BaseToolConfig):
+    model_config = get_configuration_dict()
+
+    assistant_id: str = Field(
+        default="",
+        description="The unique identifier of the assistant to use for the sub-agent.",
+    )
+    chat_id: str | None = Field(
+        default=None,
+        description="The chat ID to use for the sub-agent conversation. If None, a new chat will be created.",
+    )
+    reuse_chat: bool = Field(
+        default=True,
+        description="Whether to reuse the existing chat or create a new one for each sub-agent call.",
+    )
+
+    tool_description_for_system_prompt: str = Field(
+        default="",
+        description="Description of the tool that will be included in the system prompt.",
+    )
+    tool_description: str = Field(
+        default="",
+        description="Description of the tool that will be included in the tools sent to the model.",
+    )
+    param_description_sub_agent_user_message: str = Field(
+        default=DEFAULT_PARAM_DESCRIPTION_SUB_AGENT_USER_MESSAGE,
+        description="Description of the user message parameter that will be sent to the model.",
+    )
+    tool_format_information_for_system_prompt: str = Field(
+        default=DEFAULT_FORMAT_INFORMATION_SUB_AGENT_SYSTEM_MESSAGE,
+        description="Format information that will be included in the system prompt to guide response formatting.",
+    )
+    tool_description_for_user_prompt: str = Field(
+        default="",
+        description="Description of the tool that will be included in the user prompt.",
+    )
+    tool_format_information_for_user_prompt: str = Field(
+        default="",
+        description="Format information that will be included in the user prompt to guide response formatting.",
+    )
+
+    poll_interval: float = Field(
+        default=1.0,
+        description="Time interval in seconds between polling attempts when waiting for sub-agent response.",
+    )
+    max_wait: float = Field(
+        default=120.0,
+        description="Maximum time in seconds to wait for the sub-agent response before timing out.",
+    )

unique_toolkit/agentic/tools/a2a/{service.py → tool/service.py}

@@ -1,3 +1,5 @@
+import asyncio
+import contextlib
 from typing import Protocol, override
 
 import unique_sdk
@@ -5,16 +7,16 @@ 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.memory import (
+from unique_toolkit.agentic.tools.a2a.tool._memory import (
     get_sub_agent_short_term_memory_manager,
 )
-from unique_toolkit.agentic.tools.a2a.schema import (
+from unique_toolkit.agentic.tools.a2a.tool._schema import (
     SubAgentShortTermMemorySchema,
     SubAgentToolInput,
 )
+from unique_toolkit.agentic.tools.a2a.tool.config import (
+    SubAgentToolConfig,
+)
 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
@@ -34,10 +36,17 @@ 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,
+        sub_agent_assistant_id: str,
+        sequence_number: int,
     ) -> None: ...
 
+    """
+    Notify the subscriber that a sub agent response has been received.
+    Important: The subscriber should NOT modify the response in place.
+    The sequence number is a 1-indexed counter that is incremented for each concurrent run of the same sub agent.
+    """
+
 
 class SubAgentTool(Tool[SubAgentToolConfig]):
     name: str = "SubAgentTool"
@@ -66,12 +75,18 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
         self._subscribers: list[SubAgentResponseSubscriber] = []
         self._should_run_evaluation = False
 
-    def display_name(self) -> str:
-        return self._display_name
+        # Synchronization state
+        self._sequence_number = 1
+        self._lock = asyncio.Lock()
 
     def subscribe(self, subscriber: SubAgentResponseSubscriber) -> None:
         self._subscribers.append(subscriber)
 
+    @override
+    def display_name(self) -> str:
+        return self._display_name
+
+    @override
     def tool_description(self) -> LanguageModelToolDescription:
         tool_input_model_with_description = create_model(
             "SubAgentToolInput",
@@ -87,27 +102,98 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             parameters=tool_input_model_with_description,
         )
 
+    @override
     def tool_description_for_system_prompt(self) -> str:
         return self.config.tool_description_for_system_prompt
 
+    @override
     def tool_format_information_for_system_prompt(self) -> str:
         return self.config.tool_format_information_for_system_prompt
 
+    @override
     def tool_description_for_user_prompt(self) -> str:
         return self.config.tool_description_for_user_prompt
 
+    @override
     def tool_format_information_for_user_prompt(self) -> str:
         return self.config.tool_format_information_for_user_prompt
 
+    @override
     def evaluation_check_list(self) -> list[EvaluationMetricName]:
         return [EvaluationMetricName.SUB_AGENT] if self._should_run_evaluation else []
 
+    @override
     def get_evaluation_checks_based_on_tool_response(
         self,
         tool_response: ToolCallResponse,
     ) -> list[EvaluationMetricName]:
         return []
 
+    @override
+    async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
+        tool_input = SubAgentToolInput.model_validate(tool_call.arguments)
+
+        if self._lock.locked():
+            await self._notify_progress(
+                tool_call=tool_call,
+                message=f"Waiting for another run of `{self.display_name()}` to finish",
+                state=ProgressState.STARTED,
+            )
+
+        # When reusing the chat id, executing the sub agent in parrallel leads to race conditions and undefined behavior.
+        # To avoid this, we use a lock to serialize the execution of the same sub agent.
+        context = self._lock if self.config.reuse_chat else contextlib.nullcontext()
+
+        async with context:
+            sequence_number = self._sequence_number
+            self._sequence_number += 1
+
+            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 self._execute_and_handle_timeout(
+                tool_user_message=tool_input.user_message,
+                chat_id=chat_id,
+                tool_call=tool_call,
+            )
+
+            self._should_run_evaluation |= (
+                response["assessment"] is not None and len(response["assessment"]) > 0
+            )  # Run evaluation if any sub agent returned an assessment
+
+            self._notify_subscribers(response, sequence_number)
+
+            if chat_id is None:
+                await self._save_chat_id(response["chatId"])
+
+            if response["text"] is None:
+                raise ValueError("No response returned from sub agent")
+
+            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,
+        agent_chunks_handler: AgentChunksHandler,
+    ) -> LanguageModelMessage: ...  # Empty as method is deprecated
+
     async def _get_chat_id(self) -> str | None:
         if not self.config.reuse_chat:
             return None
@@ -145,8 +231,21 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
                 state=state,
             )
 
+    def _notify_subscribers(
+        self, response: unique_sdk.Space.Message, sequence_number: int
+    ) -> None:
+        for subsciber in self._subscribers:
+            subsciber.notify_sub_agent_response(
+                sub_agent_assistant_id=self.config.assistant_id,
+                response=response,
+                sequence_number=sequence_number,
+            )
+
     async def _execute_and_handle_timeout(
-        self, tool_user_message: str, chat_id: str, tool_call: LanguageModelFunction
+        self,
+        tool_user_message: str,
+        chat_id: str | None,
+        tool_call: LanguageModelFunction,
     ) -> unique_sdk.Space.Message:
         try:
             return await send_message_and_wait_for_completion(
@@ -170,61 +269,5 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
                 "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 self._execute_and_handle_timeout(
-            tool_user_message=tool_input.user_message,  # type: ignore
-            chat_id=chat_id,  # type: ignore
-            tool_call=tool_call,
-        )
-
-        self._should_run_evaluation = (
-            response["assessment"] is not None and len(response["assessment"]) > 0
-        )
-
-        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")
-
-        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,
-        agent_chunks_handler: AgentChunksHandler,
-    ) -> LanguageModelMessage: ...  # Empty as method is deprecated
-
 
 ToolFactory.register_tool(SubAgentTool, SubAgentToolConfig)