unique_toolkit 0.7.7__py3-none-any.whl → 1.23.0__py3-none-any.whl

unique_toolkit/agentic/tools/a2a/response_watcher/service.py

@@ -0,0 +1,91 @@
+import datetime
+import json
+from typing import NamedTuple
+
+import unique_sdk
+
+
+def _clone_message(
+    message: unique_sdk.Space.Message,
+) -> unique_sdk.Space.Message:
+    # copy.deepcopy does not work for instances of UniqueObject
+    return json.loads(json.dumps(message))
+
+
+class SubAgentResponse(NamedTuple):
+    assistant_id: str
+    name: str
+    sequence_number: int
+    message: unique_sdk.Space.Message
+    timestamp: datetime.datetime
+
+    def clone(self) -> "SubAgentResponse":
+        return SubAgentResponse(
+            assistant_id=self.assistant_id,
+            name=self.name,
+            sequence_number=self.sequence_number,
+            message=_clone_message(self.message),
+            timestamp=self.timestamp,
+        )
+
+
+class SubAgentResponseWatcher:
+    """
+    Save and retrieve sub agent responses immutably.
+    """
+
+    def __init__(self) -> None:
+        self._response_registry: dict[str, list[SubAgentResponse]] = {}
+
+    def notify_response(
+        self,
+        assistant_id: str,
+        name: str,
+        sequence_number: int,
+        response: unique_sdk.Space.Message,
+        timestamp: datetime.datetime,
+    ) -> None:
+        if assistant_id not in self._response_registry:
+            self._response_registry[assistant_id] = []
+
+        response = _clone_message(response)
+
+        self._response_registry[assistant_id].append(
+            SubAgentResponse(
+                assistant_id=assistant_id,
+                name=name,
+                sequence_number=sequence_number,
+                message=response,
+                timestamp=timestamp,
+            )
+        )
+
+    def get_responses(self, assistant_id: str) -> list[SubAgentResponse]:
+        return _sort_responses(  # Always return a consistent order
+            [
+                response.clone()
+                for response in self._response_registry.get(assistant_id, [])
+            ],
+        )
+
+    def get_all_responses(self) -> list[SubAgentResponse]:
+        return _sort_responses(
+            [
+                response.clone()
+                for sub_agent_responses in self._response_registry.values()
+                for response in sub_agent_responses
+            ],
+        )
+
+
+def _sort_responses(
+    responses: list[SubAgentResponse],
+) -> list[SubAgentResponse]:
+    return sorted(
+        responses,
+        key=lambda response: (
+            response.timestamp,
+            response.assistant_id,
+            response.sequence_number,
+        ),
+    )

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/tool/_memory.py

@@ -0,0 +1,26 @@
+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.tool._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/agentic/tools/a2a/tool/_schema.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+
+class SubAgentToolInput(BaseModel):
+    user_message: str
+
+
+class SubAgentShortTermMemorySchema(BaseModel):
+    chat_id: str

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

@@ -0,0 +1,73 @@
+from typing import Literal
+
+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()
+
+
+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.",
+    )
+    use_sub_agent_references: bool = Field(
+        default=True,
+        description="Whether this sub agent's references should be used in the main agent's response.",
+    )
+    forced_tools: list[str] | None = Field(
+        default=None,
+        description="The list of tool names that will be forced to be called for this sub-agent.",
+    )
+
+    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="",
+        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.",
+    )
+    stop_condition: Literal["stoppedStreamingAt", "completedAt"] = Field(
+        default="completedAt",
+        description="The condition that will be used to stop the polling for the sub-agent response.",
+    )

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

@@ -0,0 +1,306 @@
+import asyncio
+import contextlib
+import logging
+import re
+from datetime import datetime
+from typing import 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._common.referencing import (
+    get_all_ref_numbers,
+    remove_all_refs,
+    replace_ref_number,
+)
+from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
+from unique_toolkit.agentic.tools.a2a.response_watcher import SubAgentResponseWatcher
+from unique_toolkit.agentic.tools.a2a.tool._memory import (
+    get_sub_agent_short_term_memory_manager,
+)
+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.factory import ToolFactory
+from unique_toolkit.agentic.tools.schemas import ToolCallResponse
+from unique_toolkit.agentic.tools.tool import Tool
+from unique_toolkit.agentic.tools.tool_progress_reporter import (
+    ProgressState,
+    ToolProgressReporter,
+)
+from unique_toolkit.app import ChatEvent
+from unique_toolkit.language_model import (
+    LanguageModelFunction,
+    LanguageModelToolDescription,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class SubAgentTool(Tool[SubAgentToolConfig]):
+    name: str = "SubAgentTool"
+
+    def __init__(
+        self,
+        configuration: SubAgentToolConfig,
+        event: ChatEvent,
+        tool_progress_reporter: ToolProgressReporter | None = None,
+        name: str = "SubAgentTool",
+        display_name: str = "SubAgentTool",
+        response_watcher: SubAgentResponseWatcher | None = None,
+    ):
+        super().__init__(configuration, event, tool_progress_reporter)
+        self._user_id = event.user_id
+        self._company_id = event.company_id
+
+        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,
+            user_id=self._user_id,
+            chat_id=event.payload.chat_id,
+            assistant_id=self.config.assistant_id,
+        )
+        self._should_run_evaluation = False
+
+        self._response_watcher = response_watcher
+
+        # Synchronization state
+        self._sequence_number = 1
+        self._lock = asyncio.Lock()
+
+    @staticmethod
+    def get_sub_agent_reference_format(
+        name: str, sequence_number: int, reference_number: int
+    ) -> str:
+        return f"<sup><name>{name} {sequence_number}</name>{reference_number}</sup>"
+
+    @staticmethod
+    def get_sub_agent_reference_re(
+        name: str, sequence_number: int, reference_number: int
+    ) -> str:
+        return rf"<sup>\s*<name>\s*{re.escape(name)}\s*{sequence_number}\s*</name>\s*{reference_number}\s*</sup>"
+
+    @override
+    def display_name(self) -> str:
+        return self._display_name
+
+    @override
+    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,
+        )
+
+    @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)
+        timestamp = datetime.now()
+
+        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_watcher(response, sequence_number, timestamp)
+
+            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")
+
+            response_text_with_references = self._prepare_response_references(
+                response=response["text"],
+                sequence_number=sequence_number,
+            )
+
+            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_with_references,
+            )
+
+    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
+
+    def _prepare_response_references(self, response: str, sequence_number: int) -> str:
+        if not self.config.use_sub_agent_references:
+            # Remove all references from the response
+            response = remove_all_refs(response)
+            return response
+
+        for ref_number in get_all_ref_numbers(response):
+            reference = self.get_sub_agent_reference_format(
+                name=self.name,
+                sequence_number=sequence_number,
+                reference_number=ref_number,
+            )
+            response = replace_ref_number(
+                text=response, ref_number=ref_number, replacement=reference
+            )
+        return response
+
+    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 _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=self._display_name,
+                message=message,
+                state=state,
+            )
+
+    def _notify_watcher(
+        self,
+        response: unique_sdk.Space.Message,
+        sequence_number: int,
+        timestamp: datetime,
+    ) -> None:
+        if self._response_watcher is not None:
+            self._response_watcher.notify_response(
+                assistant_id=self.config.assistant_id,
+                name=self.name,
+                sequence_number=sequence_number,
+                response=response,
+                timestamp=timestamp,
+            )
+        else:
+            logger.warning(
+                "No response watcher found for sub agent %s (assistant_id: %s)",
+                self.name,
+                self.config.assistant_id,
+            )
+
+    async def _execute_and_handle_timeout(
+        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(
+                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,
+                tool_choices=self.config.forced_tools,
+                max_wait=self.config.max_wait,
+                stop_condition=self.config.stop_condition,
+            )
+        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
+
+
+ToolFactory.register_tool(SubAgentTool, SubAgentToolConfig)

unique_toolkit/agentic/tools/agent_chunks_hanlder.py

@@ -0,0 +1,65 @@
+from typing_extensions import deprecated
+
+from unique_toolkit.content.schemas import ContentChunk, ContentReference
+
+
+@deprecated("do not use this its only used in old tools")
+class AgentChunksHandler:
+    def __init__(self):
+        self._tool_chunks = {}
+        self._chunks: list[ContentChunk] = []
+        self._references: list[list[ContentReference]] = []
+
+    @property
+    def chunks(self) -> list[ContentChunk]:
+        return self._chunks
+
+    @property
+    def tool_chunks(self) -> dict:
+        return self._tool_chunks
+
+    def extend(self, chunks: list[ContentChunk]):
+        self._chunks.extend(chunks)
+
+    def replace(self, chunks: list[ContentChunk]):
+        self._chunks = chunks
+
+    def add_references(
+        self,
+        references: list[ContentReference],
+    ):
+        self._references.append(references)
+
+    @property
+    def all_references(
+        self,
+    ) -> list[list[ContentReference]]:
+        return self._references
+
+    @property
+    def latest_references(
+        self,
+    ) -> list[ContentReference]:
+        if not self._references:
+            return []
+        return self._references[-1]
+
+    @property
+    def latest_referenced_chunks(self) -> list[ContentChunk]:
+        if not self._references:
+            return []
+        return self._get_referenced_chunks_from_references(self._references[-1])
+
+    def _get_referenced_chunks_from_references(
+        self,
+        references: list[ContentReference],
+    ) -> list[ContentChunk]:
+        """
+        Get _referenced_chunks by matching sourceId from _references with merged id and chunk_id from _chunks.
+        """
+        referenced_chunks: list[ContentChunk] = []
+        for ref in references:
+            for chunk in self._chunks:
+                if ref.source_id == str(chunk.id) + "_" + str(chunk.chunk_id):
+                    referenced_chunks.append(chunk)
+        return referenced_chunks