unique_toolkit 1.8.1__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/config.py

@@ -1,3 +1,5 @@
+from typing import Literal
+
 from pydantic import Field
 
 from unique_toolkit._common.pydantic_helpers import get_configuration_dict
@@ -7,10 +9,6 @@ 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()
@@ -27,6 +25,14 @@ class SubAgentToolConfig(BaseToolConfig):
         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="",
@@ -41,7 +47,7 @@ class SubAgentToolConfig(BaseToolConfig):
         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,
+        default="",
         description="Format information that will be included in the system prompt to guide response formatting.",
     )
     tool_description_for_user_prompt: str = Field(
@@ -61,3 +67,7 @@ class SubAgentToolConfig(BaseToolConfig):
         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

@@ -1,12 +1,21 @@
 import asyncio
 import contextlib
-from typing import Protocol, override
+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,
 )
@@ -17,7 +26,6 @@ from unique_toolkit.agentic.tools.a2a.tool._schema import (
 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
 from unique_toolkit.agentic.tools.tool import Tool
@@ -30,22 +38,8 @@ from unique_toolkit.language_model import (
     LanguageModelFunction,
     LanguageModelToolDescription,
 )
-from unique_toolkit.language_model.schemas import LanguageModelMessage
 
-
-class SubAgentResponseSubscriber(Protocol):
-    def notify_sub_agent_response(
-        self,
-        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.
-    """
+logger = logging.getLogger(__name__)
 
 
 class SubAgentTool(Tool[SubAgentToolConfig]):
@@ -58,6 +52,7 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
         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
@@ -72,15 +67,25 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             chat_id=event.payload.chat_id,
             assistant_id=self.config.assistant_id,
         )
-        self._subscribers: list[SubAgentResponseSubscriber] = []
         self._should_run_evaluation = False
 
+        self._response_watcher = response_watcher
+
         # Synchronization state
         self._sequence_number = 1
         self._lock = asyncio.Lock()
 
-    def subscribe(self, subscriber: SubAgentResponseSubscriber) -> None:
-        self._subscribers.append(subscriber)
+    @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:
@@ -132,6 +137,7 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
     @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(
@@ -167,7 +173,7 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
                 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)
+            self._notify_watcher(response, sequence_number, timestamp)
 
             if chat_id is None:
                 await self._save_chat_id(response["chatId"])
@@ -175,6 +181,11 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             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,
@@ -184,16 +195,9 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             return ToolCallResponse(
                 id=tool_call.id,  # type: ignore
                 name=tool_call.name,
-                content=response["text"],
+                content=response_text_with_references,
             )
 
-    @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
@@ -209,6 +213,23 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
 
         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
@@ -231,14 +252,25 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
                 state=state,
             )
 
-    def _notify_subscribers(
-        self, response: unique_sdk.Space.Message, sequence_number: int
+    def _notify_watcher(
+        self,
+        response: unique_sdk.Space.Message,
+        sequence_number: int,
+        timestamp: datetime,
     ) -> None:
-        for subsciber in self._subscribers:
-            subsciber.notify_sub_agent_response(
-                sub_agent_assistant_id=self.config.assistant_id,
-                response=response,
+        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(
@@ -255,8 +287,9 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
                 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="completedAt",
+                stop_condition=self.config.stop_condition,
             )
         except TimeoutError as e:
             await self._notify_progress(

unique_toolkit/agentic/tools/config.py

@@ -1,9 +1,10 @@
 import json
 from enum import StrEnum
-from typing import Any, Dict
+from typing import Annotated, Any, Dict
 
 from pydantic import (
     BaseModel,
+    BeforeValidator,
     Field,
     ValidationInfo,
     model_validator,
@@ -33,6 +34,16 @@ class ToolSelectionPolicy(StrEnum):
     BY_USER = "ByUser"
 
 
+def handle_undefined_icon(value: Any) -> ToolIcon:
+    try:
+        if isinstance(value, str):
+            return ToolIcon(value)
+        else:
+            return ToolIcon.BOOK
+    except ValueError:
+        return ToolIcon.BOOK
+
+
 class ToolBuildConfig(BaseModel):
     model_config = get_configuration_dict()
     """Main tool configuration"""
@@ -40,7 +51,10 @@ class ToolBuildConfig(BaseModel):
     name: str
     configuration: BaseToolConfig
     display_name: str = ""
-    icon: ToolIcon = ToolIcon.BOOK
+    icon: Annotated[ToolIcon, BeforeValidator(handle_undefined_icon)] = Field(
+        default=ToolIcon.BOOK,
+        description="The icon name that will be used to display the tool in the user interface.",
+    )
     selection_policy: ToolSelectionPolicy = Field(
         default=ToolSelectionPolicy.BY_USER,
     )

unique_toolkit/agentic/tools/factory.py

@@ -11,6 +11,10 @@ class ToolFactory:
     tool_map: dict[str, type[Tool]] = {}
     tool_config_map: dict[str, Callable] = {}
 
+    @classmethod
+    def register_tool_config(cls, tool_name: str, tool_config: type[BaseToolConfig]):
+        cls.tool_config_map[tool_name] = tool_config
+
     @classmethod
     def register_tool(
         cls,

unique_toolkit/agentic/tools/mcp/tool_wrapper.py

@@ -1,8 +1,8 @@
 import json
+import logging
 from typing import Any, Dict
 
 import unique_sdk
-from pydantic import BaseModel, Field, create_model
 
 from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
 from unique_toolkit.agentic.tools.mcp.models import MCPToolConfig
@@ -13,12 +13,13 @@ from unique_toolkit.agentic.tools.tool_progress_reporter import (
     ToolProgressReporter,
 )
 from unique_toolkit.app.schemas import ChatEvent, McpServer, McpTool
-from unique_toolkit.language_model import LanguageModelMessage
 from unique_toolkit.language_model.schemas import (
     LanguageModelFunction,
     LanguageModelToolDescription,
 )
 
+logger = logging.getLogger(__name__)
+
 
 class MCPToolWrapper(Tool[MCPToolConfig]):
     """Wrapper class for MCP tools that implements the Tool interface"""
@@ -39,39 +40,16 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
     def tool_description(self) -> LanguageModelToolDescription:
         """Convert MCP tool schema to LanguageModelToolDescription"""
         # Create a Pydantic model from the MCP tool's input schema
-        parameters_model = self._create_parameters_model()
+        logger.info(
+            "MCP tool %s schema %s", self._mcp_tool.name, self._mcp_tool.input_schema
+        )
 
         return LanguageModelToolDescription(
             name=self.name,
             description=self._mcp_tool.description or "",
-            parameters=parameters_model,
+            parameters=self._mcp_tool.input_schema,
         )
 
-    def _create_parameters_model(self) -> type[BaseModel]:
-        """Create a Pydantic model from MCP tool's input schema"""
-        properties = self._mcp_tool.input_schema.get("properties", {})
-        required_fields = self._mcp_tool.input_schema.get("required", [])
-
-        # Convert JSON schema properties to Pydantic fields
-        fields = {}
-        for prop_name, prop_schema in properties.items():
-            field_type = self._json_schema_to_python_type(prop_schema)
-            field_description = prop_schema.get("description", "")
-
-            if prop_name in required_fields:
-                fields[prop_name] = (
-                    field_type,
-                    Field(description=field_description),
-                )
-            else:
-                fields[prop_name] = (
-                    field_type,
-                    Field(default=None, description=field_description),
-                )
-
-        # Create dynamic model
-        return create_model(f"{self.name}Parameters", **fields)
-
     def _json_schema_to_python_type(self, schema: Dict[str, Any]) -> type:
         """Convert JSON schema type to Python type"""
         json_type = schema.get("type", "string")
@@ -120,12 +98,6 @@ class MCPToolWrapper(Tool[MCPToolConfig]):
         """Return evaluation checks based on tool response"""
         return []
 
-    def get_tool_call_result_for_loop_history(
-        self,
-        tool_response: ToolCallResponse,
-    ) -> LanguageModelMessage:
-        raise NotImplementedError("function is not supported")
-
     async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
         """Execute the MCP tool using SDK to call public API"""
         self.logger.info(f"Running MCP tool: {self.name}")

unique_toolkit/agentic/tools/openai_builtin/__init__.py

@@ -0,0 +1,11 @@
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter import (
+    OpenAICodeInterpreterConfig,
+    OpenAICodeInterpreterTool,
+)
+from unique_toolkit.agentic.tools.openai_builtin.manager import OpenAIBuiltInToolManager
+
+__all__ = [
+    "OpenAIBuiltInToolManager",
+    "OpenAICodeInterpreterTool",
+    "OpenAICodeInterpreterConfig",
+]

unique_toolkit/agentic/tools/openai_builtin/base.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+from enum import StrEnum
+from typing import Generic, TypeVar
+
+from openai.types.responses.tool_param import CodeInterpreter
+
+from unique_toolkit.agentic.tools.schemas import ToolPrompts
+
+
+class OpenAIBuiltInToolName(StrEnum):
+    CODE_INTERPRETER = "code_interpreter"
+
+
+BuiltInToolType = CodeInterpreter  # Add other tool types when needed
+ToolType = TypeVar("ToolType", bound=BuiltInToolType)
+
+
+class OpenAIBuiltInTool(ABC, Generic[ToolType]):
+    @property
+    @abstractmethod
+    def name(self) -> OpenAIBuiltInToolName:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def tool_description(self) -> BuiltInToolType:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def get_tool_prompts(self) -> ToolPrompts:
+        raise NotImplementedError()

unique_toolkit/agentic/tools/openai_builtin/code_interpreter/__init__.py

@@ -0,0 +1,8 @@
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.config import (
+    OpenAICodeInterpreterConfig,
+)
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.service import (
+    OpenAICodeInterpreterTool,
+)
+
+__all__ = ["OpenAICodeInterpreterConfig", "OpenAICodeInterpreterTool"]

unique_toolkit/agentic/tools/openai_builtin/code_interpreter/config.py

@@ -0,0 +1,57 @@
+from pydantic import Field
+
+from unique_toolkit.agentic.tools.factory import ToolFactory
+from unique_toolkit.agentic.tools.openai_builtin.base import (
+    OpenAIBuiltInToolName,
+)
+from unique_toolkit.agentic.tools.schemas import BaseToolConfig
+
+DEFAULT_TOOL_DESCRIPTION = "Use this tool to run python code, e.g to generate plots, process excel files, perform calculations, etc."
+
+DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT = """
+Use this tool to run python code, e.g to generate plots, process excel files, perform calculations, etc.
+Instructions:
+- All files uploaded to the chat are available in the code interpreter under the path `/mnt/data/<filename>
+- All files generated through code should be saved in the `/mnt/data` folder
+
+Instructions for displaying images and files in the chat:
+Once files are generated in the `/mnt/data` folder you MUST reference them in the chat using markdown syntax in order to display them in the chat.
+
+- If you want to display an image, use the following syntax: `![Image Name](sandbox:/mnt/data/<filename>)`
+    - Images will be converted and shown in the chat.
+    - Do NOT display an extra download link for images a part from the markdown above.
+    - Not using markdown syntax will FAIL to show images to the user. 
+    - YOU MUST use the syntax above to display images, otherwise the image will not be displayed in the chat.
+    - For displaying a link to a file, use the following syntax: `[filename](sandbox:/mnt/data/<filename>)`
+        - Files are converted to references the user can click on to download the file
+
+You MUST always use this syntax, otherwise the files will not be displayed in the chat.
+""".strip()
+
+DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT = ""
+
+DEFAULT_TOOL_FORMAT_INFORMATION_FOR_USER_PROMPT = ""
+
+DEFAULT_TOOL_DESCRIPTION_FOR_USER_PROMPT = ""
+
+
+class OpenAICodeInterpreterConfig(BaseToolConfig):
+    upload_files_in_chat: bool = Field(default=True)
+
+    tool_description: str = DEFAULT_TOOL_DESCRIPTION
+    tool_description_for_system_prompt: str = DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT
+    tool_format_information_for_system_prompt: str = (
+        DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT
+    )
+    tool_description_for_user_prompt: str = DEFAULT_TOOL_DESCRIPTION_FOR_USER_PROMPT
+    tool_format_information_for_user_prompt: str = (
+        DEFAULT_TOOL_FORMAT_INFORMATION_FOR_USER_PROMPT
+    )
+
+    expires_after_minutes: int = 20
+    use_auto_container: bool = False
+
+
+ToolFactory.register_tool_config(
+    OpenAIBuiltInToolName.CODE_INTERPRETER, OpenAICodeInterpreterConfig
+)