unique-toolkit 2026.30.0.dev15__tar.gz → 2026.30.0.dev17__tar.gz

{unique_toolkit-2026.30.0.dev15 → unique_toolkit-2026.30.0.dev17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 2026.30.0.dev15
+Version: 2026.30.0.dev17
 Summary: 
 Author: Cedric Klinkert, Martin Fadler, Sadique Sheik, Fabian Schläpfer, Pascal Hauri, Ejiroghene Oki-Peter, Rami Azouz, Ricardo Visini, Sanzio Monti
 Author-email: Cedric Klinkert <cedric.klinkert@unique.ch>, Martin Fadler <martin.fadler@unique.ch>, Sadique Sheik <sadique@unique.ch>, Fabian Schläpfer <fabian@unique.ch>, Pascal Hauri <pascal@unique.ch>, Ejiroghene Oki-Peter <oki-peter.ejiroghene@unique.ai>, Rami Azouz <rami.ext@unique.ch>, Ricardo Visini <ricardo.visini@unique.ch>, Sanzio Monti <sanzio.monti@unique.ai>
@@ -16,7 +16,7 @@ Requires-Dist: sseclient>=0.0.27,<0.0.28
 Requires-Dist: openai>=1.105.0,<3
 Requires-Dist: platformdirs>=4.0.0,<5
 Requires-Dist: pillow>=12.1.1
-Requires-Dist: unique-sdk>=2026.30.0.dev6,<2026.30.0rc0
+Requires-Dist: unique-sdk>=2026.30.0.dev8,<2026.30.0rc0
 Requires-Dist: jambo>=0.1.2,<0.2
 Requires-Dist: jinja2>=3.1.6,<4
 Requires-Dist: docxtpl>=0.20.1,<0.21

{unique_toolkit-2026.30.0.dev15 → unique_toolkit-2026.30.0.dev17}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "2026.30.0.dev15"
+version = "2026.30.0.dev17"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12"
@@ -28,7 +28,7 @@ dependencies = [
     "openai>=1.105.0,<3",
     "platformdirs>=4.0.0,<5",
     "pillow>=12.1.1",
-    "unique-sdk>=2026.30.0.dev6,<2026.30.0rc0",
+    "unique-sdk>=2026.30.0.dev8,<2026.30.0rc0",
     "jambo>=0.1.2,<0.2",
     "jinja2>=3.1.6,<4",
     "docxtpl>=0.20.1,<0.21",

{unique_toolkit-2026.30.0.dev15 → unique_toolkit-2026.30.0.dev17}/unique_toolkit/agentic/history_manager/loop_token_reducer.py

@@ -37,6 +37,37 @@ from unique_toolkit.services.factory import UniqueServiceFactory
 MAX_INPUT_TOKENS_SAFETY_PERCENTAGE = env_settings.input_correction_reduction_factor
 SAFETY_MARGIN_FOR_AGGRESSIVE_REDUCTION = env_settings.input_aggressive_reduction_factor
 
+PLAIN_TEXT_TRUNCATION_MARKER = "\n\n[... tool response too long and was truncated]"
+# TODO: Make this configurable
+MIN_PLAIN_TEXT_CHARS_TO_KEEP = 1000
+
+
+def _truncate_plain_text(text: str, overshoot_factor: float) -> str:
+    """Truncate plain text so it helps a tool response fit the token budget.
+
+    Keeps the head and appends a marker. The amount kept scales with the
+    overshoot (mirroring the source-reduction math) and never drops below
+    ``MIN_PLAIN_TEXT_CHARS_TO_KEEP``. Truncation is on characters rather than
+    tokens because the encoder is encode-only; on prose the character ratio is a
+    good proxy for the token ratio, and the outer loop re-counts and re-enters
+    until the payload fits or no further progress is possible. Any prior marker
+    is stripped first so repeated rounds neither stack markers nor re-flag an
+    already-floored string. Returns *text* unchanged when no truncation applies.
+    """
+    base = text.removesuffix(PLAIN_TEXT_TRUNCATION_MARKER)
+
+    # Copied as for the case where the tool returns chunks (below)
+    aggressive_divisor = overshoot_factor * SAFETY_MARGIN_FOR_AGGRESSIVE_REDUCTION
+    divisor = max(
+        aggressive_divisor if aggressive_divisor >= 1.2 else overshoot_factor, 1.0
+    )
+    chars_to_keep = max(MIN_PLAIN_TEXT_CHARS_TO_KEEP, int(len(base) // divisor))
+
+    if chars_to_keep >= len(base):
+        return text
+
+    return base[:chars_to_keep] + PLAIN_TEXT_TRUNCATION_MARKER
+
 
 class SourceReductionResult(BaseModel):
     message: LanguageModelToolMessage
@@ -118,7 +149,9 @@ class LoopTokenReducer:
         token_count = self._count_message_tokens(messages)
         self._log_token_usage(token_count)
 
-        while self._exceeds_token_limit(token_count):
+        while self._exceeds_token_limit(token_count) and self._can_reduce_history(
+            loop_history
+        ):
             token_count_before_reduction = token_count
             loop_history[:] = self._handle_token_limit_exceeded(
                 loop_history, token_count
@@ -141,18 +174,30 @@ class LoopTokenReducer:
         return messages
 
     def _exceeds_token_limit(self, token_count: int) -> bool:
-        """Check if token count exceeds the maximum allowed limit and if at least one tool call has more than one source."""
-        # At least one tool call should have more than one chunk as answer
+        """Whether the payload is over the effective input-token limit.
+
+        This is purely a size check. Whether anything can actually be reduced to
+        get back under the limit is a separate concern (:meth:`_can_reduce_history`).
+        """
+        # TODO: This is not fully correct at the moment as the token_count
+        # include system_prompt and user question already
+        return token_count > self._effective_token_limit
+
+    def _can_reduce_history(self, loop_history: list[LanguageModelMessage]) -> bool:
+        """Whether there is any tool response left to reduce.
+
+        Two independent reduction levers exist:
+          * a tool call with more than one chunk can shed sources, and
+          * a chunk-less plain-text tool response can be truncated.
+        When neither applies the reduction loop cannot make progress and stops.
+        """
         has_multiple_chunks_for_a_tool_call = any(
             len(chunks) > 1
             for chunks in self._reference_manager.get_chunks_of_all_tools()
         )
-        # TODO: This is not fully correct at the moment as the token_count
-        # include system_prompt and user question already
-        # TODO: There is a problem if we exceed but only have one chunk per tool call
-        exceeds_limit = token_count > self._effective_token_limit
-
-        return has_multiple_chunks_for_a_tool_call and exceeds_limit
+        return has_multiple_chunks_for_a_tool_call or any(
+            self._can_truncate_message(message) for message in loop_history
+        )
 
     def _count_message_tokens(self, messages: LanguageModelMessages) -> int:
         """Count tokens in messages using the configured encoding model."""
@@ -211,7 +256,7 @@ class LoopTokenReducer:
             f"(overshoot factor: {overshoot_factor:.2f}x). Reducing number of sources per tool call.",
         )
 
-        return self._reduce_message_length_by_reducing_sources_in_tool_response(
+        return self._reduce_message_length_by_reducing_tool_responses_content(
             loop_history, overshoot_factor
         )
 
@@ -451,7 +496,7 @@ class LoopTokenReducer:
 
         return limited_history_messages[idx:]
 
-    def _reduce_message_length_by_reducing_sources_in_tool_response(
+    def _reduce_message_length_by_reducing_tool_responses_content(
         self,
         history: list[LanguageModelMessage],
         overshoot_factor: float,
@@ -473,21 +518,41 @@ class LoopTokenReducer:
         source_offset = max(0, self._max_db_source_number + 1)
 
         for message in history:
-            if self._should_reduce_message(message):
-                result = self._reduce_sources_in_tool_message(
+            if not self._should_reduce_message(message):
+                result = message
+            elif self._has_sources(message):
+                source_result = self._reduce_sources_in_tool_message(
                     message,
                     chunk_offset,
                     source_offset,
                     overshoot_factor,
                 )
-                content_chunks_reduced.extend(result.reduced_chunks)
-                history_reduced.append(result.message)
-                chunk_offset = result.chunk_offset
-                source_offset = result.source_offset
+                content_chunks_reduced.extend(source_result.reduced_chunks)
+                chunk_offset = source_result.chunk_offset
+                source_offset = source_result.source_offset
+                result = source_result.message
+            elif self._can_truncate_message(message):
+                assert isinstance(message.content, str)
+                truncated = _truncate_plain_text(message.content, overshoot_factor)
+                if truncated != message.content:
+                    self._logger.warning(
+                        f"Truncated plain-text tool response '{message.name}' to "
+                        f"{len(truncated)} chars "
+                        f"(overshoot factor: {overshoot_factor:.2f}x)."
+                    )
+                result = LanguageModelToolMessage(
+                    content=truncated,
+                    tool_call_id=message.tool_call_id,
+                    name=message.name,
+                )
             else:
-                history_reduced.append(message)
+                result = message
+
+            history_reduced.append(result)
+
+        if len(content_chunks_reduced) > 0:
+            self._reference_manager.replace(chunks=content_chunks_reduced)
 
-        self._reference_manager.replace(chunks=content_chunks_reduced)
         return history_reduced
 
     def _should_reduce_message(
@@ -498,6 +563,19 @@ class LoopTokenReducer:
             message, LanguageModelToolMessage
         )
 
+    def _has_sources(self, message: LanguageModelToolMessage) -> bool:
+        return len(self._reference_manager.get_chunks_of_tool(message.tool_call_id)) > 0
+
+    def _can_truncate_message(self, message: LanguageModelMessage) -> bool:
+        if not self._should_reduce_message(message) or message.name == "TableSearch":
+            return False
+        if not isinstance(message.content, str):
+            return False
+        if self._has_sources(message):
+            return False
+        base = message.content.removesuffix(PLAIN_TEXT_TRUNCATION_MARKER)
+        return len(base) > MIN_PLAIN_TEXT_CHARS_TO_KEEP
+
     def _reduce_sources_in_tool_message(
         self,
         message: LanguageModelToolMessage,

{unique_toolkit-2026.30.0.dev15 → unique_toolkit-2026.30.0.dev17}/unique_toolkit/agentic/tools/openai_builtin/base.py

@@ -2,7 +2,6 @@ from abc import ABC, abstractmethod
 from enum import StrEnum
 from typing import Generic, TypeVar
 
-from openai.types.responses import ResponseIncludable
 from openai.types.responses.tool_param import CodeInterpreter
 
 from unique_toolkit.agentic.tools.schemas import ToolPrompts
@@ -52,12 +51,3 @@ class OpenAIBuiltInTool(ABC, Generic[ToolType]):
         Default: False.
         """
         return False
-
-    def get_required_include_params(self) -> list[ResponseIncludable]:
-        """Return Responses API `include` values required by this tool.
-
-        Subclasses override this when they need additional data attached to the
-        response (e.g. code interpreter execution logs).  The default is an
-        empty list, meaning no extra includes are needed.
-        """
-        return []

{unique_toolkit-2026.30.0.dev15 → unique_toolkit-2026.30.0.dev17}/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/__init__.py

@@ -1,5 +1,9 @@
-from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.config import (
-    OpenAICodeInterpreterConfig,
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.activator import (
+    CodeInterpreterActivatorConfig,
+    CodeInterpreterActivatorTool,
+)
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.builder import (
+    CodeInterpreterBuilder,
 )
 from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.postprocessors import (
     DisplayCodeInterpreterFilesPostProcessor,
@@ -7,11 +11,15 @@ from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.postprocessors
     ShowExecutedCodePostprocessor,
     ShowExecutedCodePostprocessorConfig,
 )
-from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.service import (
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.tool import (
+    OpenAICodeInterpreterConfig,
     OpenAICodeInterpreterTool,
 )
 
 __all__ = [
+    "CodeInterpreterActivatorConfig",
+    "CodeInterpreterActivatorTool",
+    "CodeInterpreterBuilder",
     "OpenAICodeInterpreterConfig",
     "OpenAICodeInterpreterTool",
     "ShowExecutedCodePostprocessor",

unique_toolkit-2026.30.0.dev17/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/activator/__init__.py

@@ -0,0 +1,11 @@
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.activator.config import (
+    CodeInterpreterActivatorConfig,
+)
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.activator.service import (
+    CodeInterpreterActivatorTool,
+)
+
+__all__ = [
+    "CodeInterpreterActivatorConfig",
+    "CodeInterpreterActivatorTool",
+]

unique_toolkit-2026.30.0.dev17/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/activator/config.py

@@ -0,0 +1,23 @@
+from pydantic import Field
+
+from unique_toolkit.agentic.tools.schemas import (
+    BaseToolConfig,
+)
+
+_DEFAULT_TOOL_DESCRIPTION = """
+Activate the code interpreter to run Python in a secure environment for
+calculations, data analysis, processing uploaded files (including spreadsheets),
+plotting and data visualization, and generating downloadable files. 
+
+CRITICAL: Act as though these capabilities are already available: whenever they are useful, call this tool
+silently and proceed. Never mention this tool, its activation, or its setup to the user.
+""".strip()
+
+_DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT = _DEFAULT_TOOL_DESCRIPTION
+
+
+class CodeInterpreterActivatorConfig(BaseToolConfig):
+    tool_description: str = Field(default=_DEFAULT_TOOL_DESCRIPTION)
+    tool_description_for_system_prompt: str = (
+        _DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT
+    )

unique_toolkit-2026.30.0.dev17/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/activator/service.py

@@ -0,0 +1,110 @@
+import asyncio
+import logging
+from typing import override
+
+from pydantic import BaseModel, Field
+
+from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.activator.config import (
+    CodeInterpreterActivatorConfig,
+)
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.builder import (
+    CodeInterpreterBuilder,
+)
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.tool import (
+    OpenAICodeInterpreterTool,
+)
+from unique_toolkit.agentic.tools.schemas import (
+    ToolCallResponse,
+)
+from unique_toolkit.agentic.tools.tool import Tool
+from unique_toolkit.language_model.schemas import (
+    LanguageModelFunction,
+    LanguageModelToolDescription,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class _CodeParams(BaseModel):
+    arg: bool = Field(True, description="Ignored Argument.")
+
+
+class CodeInterpreterActivatorTool(Tool[CodeInterpreterActivatorConfig]):
+    NAME = "ActivatePython"
+    DISPLAY_NAME = "Activate Python"
+
+    name = NAME
+
+    def __init__(
+        self,
+        config: CodeInterpreterActivatorConfig,
+        builder: CodeInterpreterBuilder,
+    ) -> None:
+        super().__init__(config=config)
+        self._builder = builder
+        self._built_tool: OpenAICodeInterpreterTool | None = None
+        self._activation_lock = asyncio.Lock()
+
+    @property
+    def is_activated(self) -> bool:
+        return self._built_tool is not None
+
+    def get_activated_tool(self) -> OpenAICodeInterpreterTool:
+        if self._built_tool is None:
+            raise RuntimeError(
+                "Code interpreter tool is not activated yet; call `run` first."
+            )
+        return self._built_tool
+
+    @override
+    def display_name(self) -> str:
+        return self.DISPLAY_NAME
+
+    @override
+    def is_enabled(self) -> bool:
+        return True
+
+    @override
+    def is_exclusive(self) -> bool:
+        return False
+
+    @override
+    def is_capability(self) -> bool:
+        return True
+
+    @override
+    def tool_description(self) -> LanguageModelToolDescription:
+        return LanguageModelToolDescription(
+            name=self.name,
+            description=self.config.tool_description,
+            parameters=_CodeParams,
+        )
+
+    @override
+    def tool_description_for_system_prompt(self) -> str:
+        return self.config.tool_description_for_system_prompt
+
+    @override
+    def evaluation_check_list(self) -> list[EvaluationMetricName]:
+        return []
+
+    @override
+    def get_evaluation_checks_based_on_tool_response(
+        self,
+        tool_response: ToolCallResponse,
+    ) -> list[EvaluationMetricName]:
+        return []
+
+    @override
+    async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
+        async with self._activation_lock:
+            if self._built_tool is None:
+                logger.info("Building code interpreter tool")
+                self._built_tool = await self._builder.build()
+
+        return ToolCallResponse(
+            id=tool_call.id,
+            name=self.name,
+            content="The python tool is now active. You may execute python code.",
+        )

unique_toolkit-2026.30.0.dev17/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/builder/__init__.py

@@ -0,0 +1,5 @@
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.builder.builder import (
+    CodeInterpreterBuilder,
+)
+
+__all__ = ["CodeInterpreterBuilder"]

unique_toolkit-2026.30.0.dev17/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/builder/_container.py

@@ -0,0 +1,50 @@
+import logging
+
+from openai import AsyncOpenAI
+
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.builder._memory import (
+    CodeExecutionShortTermMemorySchema,
+)
+
+logger = logging.getLogger(__name__)
+
+
+async def check_container_exists(
+    client: AsyncOpenAI,
+    memory: CodeExecutionShortTermMemorySchema,
+) -> bool:
+    try:
+        container = await client.containers.retrieve(memory.container_id)
+    # The error here is sometimes InternalServerError, and sometimes a NotFoundError. We catch everything and re-create on exception
+    except Exception:
+        logger.exception("Container %s not found", memory.container_id)
+        return False
+
+    if container.status not in ["active", "running"]:
+        logger.info(
+            "Container %s has status `%s`, recreating a new one",
+            memory.container_id,
+            container.status,
+        )
+        return False
+
+    logger.info("Container %s found in short term memory", memory.container_id)
+    return True
+
+
+async def create_container(
+    client: AsyncOpenAI,
+    chat_id: str,
+    user_id: str,
+    company_id: str,
+    expires_after_minutes: int,
+) -> str:
+    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,
+        },
+    )
+    logger.info("Created new container %s", container.id)
+    return container.id

unique_toolkit-2026.30.0.dev17/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/builder/_files.py

@@ -0,0 +1,152 @@
+import asyncio
+import logging
+
+from openai import AsyncOpenAI
+from tenacity import (
+    AsyncRetrying,
+    before_sleep_log,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from unique_toolkit import ContentService
+from unique_toolkit._common.execution import (
+    SafeTaskExecutor,
+)
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.builder._memory import (
+    CodeExecutionShortTermMemorySchema,
+)
+from unique_toolkit.content.schemas import (
+    Content,
+)
+
+logger = logging.getLogger(__name__)
+
+UPLOAD_MAX_RETRIES = 2
+UPLOAD_RETRY_BASE_DELAY = 0.5
+
+
+def build_upload_retry() -> AsyncRetrying:
+    """Exponential-backoff retry policy for transient upload/download failures.
+
+    Matches the pattern used in the ``DisplayCodeInterpreterFilesPostProcessor``
+    so that every outbound I/O call gets the same behaviour: up to
+    ``UPLOAD_MAX_RETRIES`` extra attempts, doubling the wait each time,
+    with a WARNING log before each sleep.
+    """
+    return AsyncRetrying(
+        stop=stop_after_attempt(1 + UPLOAD_MAX_RETRIES),
+        wait=wait_exponential(multiplier=UPLOAD_RETRY_BASE_DELAY),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+        reraise=True,
+    )
+
+
+def check_file_already_uploaded(
+    content_id: str,
+    memory: CodeExecutionShortTermMemorySchema,
+) -> bool:
+    if content_id not in memory.file_paths:
+        logger.info("File with id %s not in short term memory", content_id)
+        return False
+
+    return True
+
+
+async def upload_file_to_container(
+    client: AsyncOpenAI,
+    content_id: str,
+    filename: str,
+    content_service: ContentService,
+    container_id: str,
+) -> str:
+    logger.info(
+        "Uploading file %s (%s) to container %s",
+        content_id,
+        filename,
+        container_id,
+    )
+
+    file_content = await build_upload_retry()(
+        content_service.download_content_to_bytes_async,
+        content_id=content_id,
+    )
+    logger.info(
+        "Downloaded %d bytes for file %s; uploading to container %s",
+        len(file_content),
+        content_id,
+        container_id,
+    )
+
+    openai_file = await build_upload_retry()(
+        client.containers.files.create,
+        container_id=container_id,
+        file=(filename, file_content),
+    )
+    logger.info(
+        "File %s successfully uploaded as OpenAI file %s in container %s",
+        content_id,
+        openai_file.id,
+        container_id,
+    )
+
+    return openai_file.path
+
+
+async def upload_files_to_container(
+    client: AsyncOpenAI,
+    uploaded_files: list[Content],
+    memory: CodeExecutionShortTermMemorySchema,
+    content_service: ContentService,
+) -> tuple[CodeExecutionShortTermMemorySchema, bool]:
+    async def _check_and_upload(content: Content) -> str | None:
+        if check_file_already_uploaded(content_id=content.id, memory=memory):
+            return None
+
+        return await upload_file_to_container(
+            client=client,
+            content_id=content.id,
+            filename=content.key,
+            content_service=content_service,
+            container_id=memory.container_id,
+        )
+
+    # Deduplicate
+    unique_contents = {content.id: content for content in uploaded_files}.values()
+
+    executor = SafeTaskExecutor(logger=logger)
+
+    results = await asyncio.gather(
+        *(
+            executor.execute_async(_check_and_upload, content)
+            for content in unique_contents
+        ),
+    )
+
+    updated = False
+    for content, result in zip(unique_contents, results):
+        if result.success and (filepath := result.unpack()) is not None:
+            memory.file_paths[content.id] = filepath
+            updated = True
+
+    return memory, updated
+
+
+async def resolve_kb_contents(
+    content_service: ContentService,
+    content_ids: list[str],
+) -> list[Content]:
+    contents = await content_service.search_contents_async(
+        where={"id": {"in": content_ids}},
+    )
+
+    found_ids = {c.id for c in contents}
+    missing = [content_id for content_id in content_ids if content_id not in found_ids]
+    if missing:
+        logger.warning(
+            "additional_uploaded_documents: %d content ids not found or not accessible in KB: %s",
+            len(missing),
+            missing,
+        )
+
+    return contents

unique_toolkit-2026.30.0.dev17/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/builder/_memory.py

@@ -0,0 +1,35 @@
+from openai import BaseModel
+
+from unique_toolkit import ShortTermMemoryService
+from unique_toolkit.agentic.short_term_memory_manager.persistent_short_term_memory_manager import (
+    PersistentShortMemoryManager,
+)
+
+SHORT_TERM_MEMORY_NAME = "container_code_execution"
+
+
+class CodeExecutionShortTermMemorySchema(BaseModel):
+    container_id: str
+    file_paths: dict[str, str] = {}
+
+
+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