unique_toolkit 0.7.7__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/schemas.py

@@ -0,0 +1,141 @@
+import base64
+import gzip
+import re
+from typing import Any, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator
+
+from unique_toolkit._common.pydantic_helpers import get_configuration_dict
+from unique_toolkit.agentic.tools.utils.source_handling.schema import SourceFormatConfig
+from unique_toolkit.content.schemas import ContentChunk
+
+
+# TODO: this needs to be more general as the tools can potentially return anything maybe make a base class and then derive per "type" of tool
+class ToolCallResponse(BaseModel):
+    id: str
+    name: str
+    content: str = ""
+    debug_info: Optional[dict] = None  # TODO: Make the default {}
+    content_chunks: Optional[list[ContentChunk]] = None  # TODO: Make the default []
+    reasoning_result: Optional[dict] = None  # TODO: Make the default {}
+    error_message: str = ""
+
+    @property
+    def successful(self) -> bool:
+        return self.error_message == ""
+
+
+class BaseToolConfig(BaseModel):
+    model_config = get_configuration_dict()
+    # TODO: add a check for the parameters to all be consistent within the tool config
+    pass
+
+
+class Source(BaseModel):
+    """Represents the sources in the tool call response that the llm will see
+
+    Args:
+        source_number: The number of the source
+        content: The content of the source
+    """
+
+    model_config = ConfigDict(
+        validate_by_alias=True, serialize_by_alias=True, validate_by_name=True
+    )
+
+    source_number: int | None = Field(
+        default=None,
+        serialization_alias="[source_number] - Used for citations!",
+        validation_alias="[source_number] - Used for citations!",
+    )
+    content: str = Field(
+        serialization_alias="[content] - Content of source",
+        validation_alias="[content] - Content of source",
+    )
+    order: int = Field(
+        serialization_alias="[order] - Index in the document!",
+        validation_alias="[order] - Index in the document!",
+    )
+    chunk_id: str | None = Field(
+        default=None,
+        serialization_alias="[chunk_id] - IGNORE",
+        validation_alias="[chunk_id] - IGNORE",
+    )
+    id: str = Field(
+        serialization_alias="[id] - IGNORE",
+        validation_alias="[id] - IGNORE",
+    )
+    key: str | None = Field(
+        default=None,
+        serialization_alias="[key] - IGNORE",
+        validation_alias="[key] - IGNORE",
+    )
+    metadata: dict[str, str] | str | None = Field(
+        default=None,
+        serialization_alias="[metadata] - Formatted metadata",
+        validation_alias="[metadata] - Formatted metadata",
+    )
+    url: str | None = Field(
+        default=None,
+        serialization_alias="[url] - IGNORE",
+        validation_alias="[url] - IGNORE",
+    )
+
+    @field_validator("metadata", mode="before")
+    def _metadata_str_to_dict(
+        cls, v: str | dict[str, str] | None
+    ) -> dict[str, str] | None:
+        """
+        Accept   • dict   → keep as-is
+                 • str    → parse tag-string back to dict
+        """
+        if v is None or isinstance(v, dict):
+            return v
+
+        # v is the rendered string.  Build a dict by matching the
+        # patterns defined in SourceFormatConfig.sections.
+        cfg = SourceFormatConfig()  # or inject your app-wide config
+        out: dict[str, str] = {}
+        for key, tmpl in cfg.sections.items():
+            pattern = cfg.template_to_pattern(tmpl)
+            m = re.search(pattern, v, flags=re.S)
+            if m:
+                out[key] = m.group(1).strip()
+
+        return out if out else v  # type: ignore
+
+    # Compression + Base64 for url to hide it from the LLM
+    @field_serializer("url")
+    def serialize_url(self, value: str | None) -> str | None:
+        if value is None:
+            return None
+        # Compress then base64 encode
+        compressed = gzip.compress(value.encode())
+        return base64.b64encode(compressed).decode()
+
+    @field_validator("url", mode="before")
+    @classmethod
+    def validate_url(cls, value: Any) -> str | None:
+        if value is None or isinstance(value, str) and not value:
+            return None
+        if isinstance(value, str):
+            try:
+                # Try to decode base64 then decompress
+                decoded_bytes = base64.b64decode(value.encode())
+                decompressed = gzip.decompress(decoded_bytes).decode()
+                return decompressed
+            except Exception:
+                # If decoding/decompression fails, assume it's plain text
+                return value
+        return str(value)
+
+
+class ToolPrompts(BaseModel):
+    name: str
+    display_name: str
+    tool_system_prompt: str
+    tool_format_information_for_system_prompt: str
+    tool_user_prompt: str
+    tool_format_information_for_user_prompt: str
+    tool_description: str
+    input_model: dict[str, Any]