unique_toolkit 0.8.29__py3-none-any.whl → 0.8.30__py3-none-any.whl

unique_toolkit/tools/a2a/__init__.py

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

unique_toolkit/tools/a2a/config.py

@@ -0,0 +1,29 @@
+from unique_toolkit.tools.schemas import BaseToolConfig
+
+from unique_toolkit.tools.config import get_configuration_dict
+
+
+DEFAULT_PARAM_DESCRIPTION_SUB_AGENT_USER_MESSAGE = """
+This is the message that will be sent to the sub-agent.
+""".strip()
+
+
+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

unique_toolkit/tools/a2a/manager.py

@@ -0,0 +1,50 @@
+from logging import Logger
+
+from unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.tools.a2a.service import SubAgentTool, ToolProgressReporter
+
+from unique_toolkit.tools.a2a.config import SubAgentToolConfig
+from unique_toolkit.tools.config import ToolBuildConfig
+from unique_toolkit.tools.schemas import BaseToolConfig
+from unique_toolkit.tools.tool import Tool
+
+
+class A2AManager:
+    def __init__(
+        self,
+        logger: Logger,
+        tool_progress_reporter: ToolProgressReporter,
+    ):
+        self._logger = logger
+        self._tool_progress_reporter = tool_progress_reporter
+
+    def get_all_sub_agents(
+        self, tool_configs: list[ToolBuildConfig], event: ChatEvent
+    ) -> tuple[list[ToolBuildConfig], list[Tool[BaseToolConfig]]]:
+        sub_agents = []
+
+        for tool_config in tool_configs:
+            if not tool_config.is_sub_agent:
+                continue
+
+            if not isinstance(tool_config.configuration, SubAgentToolConfig):
+                self._logger.error(
+                    "tool_config.configuration must be of type SubAgentToolConfig"
+                )
+                continue
+
+            sub_agent_tool_config: SubAgentToolConfig = tool_config.configuration
+
+            sub_agents.append(
+                SubAgentTool(
+                    configuration=sub_agent_tool_config,
+                    event=event,
+                    tool_progress_reporter=self._tool_progress_reporter,
+                )
+            )
+
+        filtered_tool_config = [
+            tool_config for tool_config in tool_configs if not tool_config.is_sub_agent
+        ]
+
+        return filtered_tool_config, sub_agents

unique_toolkit/tools/a2a/memory.py

@@ -0,0 +1,27 @@
+from unique_toolkit import ShortTermMemoryService
+
+from unique_toolkit.short_term_memory.persistent_short_term_memory_manager import (
+    PersistentShortMemoryManager,
+)
+from unique_toolkit.tools.a2a.schema import SubAgentShortTermMemorySchema
+
+
+def _get_short_term_memory_name(assistant_id: str) -> str:
+    return f"sub_agent_chat_id_{assistant_id}"
+
+
+def get_sub_agent_short_term_memory_manager(
+    company_id: str, user_id: str, chat_id: str, assistant_id: str
+) -> PersistentShortMemoryManager[SubAgentShortTermMemorySchema]:
+    short_term_memory_service = ShortTermMemoryService(
+        company_id=company_id,
+        user_id=user_id,
+        chat_id=chat_id,
+        message_id=None,
+    )
+    short_term_memory_manager = PersistentShortMemoryManager(
+        short_term_memory_service=short_term_memory_service,
+        short_term_memory_schema=SubAgentShortTermMemorySchema,
+        short_term_memory_name=_get_short_term_memory_name(assistant_id),
+    )
+    return short_term_memory_manager

unique_toolkit/tools/a2a/schema.py

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

unique_toolkit/tools/a2a/service.py

@@ -0,0 +1,154 @@
+from pydantic import Field, create_model
+from unique_sdk.utils.chat_in_space import send_message_and_wait_for_completion
+from unique_toolkit.app import ChatEvent
+from unique_toolkit.language_model import (
+    LanguageModelFunction,
+    LanguageModelMessage,
+    LanguageModelToolDescription,
+)
+
+from unique_toolkit.evaluators.schemas import EvaluationMetricName
+from unique_toolkit.tools.schemas import ToolCallResponse
+from unique_toolkit.tools.a2a.config import SubAgentToolConfig
+from unique_toolkit.tools.a2a.memory import (
+    get_sub_agent_short_term_memory_manager,
+)
+from unique_toolkit.tools.a2a.schema import (
+    SubAgentShortTermMemorySchema,
+    SubAgentToolCallResponse,
+    SubAgentToolInput,
+)
+from unique_toolkit.tools.tool import Tool
+
+from unique_toolkit.tools.tool_progress_reporter import (
+    ProgressState,
+    ToolProgressReporter,
+)
+
+
+class SubAgentTool(Tool[SubAgentToolConfig]):
+    name: str = "SubAgentTool"
+
+    def __init__(
+        self,
+        configuration: SubAgentToolConfig,
+        event: ChatEvent,
+        tool_progress_reporter: ToolProgressReporter | None = None,
+    ):
+        super().__init__(configuration, event, tool_progress_reporter)
+        self._user_id = event.user_id
+        self._company_id = event.company_id
+        self.name = configuration.name
+
+        self._short_term_memory_manager = get_sub_agent_short_term_memory_manager(
+            company_id=self._company_id,
+            user_id=self._user_id,
+            chat_id=event.payload.chat_id,
+            assistant_id=self.config.assistant_id,
+        )
+
+    def tool_description(self) -> LanguageModelToolDescription:
+        tool_input_model_with_description = create_model(
+            "SubAgentToolInput",
+            user_message=(
+                str,
+                Field(description=self.config.param_description_sub_agent_user_message),
+            ),
+        )
+
+        return LanguageModelToolDescription(
+            name=self.name,
+            description=self.config.tool_description,
+            parameters=tool_input_model_with_description,
+        )
+
+    def tool_description_for_system_prompt(self) -> str:
+        return self.config.tool_description_for_system_prompt
+
+    def tool_format_information_for_system_prompt(self) -> str:
+        return self.config.tool_format_information_for_system_prompt
+
+    def tool_description_for_user_prompt(self) -> str:
+        return self.config.tool_description_for_user_prompt
+
+    def tool_format_information_for_user_prompt(self) -> str:
+        return self.config.tool_format_information_for_user_prompt
+
+    def evaluation_check_list(self) -> list[EvaluationMetricName]:
+        return []
+
+    def get_evaluation_checks_based_on_tool_response(
+        self,
+        tool_response: ToolCallResponse,
+    ) -> list[EvaluationMetricName]:
+        return []
+
+    async def _get_chat_id(self) -> str | None:
+        if not self.config.reuse_chat:
+            return None
+
+        if self.config.chat_id is not None:
+            return self.config.chat_id
+
+        # Check if there is a saved chat id in short term memory
+        short_term_memory = await self._short_term_memory_manager.load_async()
+
+        if short_term_memory is not None:
+            return short_term_memory.chat_id
+
+        return None
+
+    async def _save_chat_id(self, chat_id: str) -> None:
+        if not self.config.reuse_chat:
+            return
+
+        await self._short_term_memory_manager.save_async(
+            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:
+            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,
+            )
+
+        # 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
+            chat_id=chat_id,  # type: ignore
+            poll_interval=self.config.poll_interval,
+            max_wait=self.config.max_wait,
+        )
+
+        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")
+
+        self._text = response["text"]
+        return ToolCallResponse(
+            id=tool_call.id,  # type: ignore
+            name=tool_call.name,
+            content=response["text"],
+        )
+
+    def get_tool_call_result_for_loop_history(
+        self,
+        tool_response: ToolCallResponse,
+    ) -> LanguageModelMessage:
+        return ToolCallResponse(
+            id=tool_call.id,  # type: ignore
+            name=tool_call.name,
+            content=response["text"],
+        )

unique_toolkit/tools/config.py

@@ -72,6 +72,7 @@ class ToolBuildConfig(BaseModel):
         default=False,
         description="This tool must be chosen by the user and no other tools are used for this iteration.",
     )
+    is_sub_agent: bool = False
 
     is_enabled: bool = Field(default=True)
 

unique_toolkit/tools/mcp/models.py

@@ -1,5 +1,4 @@
 from typing import Any, Dict, Optional
-
 from unique_toolkit.tools.schemas import BaseToolConfig
 
 

unique_toolkit/tools/test/test_mcp_manager.py

@@ -15,6 +15,8 @@ from unique_toolkit.tools.tool import Tool
 from unique_toolkit.tools.tool_manager import ToolManager, ToolManagerConfig
 from unique_toolkit.tools.tool_progress_reporter import ToolProgressReporter
 
+from unique_toolkit.tools.a2a.manager import A2AManager
+
 
 class MockParameters(BaseModel):
     pass
@@ -161,6 +163,14 @@ class TestMCPManager:
             event=self.event,
             tool_progress_reporter=tool_progress_reporter,
         )
+    
+    @pytest.fixture
+    def a2a_manager(self,tool_progress_reporter):
+        """Create MCP manager fixture"""
+        return A2AManager(
+            logger=self.logger,
+            tool_progress_reporter=tool_progress_reporter,
+        )
 
     @pytest.fixture
     def tool_manager_config(self, internal_tools):
@@ -168,19 +178,23 @@ class TestMCPManager:
         return ToolManagerConfig(tools=internal_tools, max_tool_calls=10)
 
     @pytest.fixture
-    def tool_manager(self, tool_manager_config, mcp_manager, tool_progress_reporter):
+    def tool_manager(self, tool_manager_config, mcp_manager, a2a_manager,tool_progress_reporter):
         """Create tool manager fixture"""
+
+
         return ToolManager(
             logger=self.logger,
             config=tool_manager_config,
             event=self.event,
             tool_progress_reporter=tool_progress_reporter,
             mcp_manager=mcp_manager,
+            a2a_manager=a2a_manager
         )
 
     def test_tool_manager_initialization(self, tool_manager):
         """Test tool manager is initialized correctly"""
         assert tool_manager is not None
+
         assert (
             len(tool_manager.get_tools()) >= 2
         )  # Should have both internal and MCP tools
@@ -255,12 +269,20 @@ class TestMCPManager:
         """Test the _init__tools method behavior with different scenarios"""
 
         # Test 1: Normal initialization with both tool types
+
+        
+        a2a_manager = A2AManager(
+            logger=self.logger,
+            tool_progress_reporter=tool_progress_reporter,
+        )
+
         tool_manager = ToolManager(
             logger=self.logger,
             config=tool_manager_config,
             event=self.event,
             tool_progress_reporter=tool_progress_reporter,
             mcp_manager=mcp_manager,
+            a2a_manager=a2a_manager
         )
 
         # Verify both tools are loaded
@@ -285,12 +307,20 @@ class TestMCPManager:
         event_with_disabled.payload.tool_choices = ["internal_search", "mcp_test_tool"]
         event_with_disabled.payload.disabled_tools = ["internal_search"]
 
+        
+        
+        a2a_manager = A2AManager(
+            logger=self.logger,
+            tool_progress_reporter=tool_progress_reporter,
+        )
+
         tool_manager = ToolManager(
             logger=self.logger,
             config=tool_manager_config,
             event=event_with_disabled,
             tool_progress_reporter=tool_progress_reporter,
             mcp_manager=mcp_manager,
+            a2a_manager=a2a_manager,
         )
 
         # Should only have MCP tool, internal tool should be filtered out
@@ -315,12 +345,20 @@ class TestMCPManager:
         event_with_limited_choices.payload.tool_choices = ["internal_search"]
         event_with_limited_choices.payload.disabled_tools = []
 
+
+        
+        a2a_manager = A2AManager(
+            logger=self.logger,
+            tool_progress_reporter=tool_progress_reporter,
+        )
+
         tool_manager = ToolManager(
             logger=self.logger,
             config=tool_manager_config,
             event=event_with_limited_choices,
             tool_progress_reporter=tool_progress_reporter,
             mcp_manager=mcp_manager,
+            a2a_manager=a2a_manager,
         )
 
         # Should only have internal search tool
@@ -348,12 +386,20 @@ class TestMCPManager:
             tools=[exclusive_tool_config], max_tool_calls=10
         )
 
+        
+        
+        a2a_manager = A2AManager(
+            logger=self.logger,
+            tool_progress_reporter=tool_progress_reporter,
+        )
+
         tool_manager = ToolManager(
             logger=self.logger,
             config=config_with_exclusive,
             event=self.event,
             tool_progress_reporter=tool_progress_reporter,
             mcp_manager=mcp_manager,
+            a2a_manager=a2a_manager
         )
 
         # Should only have the exclusive tool, MCP tools should be ignored
@@ -383,12 +429,20 @@ class TestMCPManager:
             tools=[disabled_tool_config], max_tool_calls=10
         )
 
+
+        
+        a2a_manager = A2AManager(
+            logger=self.logger,
+            tool_progress_reporter=tool_progress_reporter,
+        )
+
         tool_manager = ToolManager(
             logger=self.logger,
             config=config_with_disabled,
             event=self.event,
             tool_progress_reporter=tool_progress_reporter,
             mcp_manager=mcp_manager,
+            a2a_manager=a2a_manager
         )
 
         # Should only have MCP tool, disabled internal tool should be filtered out

unique_toolkit/tools/tool_manager.py

@@ -2,7 +2,10 @@ import asyncio
 from logging import Logger, getLogger
 from typing import Any
 
+
 from pydantic import BaseModel, Field
+from unique_toolkit.tools.a2a.manager import A2AManager
+
 
 from unique_toolkit.app.schemas import ChatEvent
 from unique_toolkit.evals.schemas import EvaluationMetricName
@@ -71,6 +74,7 @@ class ToolManager:
         event: ChatEvent,
         tool_progress_reporter: ToolProgressReporter,
         mcp_manager: MCPManager,
+        a2a_manager: A2AManager,
     ):
         self._logger = logger
         self._config = config
@@ -81,6 +85,7 @@ class ToolManager:
         # this needs to be a set of strings to avoid duplicates
         self._tool_evaluation_check_list: set[EvaluationMetricName] = set()
         self._mcp_manager = mcp_manager
+        self._a2a_manager = a2a_manager
         self._init__tools(event)
 
     def _init__tools(self, event: ChatEvent) -> None:
@@ -90,6 +95,10 @@ class ToolManager:
         self._logger.info(f"Tool choices: {tool_choices}")
         self._logger.info(f"Tool configs: {tool_configs}")
 
+        tool_configs, sub_agents = self._a2a_manager.get_all_sub_agents(
+            tool_configs, event
+        )
+
         # Build internal tools from configurations
         internal_tools = [
             ToolFactory.build_tool_with_settings(
@@ -105,7 +114,7 @@ class ToolManager:
         # Get MCP tools (these are already properly instantiated)
         mcp_tools = self._mcp_manager.get_all_mcp_tools()
         # Combine both types of tools
-        self.available_tools = internal_tools + mcp_tools
+        self.available_tools = internal_tools + mcp_tools + sub_agents
 
         for t in self.available_tools:
             if t.is_exclusive():

{unique_toolkit-0.8.29.dist-info → unique_toolkit-0.8.30.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 0.8.29
+Version: 0.8.30
 Summary: 
 License: Proprietary
 Author: Martin Fadler
@@ -114,6 +114,10 @@ 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).
+
+## [0.8.30] - 2025-08-28
+- Added A2A manager
+
 ## [0.8.29] - 2025-08-27
 - Include `MessageExecution` and `MessageLog` in toolkit 
 
@@ -126,6 +130,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [0.8.26] - 2025-08-27
 - Optimized MCP manager
 
+## [0.8.26] - 2025-08-27
+- Optimized MCP manager
+
 ## [0.8.25] - 2025-08-27
 - Load environment variables automatically from plattform dirs or environment
 - General Endpoint definition utility

{unique_toolkit-0.8.29.dist-info → unique_toolkit-0.8.30.dist-info}/RECORD

@@ -99,23 +99,29 @@ unique_toolkit/short_term_memory/service.py,sha256=5PeVBu1ZCAfyDb2HLVvlmqSbyzBBu
 unique_toolkit/smart_rules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unique_toolkit/smart_rules/compile.py,sha256=cxWjb2dxEI2HGsakKdVCkSNi7VK9mr08w5sDcFCQyWI,9553
 unique_toolkit/thinking_manager/thinking_manager.py,sha256=AJfmrTXTr-DxBnJ2_zYYpYo57kr5deqT0LiZb8UdaDQ,4175
-unique_toolkit/tools/config.py,sha256=E4lFQA4gCi3_j4hcdnd5YUHhbaWTyIGWrgD18QardQw,3796
+unique_toolkit/tools/a2a/__init__.py,sha256=BIPxaqnfTD9IywxrqTLbGnkqNQrfLnKbdzF5ELYzCxM,171
+unique_toolkit/tools/a2a/config.py,sha256=NjLBt8Bmxmh0XSU6W20_c-0jZoj4x8EX7Uodgyg1F7Y,866
+unique_toolkit/tools/a2a/manager.py,sha256=nUEhjQmBh49rHWXAMDiRxR62vaPQdbSaBwaa_upUSEA,1667
+unique_toolkit/tools/a2a/memory.py,sha256=cP4vaVAtL0IrlBT8QlYl72M6HCrds_Q7QRCIrwwP71A,1005
+unique_toolkit/tools/a2a/schema.py,sha256=1R7qIu2l2qnUJDGRuUnZLqPPkHnT3x4d3d4PALoBzcY,296
+unique_toolkit/tools/a2a/service.py,sha256=bQv6gooQyP_wvwEV_7Bhkvq4kHcLXtQfMQWov5UPT9A,5178
+unique_toolkit/tools/config.py,sha256=nYwglwUSb3fxhqYxI83k-qoecuF5Zcol5FMOxHvTNeE,3827
 unique_toolkit/tools/factory.py,sha256=w3uNHuYBIJ330Xi8PTdAkr8G3OMbQH2cBgvk5UT16oE,1253
 unique_toolkit/tools/mcp/__init__.py,sha256=RLF_p-LDRC7GhiB3fdCi4u3bh6V9PY_w26fg61BLyco,122
 unique_toolkit/tools/mcp/manager.py,sha256=lQcCsfCMqW2j6uCwNmgYoQDwjm8YosvCh4-apw3KZKs,2683
-unique_toolkit/tools/mcp/models.py,sha256=f9OXcGadBx2tDlhYFu-7oEkl5p9ppoDKTqZbca_NblQ,714
+unique_toolkit/tools/mcp/models.py,sha256=6-qFq3ip2OhqRC5H6YzhsC27w5CYrWln75dWmeIovRk,713
 unique_toolkit/tools/mcp/tool_wrapper.py,sha256=w7Fbo4FSMYvtgSq7Sqt1dmAPvqHjoBQS-TwnSp0p6j0,9927
 unique_toolkit/tools/schemas.py,sha256=rArQccbfIv7CWcozClAZ-BVlOwAsjpgL8KUab_WeO3k,4817
-unique_toolkit/tools/test/test_mcp_manager.py,sha256=jnVKxqJyeJBDb6LvI8bM9dd8XPxrbXV0VCOHMhCdlug,14398
+unique_toolkit/tools/test/test_mcp_manager.py,sha256=kzEq4cdhV4vI8FvqKBCvMVXSyJRMrUdedg14Ps5g8ko,15678
 unique_toolkit/tools/test/test_tool_progress_reporter.py,sha256=GTtmBqOUo0-4fh_q0lRgxDhwKeankc3FHFD5ULZAm4Y,6299
 unique_toolkit/tools/tool.py,sha256=E0qORpQBrdksAylf_RuPaEbXKi_AlswVGv-67t3DNZI,5879
-unique_toolkit/tools/tool_manager.py,sha256=_Qt5fUGvA82YHn5LvqNJurLYW41jH9endNwsiSeyXyM,9948
+unique_toolkit/tools/tool_manager.py,sha256=9gncHvsv_8bHAeufV99u6V_lskoT5QG8vT_RP5AA_-0,10208
 unique_toolkit/tools/tool_progress_reporter.py,sha256=ixud9VoHey1vlU1t86cW0-WTvyTwMxNSWBon8I11SUk,7955
 unique_toolkit/tools/utils/execution/execution.py,sha256=vjG2Y6awsGNtlvyQAGCTthQ5thWHYnn-vzZXaYLb3QE,7922
 unique_toolkit/tools/utils/source_handling/schema.py,sha256=vzAyf6ZWNexjMO0OrnB8y2glGkvAilmGGQXd6zcDaKw,870
 unique_toolkit/tools/utils/source_handling/source_formatting.py,sha256=C7uayNbdkNVJdEARA5CENnHtNY1SU6etlaqbgHNyxaQ,9152
 unique_toolkit/tools/utils/source_handling/tests/test_source_formatting.py,sha256=oM5ZxEgzROrnX1229KViCAFjRxl9wCTzWZoinYSHleM,6979
-unique_toolkit-0.8.29.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
-unique_toolkit-0.8.29.dist-info/METADATA,sha256=14-dyP0BCriYbAGs3mQBbGUG-SfuOmJDX6lq1Puf2Gg,29245
-unique_toolkit-0.8.29.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-unique_toolkit-0.8.29.dist-info/RECORD,,
+unique_toolkit-0.8.30.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
+unique_toolkit-0.8.30.dist-info/METADATA,sha256=3HA-gaW3UOJjBHHeFK9DyuMzHzFi529HIABflY9Y5CI,29342
+unique_toolkit-0.8.30.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+unique_toolkit-0.8.30.dist-info/RECORD,,