unique_toolkit 1.15.0__py3-none-any.whl → 1.16.1__py3-none-any.whl

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
+)

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