unique_toolkit 0.7.9__py3-none-any.whl → 1.33.3__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,158 @@
+import re
+from enum import StrEnum
+from typing import Annotated, Generic, Literal, TypeVar
+
+from pydantic import Field
+from pydantic.main import BaseModel
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit.agentic.tools.schemas import BaseToolConfig
+
+
+class SubAgentSystemReminderType(StrEnum):
+    FIXED = "fixed"
+    REGEXP = "regexp"
+    REFERENCE = "reference"
+
+
+T = TypeVar("T", bound=SubAgentSystemReminderType)
+
+
+class SystemReminderConfig(BaseModel, Generic[T]):
+    model_config = get_configuration_dict()
+
+    type: T
+
+
+_SYSTEM_REMINDER_FIELD_DESCRIPTION = """
+The reminder to add to the tool response. The reminder can be a Jinja template and can contain the following placeholders:
+- {{ display_name }}: The display name of the sub agent.
+- {{ tool_name }}: The tool name. 
+""".strip()
+
+
+class ReferenceSystemReminderConfig(SystemReminderConfig):
+    type: Literal[SubAgentSystemReminderType.REFERENCE] = (
+        SubAgentSystemReminderType.REFERENCE
+    )
+    reminder: str = Field(
+        default="Rememeber to properly reference EACH fact from sub agent {{ display_name }}'s response with the correct format INLINE.",
+        description=_SYSTEM_REMINDER_FIELD_DESCRIPTION,
+    )
+
+
+class FixedSystemReminderConfig(SystemReminderConfig):
+    type: Literal[SubAgentSystemReminderType.FIXED] = SubAgentSystemReminderType.FIXED
+    reminder: str = Field(
+        description=_SYSTEM_REMINDER_FIELD_DESCRIPTION,
+    )
+
+
+_REGEXP_DETECTED_REMINDER_FIELD_DESCRIPTION = """
+The reminder to add to the tool response. The reminder can be a Jinja template and can contain the following placeholders:
+- {{ display_name }}: The display name of the sub agent.
+- {{ tool_name }}: The tool name. 
+- {{ text_matches }}: Will be replaced with the portions of the text that triggered the reminder.
+""".strip()
+
+
+class RegExpDetectedSystemReminderConfig(SystemReminderConfig):
+    """A system reminder that is only added if the sub agent response matches a regular expression."""
+
+    type: Literal[SubAgentSystemReminderType.REGEXP] = SubAgentSystemReminderType.REGEXP
+
+    regexp: re.Pattern[str] = Field(
+        description="The regular expression to use to detect whether the system reminder should be added.",
+    )
+    reminder: str = Field(
+        description=_REGEXP_DETECTED_REMINDER_FIELD_DESCRIPTION,
+    )
+
+
+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.",
+    )
+
+    tool_input_json_schema: str | None = Field(
+        default=None,
+        description="A custom JSON schema to send to the llm as the tool input schema.",
+    )
+    returns_content_chunks: bool = Field(
+        default=False,
+        description="If set, the sub-agent response will be interpreted as a list of content chunks.",
+    )
+
+    system_reminders_config: list[
+        Annotated[
+            FixedSystemReminderConfig
+            | RegExpDetectedSystemReminderConfig
+            | ReferenceSystemReminderConfig,
+            Field(discriminator="type"),
+        ]
+    ] = Field(
+        default=[],
+        description="Configuration for the system reminders to add to the tool response.",
+    )

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

@@ -0,0 +1,393 @@
+import asyncio
+import contextlib
+import json
+import logging
+import re
+from datetime import datetime
+from typing import cast, override
+
+import unique_sdk
+from pydantic import Field, TypeAdapter, 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._common.utils.jinja.render import render_template
+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 (
+    RegExpDetectedSystemReminderConfig,
+    SubAgentSystemReminderType,
+    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.content import ContentChunk
+from unique_toolkit.language_model import (
+    LanguageModelFunction,
+    LanguageModelToolDescription,
+)
+
+logger = logging.getLogger(__name__)
+
+_ContentChunkList = TypeAdapter(list[ContentChunk])
+
+
+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:
+        if self.config.tool_input_json_schema is not None:
+            return LanguageModelToolDescription(
+                name=self.name,
+                description=self.config.tool_description,
+                parameters=json.loads(self.config.tool_input_json_schema),
+            )
+
+        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:
+        if self.config.tool_input_json_schema is not None:
+            tool_input = json.dumps(tool_call.arguments)
+        else:
+            tool_input = SubAgentToolInput.model_validate(
+                tool_call.arguments
+            ).user_message
+
+        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,
+                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,
+                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")
+
+            if self.config.returns_content_chunks:
+                content = ""
+                content_chunks = _ContentChunkList.validate_json(response["text"])
+            else:
+                content = self._prepare_response_references(
+                    response=response["text"],
+                    sequence_number=sequence_number,
+                )
+                content_chunks = None
+
+            await self._notify_progress(
+                tool_call=tool_call,
+                message=tool_input,
+                state=ProgressState.FINISHED,
+            )
+
+            return ToolCallResponse(
+                id=tool_call.id,
+                name=tool_call.name,
+                content=_format_response(
+                    tool_name=self.name,
+                    text=content,
+                    system_reminders=self._get_system_reminders(response),
+                ),
+                content_chunks=content_chunks,
+            )
+
+    def _get_system_reminders(self, message: unique_sdk.Space.Message) -> list[str]:
+        reminders = []
+        for reminder_config in self.config.system_reminders_config:
+            if reminder_config.type == SubAgentSystemReminderType.FIXED:
+                reminders.append(
+                    render_template(
+                        reminder_config.reminder,
+                        display_name=self.display_name(),
+                        tool_name=self.name,
+                    )
+                )
+            elif (
+                reminder_config.type == SubAgentSystemReminderType.REFERENCE
+                and self.config.use_sub_agent_references
+                and message["references"] is not None
+                and len(message["references"]) > 0
+            ):
+                reminders.append(
+                    render_template(
+                        reminder_config.reminder,
+                        display_name=self.display_name(),
+                        tool_name=self.name,
+                    )
+                )
+            elif (
+                reminder_config.type == SubAgentSystemReminderType.REGEXP
+                and message["text"] is not None
+            ):
+                reminder_config = cast(
+                    RegExpDetectedSystemReminderConfig, reminder_config
+                )
+                text_matches = [
+                    match.group(0)
+                    for match in reminder_config.regexp.finditer(message["text"])
+                ]
+                if len(text_matches) > 0:
+                    reminders.append(
+                        render_template(
+                            reminder_config.reminder,
+                            display_name=self.display_name(),
+                            tool_name=self.name,
+                            text_matches=text_matches,
+                        )
+                    )
+
+        return reminders
+
+    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
+
+
+def _format_response(tool_name: str, text: str, system_reminders: list[str]) -> str:
+    if len(system_reminders) == 0:
+        return text
+
+    reponse_key = f"{tool_name} response"
+    response = {reponse_key: text, "SYSTEM_REMINDERS": system_reminders}
+
+    return json.dumps(response, indent=2)
+
+
+ToolFactory.register_tool(SubAgentTool, SubAgentToolConfig)