unique_toolkit 0.7.9__py3-none-any.whl → 1.33.3__py3-none-any.whl

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

@@ -0,0 +1,88 @@
+from pydantic import Field
+from pydantic.json_schema import SkipJsonSchema
+
+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
+
+CRUCIAL 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.
+
+WHENEVER you reference a generated file, you MUST use the following format:
+```
+**Descriptive Title of Graph/Chart/File** (<- linebreak is important) (You must choose a good user friendly title, Other markdown syntax such as `#` can be used too)
+[*Generating your {Graph/Chart/File}…*](sandbox:/mnt/data/<filename>)
+```
+IMPORTANT: Do NOT append a leading `!` even when displaying an image.
+Always use a line break between the title and the markdown!
+
+- Files with image file extensions are displayed directly in the chat, while other file extensions are shown as download links.
+- Not using syntax above 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.
+- Only the following file types are allowed to be uploaded to the platform, anything else will FAIL: PDF, DOCX, XLSX, PPTX, CSV, HTML, MD, TXT, PNG, JPG, JPEG.
+
+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_to_container: bool = Field(
+        default=True,
+        description="If set, the files uploaded to the chat will be uploaded to the container where code is executed.",
+    )
+
+    tool_description: str = Field(
+        default=DEFAULT_TOOL_DESCRIPTION,
+        description="The description of the tool that will be sent to the model.",
+    )
+    tool_description_for_system_prompt: str = Field(
+        default=DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT,
+        description="The description of the tool that will be included in the system prompt.",
+    )
+    tool_format_information_for_system_prompt: SkipJsonSchema[str] = Field(
+        # Since the tool is executed in Azure, it's not always possible to have dynamic format information
+        default=DEFAULT_TOOL_FORMAT_INFORMATION_FOR_SYSTEM_PROMPT,
+        description="The format information of the tool that will be included in the system prompt.",
+    )
+    tool_description_for_user_prompt: SkipJsonSchema[str] = (
+        Field(  # At the moment, this is not appended to the user prompt
+            default=DEFAULT_TOOL_DESCRIPTION_FOR_USER_PROMPT,
+            description="The description of the tool that will be included in the user prompt.",
+        )
+    )
+    tool_format_information_for_user_prompt: SkipJsonSchema[str] = (
+        Field(  # At the moment, this is not appended to the user prompt
+            default=DEFAULT_TOOL_FORMAT_INFORMATION_FOR_USER_PROMPT,
+            description="The format information of the tool that will be included in the user prompt.",
+        )
+    )
+
+    expires_after_minutes: int = Field(
+        default=20,
+        description="The number of minutes after which the container will be deleted.",
+    )
+    use_auto_container: bool = Field(
+        default=False,
+        description="If set, use the `auto` container setting from OpenAI. Note that this will recreate the container on each call.",
+    )
+
+
+ToolFactory.register_tool_config(
+    OpenAIBuiltInToolName.CODE_INTERPRETER, OpenAICodeInterpreterConfig
+)

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

@@ -0,0 +1,250 @@
+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,
+        is_exclusive: bool = False,
+    ) -> 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
+        self._is_exclusive = is_exclusive
+
+    @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",
+        }
+
+    @override
+    def is_enabled(self) -> bool:
+        return True
+
+    @override
+    def takes_control(self) -> bool:
+        return False
+
+    @override
+    def is_exclusive(self) -> bool:
+        return self._is_exclusive
+
+    @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,
+        is_exclusive: bool = False,
+    ) -> "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,
+        )
+
+        if config.upload_files_in_chat_to_container:
+            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, is_exclusive=is_exclusive
+        )
+
+    @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={},
+        )
+
+    @override
+    def display_name(self) -> str:
+        return self.DISPLAY_NAME

unique_toolkit/agentic/tools/openai_builtin/manager.py

@@ -0,0 +1,79 @@
+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,
+        builtin_tools: list[OpenAIBuiltInTool],
+    ):
+        self._builtin_tools = builtin_tools
+
+    @classmethod
+    async def _build_tool(
+        cls,
+        uploaded_files: list[Content],
+        content_service: ContentService,
+        user_id: str,
+        company_id: str,
+        chat_id: str,
+        client: AsyncOpenAI,
+        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=uploaded_files,
+                content_service=content_service,
+                client=client,
+                company_id=company_id,
+                user_id=user_id,
+                chat_id=chat_id,
+                is_exclusive=tool_config.is_exclusive,
+            )
+            return tool
+        else:
+            raise ValueError(f"Unknown built-in tool name: {tool_config.name}")
+
+    @classmethod
+    async def build_manager(
+        cls,
+        uploaded_files: list[Content],
+        content_service: ContentService,
+        user_id: str,
+        company_id: str,
+        chat_id: str,
+        client: AsyncOpenAI,
+        tool_configs: list[ToolBuildConfig],
+    ) -> "OpenAIBuiltInToolManager":
+        builtin_tools = []
+        for tool_config in tool_configs:
+            if tool_config.name in OpenAIBuiltInToolName and tool_config.is_enabled:
+                builtin_tools.append(
+                    await cls._build_tool(
+                        uploaded_files,
+                        content_service,
+                        user_id,
+                        company_id,
+                        chat_id,
+                        client,
+                        tool_config,
+                    )
+                )
+
+        return OpenAIBuiltInToolManager(builtin_tools)
+
+    def get_all_openai_builtin_tools(self) -> list[OpenAIBuiltInTool]:
+        return self._builtin_tools.copy()

unique_toolkit/agentic/tools/schemas.py

@@ -0,0 +1,145 @@
+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 = ""
+    system_reminder: str = Field(
+        default="",
+        description="A reminder for the agent to consider when using the tool that will be appended to the tool call response",
+    )
+
+    @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]