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

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

@@ -0,0 +1,230 @@
+import logging
+from typing import override
+
+from openai import AsyncOpenAI, BaseModel, NotFoundError
+from openai.types.responses.tool_param import CodeInterpreter
+
+from unique_toolkit import ContentService, ShortTermMemoryService
+from unique_toolkit.agentic.short_term_memory_manager.persistent_short_term_memory_manager import (
+    PersistentShortMemoryManager,
+)
+from unique_toolkit.agentic.tools.openai_builtin.base import (
+    OpenAIBuiltInTool,
+    OpenAIBuiltInToolName,
+)
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.config import (
+    OpenAICodeInterpreterConfig,
+)
+from unique_toolkit.agentic.tools.schemas import ToolPrompts
+from unique_toolkit.content.schemas import (
+    Content,
+)
+
+logger = logging.getLogger(__name__)
+
+
+_SHORT_TERM_MEMORY_NAME = "container_code_execution"
+
+
+class CodeExecutionShortTermMemorySchema(BaseModel):
+    container_id: str | None = None
+    file_ids: dict[str, str] = {}  # Mapping of unique file id to openai file id
+
+
+CodeExecutionMemoryManager = PersistentShortMemoryManager[
+    CodeExecutionShortTermMemorySchema
+]
+
+
+def _get_container_code_execution_short_term_memory_manager(
+    company_id: str, user_id: str, chat_id: str
+) -> CodeExecutionMemoryManager:
+    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=CodeExecutionShortTermMemorySchema,
+        short_term_memory_name=_SHORT_TERM_MEMORY_NAME,
+    )
+    return short_term_memory_manager
+
+
+async def _create_container_if_not_exists(
+    client: AsyncOpenAI,
+    chat_id: str,
+    user_id: str,
+    company_id: str,
+    expires_after_minutes: int,
+    memory: CodeExecutionShortTermMemorySchema | None = None,
+) -> CodeExecutionShortTermMemorySchema:
+    if memory is not None:
+        logger.info("Container found in short term memory")
+    else:
+        logger.info("No Container in short term memory, creating a new container")
+        memory = CodeExecutionShortTermMemorySchema()
+
+    container_id = memory.container_id
+
+    if container_id is not None:
+        try:
+            container = await client.containers.retrieve(container_id)
+            if container.status not in ["active", "running"]:
+                logger.info(
+                    "Container has status `%s`, recreating a new one", container.status
+                )
+                container_id = None
+        except NotFoundError:
+            container_id = None
+
+    if container_id is None:
+        memory = CodeExecutionShortTermMemorySchema()
+
+        container = await client.containers.create(
+            name=f"code_execution_{company_id}_{user_id}_{chat_id}",
+            expires_after={
+                "anchor": "last_active_at",
+                "minutes": expires_after_minutes,
+            },
+        )
+
+        memory.container_id = container.id
+
+    return memory
+
+
+async def _upload_files_to_container(
+    client: AsyncOpenAI,
+    uploaded_files: list[Content],
+    memory: CodeExecutionShortTermMemorySchema,
+    content_service: ContentService,
+    chat_id: str,
+) -> CodeExecutionShortTermMemorySchema:
+    container_id = memory.container_id
+
+    assert container_id is not None
+
+    memory = memory.model_copy(deep=True)
+
+    for file in uploaded_files:
+        upload = True
+        if file.id in memory.file_ids:
+            try:
+                _ = await client.containers.files.retrieve(
+                    container_id=container_id, file_id=memory.file_ids[file.id]
+                )
+                logger.info("File with id %s already uploaded to container", file.id)
+                upload = False
+            except NotFoundError:
+                upload = True
+
+        if upload:
+            logger.info(
+                "Uploding file %s to container %s", file.id, memory.container_id
+            )
+            file_content = content_service.download_content_to_bytes(
+                content_id=file.id, chat_id=chat_id
+            )  # TODO: Use async version when available
+
+            openai_file = await client.containers.files.create(
+                container_id=container_id,
+                file=(file.key, file_content),
+            )
+            memory.file_ids[file.id] = openai_file.id
+
+    return memory
+
+
+class OpenAICodeInterpreterTool(OpenAIBuiltInTool[CodeInterpreter]):
+    DISPLAY_NAME = "Code Interpreter"
+
+    def __init__(
+        self,
+        config: OpenAICodeInterpreterConfig,
+        container_id: str | None,
+    ):
+        self._config = config
+
+        if not config.use_auto_container and container_id is None:
+            raise ValueError("`container_id` required when not using `auto` containers")
+
+        self._container_id = container_id
+
+    @property
+    @override
+    def name(self) -> OpenAIBuiltInToolName:
+        return OpenAIBuiltInToolName.CODE_INTERPRETER
+
+    @override
+    def tool_description(self) -> CodeInterpreter:
+        if self._config.use_auto_container:
+            return {"container": {"type": "auto"}, "type": "code_interpreter"}
+
+        return {
+            "container": self._container_id,  # type: ignore
+            "type": "code_interpreter",
+        }
+
+    @classmethod
+    async def build_tool(
+        cls,
+        config: OpenAICodeInterpreterConfig,
+        uploaded_files: list[Content],
+        client: AsyncOpenAI,
+        content_service: ContentService,
+        company_id: str,
+        user_id: str,
+        chat_id: str,
+    ) -> "OpenAICodeInterpreterTool":
+        if config.use_auto_container:
+            logger.info("Using `auto` container setting")
+            return cls(config=config, container_id=None)
+
+        memory_manager = _get_container_code_execution_short_term_memory_manager(
+            company_id=company_id,
+            user_id=user_id,
+            chat_id=chat_id,
+        )
+
+        memory = await memory_manager.load_async()
+
+        memory = await _create_container_if_not_exists(
+            client=client,
+            memory=memory,
+            chat_id=chat_id,
+            user_id=user_id,
+            company_id=company_id,
+            expires_after_minutes=config.expires_after_minutes,
+        )
+
+        memory = await _upload_files_to_container(
+            client=client,
+            uploaded_files=uploaded_files,
+            content_service=content_service,
+            chat_id=chat_id,
+            memory=memory,
+        )
+
+        await memory_manager.save_async(memory)
+
+        assert memory.container_id is not None
+
+        return OpenAICodeInterpreterTool(
+            config=config, container_id=memory.container_id
+        )
+
+    @override
+    def get_tool_prompts(self) -> ToolPrompts:
+        return ToolPrompts(
+            name="the python tool",  # https://platform.openai.com/docs/guides/tools-code-interpreter
+            display_name=self.DISPLAY_NAME,
+            tool_description=self._config.tool_description,
+            tool_system_prompt=self._config.tool_description_for_system_prompt,
+            tool_format_information_for_system_prompt=self._config.tool_format_information_for_system_prompt,
+            tool_user_prompt=self._config.tool_description_for_user_prompt,
+            tool_format_information_for_user_prompt=self._config.tool_format_information_for_user_prompt,
+            input_model={},
+        )

unique_toolkit/agentic/tools/openai_builtin/manager.py

@@ -0,0 +1,62 @@
+from openai import AsyncOpenAI
+
+from unique_toolkit.agentic.tools.config import ToolBuildConfig
+from unique_toolkit.agentic.tools.openai_builtin.base import (
+    OpenAIBuiltInTool,
+    OpenAIBuiltInToolName,
+)
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter import (
+    OpenAICodeInterpreterConfig,
+    OpenAICodeInterpreterTool,
+)
+from unique_toolkit.content.schemas import Content
+from unique_toolkit.content.service import ContentService
+
+
+class OpenAIBuiltInToolManager:
+    def __init__(
+        self,
+        uploaded_files: list[Content],
+        content_service: ContentService,
+        user_id: str,
+        company_id: str,
+        chat_id: str,
+        client: AsyncOpenAI,
+    ):
+        self._uploaded_files = uploaded_files
+        self._content_service = content_service
+        self._user_id = user_id
+        self._company_id = company_id
+        self._client = client
+        self._chat_id = chat_id
+
+    async def _build_tool(self, tool_config: ToolBuildConfig) -> OpenAIBuiltInTool:
+        if tool_config.name == OpenAIBuiltInToolName.CODE_INTERPRETER:
+            assert isinstance(tool_config.configuration, OpenAICodeInterpreterConfig)
+            tool = await OpenAICodeInterpreterTool.build_tool(
+                config=tool_config.configuration,
+                uploaded_files=self._uploaded_files,
+                user_id=self._user_id,
+                company_id=self._company_id,
+                chat_id=self._chat_id,
+                content_service=self._content_service,
+                client=self._client,
+            )
+            return tool
+        else:
+            raise ValueError(f"Unknown built-in tool name: {tool_config.name}")
+
+    async def get_all_openai_builtin_tools(
+        self, tool_configs: list[ToolBuildConfig]
+    ) -> tuple[list[ToolBuildConfig], list[OpenAIBuiltInTool]]:
+        openai_builtin_tools = []
+        filtered_tool_configs = []
+
+        for tool_config in tool_configs:
+            if tool_config.name not in OpenAIBuiltInToolName:
+                filtered_tool_configs.append(tool_config)
+                continue
+
+            openai_builtin_tools.append(await self._build_tool(tool_config))
+
+        return filtered_tool_configs, openai_builtin_tools

unique_toolkit/agentic/tools/test/test_mcp_manager.py

@@ -5,6 +5,7 @@ import pytest
 from pydantic import BaseModel
 
 from tests.test_obj_factory import get_event_obj
+from unique_toolkit.agentic.tools.a2a import SubAgentResponseWatcher
 from unique_toolkit.agentic.tools.a2a.manager import A2AManager
 from unique_toolkit.agentic.tools.config import (
     ToolBuildConfig,
@@ -13,6 +14,8 @@ from unique_toolkit.agentic.tools.config import (
 )
 from unique_toolkit.agentic.tools.factory import ToolFactory
 from unique_toolkit.agentic.tools.mcp.manager import MCPManager
+from unique_toolkit.agentic.tools.mcp.models import MCPToolConfig
+from unique_toolkit.agentic.tools.mcp.tool_wrapper import MCPToolWrapper
 from unique_toolkit.agentic.tools.schemas import BaseToolConfig
 from unique_toolkit.agentic.tools.tool import Tool
 from unique_toolkit.agentic.tools.tool_manager import ToolManager, ToolManagerConfig
@@ -48,11 +51,6 @@ class MockInternalSearchTool(Tool[BaseToolConfig]):
     def tool_format_information_for_system_prompt(self) -> str:
         return "Use this tool to search for content"
 
-    def get_tool_call_result_for_loop_history(self, tool_response):
-        from unique_toolkit.language_model.schemas import LanguageModelMessage
-
-        return LanguageModelMessage(role="tool", content="Mock search result")
-
     def evaluation_check_list(self):
         return []
 
@@ -114,11 +112,22 @@ class TestMCPManager:
             name="mcp_test_tool",
             description="Test MCP tool",
             input_schema={
+                "$schema": "https://json-schema.org/draft/2020-12/schema",
                 "type": "object",
                 "properties": {
-                    "query": {"type": "string", "description": "Search query"}
+                    "folder": {
+                        "description": "Mail folder to read mails from",
+                        "default": "inbox",
+                        "type": "string",
+                    },
+                    "limit": {
+                        "description": "Number of emails to retrieve",
+                        "default": 10,
+                        "type": "number",
+                        "minimum": 1,
+                        "maximum": 50,
+                    },
                 },
-                "required": ["query"],
             },
             output_schema=None,
             annotations=None,
@@ -173,6 +182,7 @@ class TestMCPManager:
         return A2AManager(
             logger=self.logger,
             tool_progress_reporter=tool_progress_reporter,
+            response_watcher=SubAgentResponseWatcher(),
         )
 
     @pytest.fixture
@@ -277,6 +287,7 @@ class TestMCPManager:
         a2a_manager = A2AManager(
             logger=self.logger,
             tool_progress_reporter=tool_progress_reporter,
+            response_watcher=SubAgentResponseWatcher(),
         )
 
         tool_manager = ToolManager(
@@ -313,6 +324,7 @@ class TestMCPManager:
         a2a_manager = A2AManager(
             logger=self.logger,
             tool_progress_reporter=tool_progress_reporter,
+            response_watcher=SubAgentResponseWatcher(),
         )
 
         tool_manager = ToolManager(
@@ -349,6 +361,7 @@ class TestMCPManager:
         a2a_manager = A2AManager(
             logger=self.logger,
             tool_progress_reporter=tool_progress_reporter,
+            response_watcher=SubAgentResponseWatcher(),
         )
 
         tool_manager = ToolManager(
@@ -388,6 +401,7 @@ class TestMCPManager:
         a2a_manager = A2AManager(
             logger=self.logger,
             tool_progress_reporter=tool_progress_reporter,
+            response_watcher=SubAgentResponseWatcher(),
         )
 
         tool_manager = ToolManager(
@@ -429,6 +443,7 @@ class TestMCPManager:
         a2a_manager = A2AManager(
             logger=self.logger,
             tool_progress_reporter=tool_progress_reporter,
+            response_watcher=SubAgentResponseWatcher(),
         )
 
         tool_manager = ToolManager(
@@ -446,3 +461,76 @@ class TestMCPManager:
         assert "internal_search" not in tool_names
         assert "mcp_test_tool" in tool_names
         assert len(tools) == 1
+
+    def test_mcp_tool_wrapper_generates_correct_parameters(
+        self, mcp_tools, mcp_servers, tool_progress_reporter
+    ):
+        """Test that MCPToolWrapper correctly generates parameters field from MCP tool input schema"""
+        # Get the first MCP tool and server from fixtures
+        mcp_tool = mcp_tools[0]
+        mcp_server = mcp_servers[0]
+
+        # Create MCPToolConfig
+        config = MCPToolConfig(
+            server_id="test_server_id",
+            server_name="test_server",
+            mcp_source_id="test_server_id",
+        )
+
+        # Initialize the MCPToolWrapper
+        tool_wrapper = MCPToolWrapper(
+            mcp_server=mcp_server,
+            mcp_tool=mcp_tool,
+            config=config,
+            event=self.event,
+            tool_progress_reporter=tool_progress_reporter,
+        )
+
+        # Call tool_description to get the LanguageModelToolDescription
+        tool_description = tool_wrapper.tool_description()
+
+        # Verify the basic properties
+        assert tool_description.name == "mcp_test_tool"
+        assert tool_description.description == "Test MCP tool"
+
+        # Verify that parameters field is correctly set to the input_schema
+        expected_parameters = {
+            "$schema": "https://json-schema.org/draft/2020-12/schema",
+            "type": "object",
+            "properties": {
+                "folder": {
+                    "description": "Mail folder to read mails from",
+                    "default": "inbox",
+                    "type": "string",
+                },
+                "limit": {
+                    "description": "Number of emails to retrieve",
+                    "default": 10,
+                    "type": "number",
+                    "minimum": 1,
+                    "maximum": 50,
+                },
+            },
+        }
+
+        assert tool_description.parameters == expected_parameters
+
+        # Verify specific parameter properties
+        parameters = tool_description.parameters
+        assert isinstance(parameters, dict)
+        assert parameters["type"] == "object"
+        assert "properties" in parameters
+
+        # Check folder parameter
+        folder_prop = parameters["properties"]["folder"]
+        assert folder_prop["type"] == "string"
+        assert folder_prop["description"] == "Mail folder to read mails from"
+        assert folder_prop["default"] == "inbox"
+
+        # Check limit parameter
+        limit_prop = parameters["properties"]["limit"]
+        assert limit_prop["type"] == "number"
+        assert limit_prop["description"] == "Number of emails to retrieve"
+        assert limit_prop["default"] == 10
+        assert limit_prop["minimum"] == 1
+        assert limit_prop["maximum"] == 50