unique-toolkit 2026.22.0.dev0__tar.gz → 2026.22.0.dev2__tar.gz

{unique_toolkit-2026.22.0.dev0 → unique_toolkit-2026.22.0.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 2026.22.0.dev0
+Version: 2026.22.0.dev2
 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.22.0.dev0,<2026.22.0rc0
+Requires-Dist: unique-sdk>=2026.22.0.dev2,<2026.22.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.22.0.dev0 → unique_toolkit-2026.22.0.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "2026.22.0.dev0"
+version = "2026.22.0.dev2"
 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.22.0.dev0,<2026.22.0rc0",
+    "unique-sdk>=2026.22.0.dev2,<2026.22.0rc0",
     "jambo>=0.1.2,<0.2",
     "jinja2>=3.1.6,<4",
     "docxtpl>=0.20.1,<0.21",

{unique_toolkit-2026.22.0.dev0 → unique_toolkit-2026.22.0.dev2}/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/config.py

@@ -17,7 +17,36 @@ from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.postprocessors
 )
 from unique_toolkit.agentic.tools.schemas import BaseToolConfig
 
-DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT = """
+_TOOL_DESCRIPTION_FILES_SUFFIX = """
+{%- if user_uploaded_files and user_uploaded_files | length > 0 %}
+
+## User Uploaded Files
+The following files have been uploaded by the user and are available in the environment:
+{% for file in user_uploaded_files %}
+- {{ file }}
+{% endfor %}
+{% endif %}
+
+{%- if kb_uploaded_files and kb_uploaded_files | length > 0 %}
+
+## Knowledge Base Files
+These files are made available by the system but were NOT uploaded by the user:
+{% for file in kb_uploaded_files %}
+- {{ file }}
+{% endfor %}
+{% endif %}
+
+{%- if code_interpreter_artifacts and code_interpreter_artifacts | length > 0 %}
+
+## Code Interpreter Artifacts
+The following files have been generated by you using the python tool in previous turns:
+{% for file in code_interpreter_artifacts %}
+- {{ file }}
+{% endfor %}
+{% endif %}
+""".strip()
+
+_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT_BASE = """
 Use this tool to run python code, e.g to generate plots, process excel files, perform calculations, etc.
 
 Instructions:
@@ -37,7 +66,7 @@ WHENEVER you reference a generated file, you MUST use the following format:
 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. 
+- 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.
@@ -58,7 +87,11 @@ Handling User Queries:
         - REMEMBER that you can always read the content of text, csv files if needed. In this case, you MUST always limit the amount of data displayed.
 - If a tool step fails, you must recover as much as possible.
 - After exhausting all possible solutions without success, inform the user of what was tried and request clarification/help.
-""".strip()
+"""
+
+DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT = (
+    _TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT_BASE + _TOOL_DESCRIPTION_FILES_SUFFIX
+).strip()
 
 # Used when the code-execution fence feature flag (UN-17972) is enabled.
 # The frontend derives the artifact title from the filename itself, so the
@@ -71,7 +104,7 @@ Handling User Queries:
 # UN-19880: kaleido is not installed in the Azure OpenAI code interpreter environment,
 # so plotly's fig.write_image() always fails. The "Visualization library choice" section
 # steers the model to use matplotlib for PNG outputs and plotly only for HTML artifacts.
-DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT_FENCE = """
+_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT_FENCE_BASE = """
 Use this tool to run python code, e.g to generate plots, process excel files, perform calculations, etc.
 
 Instructions:
@@ -137,7 +170,11 @@ Handling User Queries:
         - REMEMBER that you can always read the content of text, csv files if needed. In this case, you MUST always limit the amount of data displayed.
 - If a tool step fails, you must recover as much as possible.
 - After exhausting all possible solutions without success, inform the user of what was tried and request clarification/help.
-""".strip()
+"""
+
+DEFAULT_TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT_FENCE = (
+    _TOOL_DESCRIPTION_FOR_SYSTEM_PROMPT_FENCE_BASE + _TOOL_DESCRIPTION_FILES_SUFFIX
+).strip()
 
 
 DEFAULT_TOOL_DESCRIPTION_FOR_USER_PROMPT = ""

unique_toolkit-2026.22.0.dev2/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/postprocessors/artifacts.py

@@ -0,0 +1,103 @@
+import logging
+from mimetypes import guess_type
+
+from openai.types.responses.response_output_text import AnnotationContainerFileCitation
+from pydantic import BaseModel
+
+from unique_toolkit import ChatService
+from unique_toolkit._common.utils.files import FileMimeType, ImageMimeType
+from unique_toolkit.content.schemas import Content
+
+_CODE_EXECUTION_ARTIFACT_METADATA_KEY: str = "codeExecutionArtifactMetadata"
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class CodeExecutionArtifactMetadata(BaseModel):
+    container_id: str
+    file_id: str
+    filepath: str
+
+
+def _kb_safe_mime(mime: str) -> str:
+    """Return a MIME type the Unique KB will accept.
+
+    Membership is defined by ``FileMimeType`` and ``ImageMimeType`` in
+    ``unique_toolkit._common.utils.files`` (same catalog as path-based helpers
+    like ``FileMimeType.is_valid_mime``, but here we already have a resolved
+    MIME string from ``mimetypes.guess_type``, so we use StrEnum value lookup
+    instead of re-parsing a path).
+
+    The KB GraphQL API rejects many code-file MIME types (e.g. ``text/x-python``
+    for ``.py`` files).  Anything not in those enums is coerced to
+    ``text/plain`` so the file can be stored and downloaded without changing
+    its bytes.
+
+    Other ``image/*`` subtypes (e.g. ``image/jpg``) still pass through unchanged.
+    """
+    try:
+        FileMimeType(mime)
+        return mime
+    except ValueError:
+        pass
+    try:
+        ImageMimeType(mime)
+        return mime
+    except ValueError:
+        pass
+    if mime.startswith("image/"):
+        return mime
+    return "text/plain"
+
+
+async def save_code_execution_artifact(
+    chat_service: ChatService,
+    file: AnnotationContainerFileCitation,
+    file_bytes: bytes,
+) -> Content:
+    raw_mime = guess_type(file.filename)[0] or "text/plain"
+    mime = _kb_safe_mime(raw_mime)
+
+    if mime != raw_mime:
+        _LOGGER.info(
+            "MIME type '%s' is not supported by the Unique KB; "
+            "uploading '%s' as 'text/plain' so the file can be stored and downloaded.",
+            raw_mime,
+            file.filename,
+        )
+
+    _LOGGER.info(
+        "Uploading '%s' to knowledge base (%d bytes, mime type %s)",
+        file.filename,
+        len(file_bytes),
+        mime,
+    )
+
+    return await chat_service.upload_to_chat_from_bytes_async(
+        content=file_bytes,
+        content_name=file.filename,
+        mime_type=mime,
+        skip_ingestion=True,
+        hide_in_chat=True,
+        metadata={
+            _CODE_EXECUTION_ARTIFACT_METADATA_KEY: CodeExecutionArtifactMetadata(
+                container_id=file.container_id,
+                file_id=file.file_id,
+                filepath=f"/mnt/data/{file.filename}",
+            ).model_dump()
+        },
+    )
+
+
+def load_code_execution_metadata(
+    content: Content,
+) -> CodeExecutionArtifactMetadata | None:
+    if (
+        content.metadata is None
+        or _CODE_EXECUTION_ARTIFACT_METADATA_KEY not in content.metadata
+    ):
+        return None
+
+    return CodeExecutionArtifactMetadata.model_validate(
+        content.metadata[_CODE_EXECUTION_ARTIFACT_METADATA_KEY]
+    )

{unique_toolkit-2026.22.0.dev0 → unique_toolkit-2026.22.0.dev2}/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/postprocessors/generated_files.py

@@ -21,7 +21,6 @@ from tenacity import (
 
 from unique_toolkit import ChatService
 from unique_toolkit._common.execution import failsafe_async
-from unique_toolkit._common.utils.files import FileMimeType, ImageMimeType
 from unique_toolkit.agentic.feature_flags.feature_flags import feature_flags
 from unique_toolkit.agentic.postprocessor.postprocessor_manager import (
     ResponsesApiPostprocessor,
@@ -30,6 +29,9 @@ from unique_toolkit.agentic.short_term_memory_manager.persistent_short_term_memo
     PersistentShortMemoryManager,
 )
 from unique_toolkit.agentic.tools.config import get_configuration_dict
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.postprocessors.artifacts import (
+    save_code_execution_artifact,
+)
 from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.schemas import (
     CodeInterpreterBlock,
     CodeInterpreterFile,
@@ -661,32 +663,13 @@ class DisplayCodeInterpreterFilesPostProcessor(
                     container_file.filename, "uploading", force_publish=True
                 )
 
-            raw_mime = guess_type(container_file.filename)[0] or "text/plain"
-            mime = _kb_safe_mime(raw_mime)
-            if mime != raw_mime:
-                self._log.info(
-                    "MIME type '%s' is not supported by the Unique KB; "
-                    "uploading '%s' as '%s' so the file can be stored and downloaded.",
-                    raw_mime,
-                    container_file.filename,
-                    mime,
-                )
-            self._log.info(
-                "Uploading '%s' to knowledge base (mime=%s, %d bytes)",
-                container_file.filename,
-                mime,
-                len(file_bytes),
-            )
-
             assert self._chat_service is not None  # Checked in __init__
             upload_t0 = time.monotonic()
             content = await self._build_retry()(
-                self._chat_service.upload_to_chat_from_bytes_async,
-                content=file_bytes,
-                content_name=container_file.filename,
-                mime_type=mime,
-                skip_ingestion=True,
-                hide_in_chat=True,
+                save_code_execution_artifact,
+                chat_service=self._chat_service,
+                file=container_file,
+                file_bytes=file_bytes,
             )
             upload_ms = (time.monotonic() - upload_t0) * 1000
             if upload_ms > 2000:
@@ -929,37 +912,6 @@ def _get_file_type(filename: str) -> CodeInterpreterFileType:
     return "document"
 
 
-def _kb_safe_mime(mime: str) -> str:
-    """Return a MIME type the Unique KB will accept.
-
-    Membership is defined by ``FileMimeType`` and ``ImageMimeType`` in
-    ``unique_toolkit._common.utils.files`` (same catalog as path-based helpers
-    like ``FileMimeType.is_valid_mime``, but here we already have a resolved
-    MIME string from ``mimetypes.guess_type``, so we use StrEnum value lookup
-    instead of re-parsing a path).
-
-    The KB GraphQL API rejects many code-file MIME types (e.g. ``text/x-python``
-    for ``.py`` files).  Anything not in those enums is coerced to
-    ``text/plain`` so the file can be stored and downloaded without changing
-    its bytes.
-
-    Other ``image/*`` subtypes (e.g. ``image/jpg``) still pass through unchanged.
-    """
-    try:
-        FileMimeType(mime)
-        return mime
-    except ValueError:
-        pass
-    try:
-        ImageMimeType(mime)
-        return mime
-    except ValueError:
-        pass
-    if mime.startswith("image/"):
-        return mime
-    return "text/plain"
-
-
 def _file_title(filename: str) -> str:
     """Derive a human-readable title from a filename.
 

{unique_toolkit-2026.22.0.dev0 → unique_toolkit-2026.22.0.dev2}/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/service.py

@@ -8,10 +8,10 @@ from openai import (
     AsyncOpenAI,
     BaseModel,
     NotFoundError,
-    OpenAIError,
 )
 from openai.types.responses import ResponseCodeInterpreterToolCall, ResponseIncludable
 from openai.types.responses.tool_param import CodeInterpreter
+from pydantic import ValidationError
 from tenacity import (
     AsyncRetrying,
     before_sleep_log,
@@ -20,7 +20,11 @@ from tenacity import (
 )
 
 from unique_toolkit import ContentService, ShortTermMemoryService
-from unique_toolkit._common.execution import SafeTaskExecutor
+from unique_toolkit._common.execution import (
+    SafeTaskExecutor,
+    failsafe_async,
+)
+from unique_toolkit._common.utils.jinja.render import render_template
 from unique_toolkit.agentic.feature_flags.feature_flags import feature_flags
 from unique_toolkit.agentic.short_term_memory_manager.persistent_short_term_memory_manager import (
     PersistentShortMemoryManager,
@@ -32,6 +36,9 @@ from unique_toolkit.agentic.tools.openai_builtin.base import (
 from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.config import (
     OpenAICodeInterpreterConfig,
 )
+from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.postprocessors.artifacts import (
+    load_code_execution_metadata,
+)
 from unique_toolkit.agentic.tools.schemas import ToolPrompts
 from unique_toolkit.content.schemas import (
     Content,
@@ -62,13 +69,13 @@ def _build_upload_retry() -> AsyncRetrying:
 _SHORT_TERM_MEMORY_NAME = "container_code_execution"
 
 
-class CodeExecutionShortTermMemorySchema(BaseModel):
+class _CodeExecutionShortTermMemorySchema(BaseModel):
     container_id: str
-    file_ids: dict[str, str] = {}  # Mapping of unique file id to openai file id
+    file_paths: dict[str, str] = {}
 
 
 CodeExecutionMemoryManager = PersistentShortMemoryManager[
-    CodeExecutionShortTermMemorySchema
+    _CodeExecutionShortTermMemorySchema
 ]
 
 
@@ -83,7 +90,7 @@ def _get_container_code_execution_short_term_memory_manager(
     )
     short_term_memory_manager = PersistentShortMemoryManager(
         short_term_memory_service=short_term_memory_service,
-        short_term_memory_schema=CodeExecutionShortTermMemorySchema,
+        short_term_memory_schema=_CodeExecutionShortTermMemorySchema,
         short_term_memory_name=_SHORT_TERM_MEMORY_NAME,
     )
     return short_term_memory_manager
@@ -91,7 +98,7 @@ def _get_container_code_execution_short_term_memory_manager(
 
 async def _check_container_exists(
     client: AsyncOpenAI,
-    memory: CodeExecutionShortTermMemorySchema,
+    memory: _CodeExecutionShortTermMemorySchema,
 ) -> bool:
     try:
         container = await client.containers.retrieve(memory.container_id)
@@ -129,27 +136,15 @@ async def _create_container(
     return container.id
 
 
-async def _check_file_already_uploaded(
-    client: AsyncOpenAI,
+def _check_file_already_uploaded(
     content_id: str,
-    memory: CodeExecutionShortTermMemorySchema,
+    memory: _CodeExecutionShortTermMemorySchema,
 ) -> bool:
-    container_id = memory.container_id
-
-    if content_id not in memory.file_ids:
+    if content_id not in memory.file_paths:
         logger.info("File with id %s not in short term memory", content_id)
         return False
 
-    file_id = memory.file_ids[content_id]
-    try:
-        await client.containers.files.retrieve(
-            container_id=container_id, file_id=file_id
-        )
-        logger.info("File with id %s already uploaded to container", content_id)
-        return True
-    except OpenAIError:  # Azure API Errors are not stable
-        logger.exception("Error while retrieving file %s", file_id)
-        return False
+    return True
 
 
 async def _upload_file_to_container(
@@ -189,20 +184,17 @@ async def _upload_file_to_container(
         container_id,
     )
 
-    return openai_file.id
+    return openai_file.path
 
 
 async def _upload_files_to_container(
     client: AsyncOpenAI,
     uploaded_files: list[Content],
-    memory: CodeExecutionShortTermMemorySchema,
+    memory: _CodeExecutionShortTermMemorySchema,
     content_service: ContentService,
-) -> tuple[CodeExecutionShortTermMemorySchema, bool]:
-
+) -> tuple[_CodeExecutionShortTermMemorySchema, bool]:
     async def _check_and_upload(content: Content) -> str | None:
-        if await _check_file_already_uploaded(
-            client=client, content_id=content.id, memory=memory
-        ):
+        if _check_file_already_uploaded(content_id=content.id, memory=memory):
             return None
 
         return await _upload_file_to_container(
@@ -227,8 +219,8 @@ async def _upload_files_to_container(
 
     updated = False
     for content, result in zip(unique_contents, results):
-        if result.success and (file_id := result.unpack()) is not None:
-            memory.file_ids[content.id] = file_id
+        if result.success and (filepath := result.unpack()) is not None:
+            memory.file_paths[content.id] = filepath
             updated = True
 
     return memory, updated
@@ -263,6 +255,9 @@ class OpenAICodeInterpreterTool(OpenAIBuiltInTool[CodeInterpreter]):
         container_id: str | None,
         company_id: str = "",
         is_exclusive: bool = False,
+        user_uploaded_files: list[str] | None = None,
+        kb_uploaded_files: list[str] | None = None,
+        code_execution_files: list[str] | None = None,
     ) -> None:
         self._config = config
 
@@ -273,6 +268,10 @@ class OpenAICodeInterpreterTool(OpenAIBuiltInTool[CodeInterpreter]):
         self._company_id = company_id
         self._is_exclusive = is_exclusive
 
+        self._user_uploaded_files = user_uploaded_files
+        self._kb_uploaded_files = kb_uploaded_files
+        self._code_interpreter_artifacts = code_execution_files
+
     @property
     @override
     def name(self) -> OpenAIBuiltInToolName:
@@ -334,7 +333,10 @@ class OpenAICodeInterpreterTool(OpenAIBuiltInTool[CodeInterpreter]):
             chat_id=chat_id,
         )
 
-        memory = await memory_manager.load_async()
+        # Ignore old memory schema
+        memory = await failsafe_async(
+            failure_return_value=None, exceptions=(ValidationError,), log_exc_info=False
+        )(memory_manager.load_async)()
 
         container_updated = False
         if memory is None or not await _check_container_exists(
@@ -347,20 +349,32 @@ class OpenAICodeInterpreterTool(OpenAIBuiltInTool[CodeInterpreter]):
                 company_id=company_id,
                 expires_after_minutes=config.expires_after_minutes,
             )
-            memory = CodeExecutionShortTermMemorySchema(container_id=container_id)
+            memory = _CodeExecutionShortTermMemorySchema(container_id=container_id)
             container_updated = True
 
+        code_execution_files = []
+        user_uploaded_files = []
+        for content in uploaded_files:
+            if (metadata := load_code_execution_metadata(content)) is not None:
+                code_execution_files.append(content)
+                if metadata.container_id == memory.container_id:
+                    # Already available in container
+                    memory.file_paths[content.id] = metadata.filepath
+            else:
+                user_uploaded_files.append(content)
+
         files_to_upload: list[Content] = []
+
         if config.upload_files_in_chat_to_container:
             files_to_upload.extend(uploaded_files)
 
+        kb_files = []
         if config.additional_uploaded_documents:
-            files_to_upload.extend(
-                await _resolve_kb_contents(
-                    content_service=content_service,
-                    content_ids=config.additional_uploaded_documents,
-                )
+            kb_files = await _resolve_kb_contents(
+                content_service=content_service,
+                content_ids=config.additional_uploaded_documents,
             )
+            files_to_upload.extend(kb_files)
 
         files_updated = False
         if files_to_upload:
@@ -374,20 +388,36 @@ class OpenAICodeInterpreterTool(OpenAIBuiltInTool[CodeInterpreter]):
         if container_updated or files_updated:
             await memory_manager.save_async(memory)
 
+        def _extract_paths(contents: list[Content]) -> list[str]:
+            return [
+                memory.file_paths[content.id]
+                for content in contents
+                if content.id in memory.file_paths
+            ]
+
         return OpenAICodeInterpreterTool(
             config=config,
             container_id=memory.container_id,
             company_id=company_id,
             is_exclusive=is_exclusive,
+            user_uploaded_files=_extract_paths(user_uploaded_files),
+            kb_uploaded_files=_extract_paths(kb_files),
+            code_execution_files=_extract_paths(code_execution_files),
         )
 
     @override
     def get_tool_prompts(self) -> ToolPrompts:
+        rendered_prompt = render_template(
+            self._config.tool_description_for_system_prompt,
+            user_uploaded_files=self._user_uploaded_files,
+            kb_uploaded_files=self._kb_uploaded_files,
+            code_interpreter_artifacts=self._code_interpreter_artifacts,
+        )
         return ToolPrompts(
             name="python",  # 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_system_prompt=rendered_prompt,
             tool_format_information_for_system_prompt="",
             tool_user_prompt=self._config.tool_description_for_user_prompt,
             tool_format_information_for_user_prompt="",

{unique_toolkit-2026.22.0.dev0 → unique_toolkit-2026.22.0.dev2}/unique_toolkit/services/chat_service.py

@@ -1986,6 +1986,19 @@ class ChatService(ChatServiceDeprecated):
             chat_id=self._content_scope_chat_id,
         )
 
+    @staticmethod
+    def _filter_images_and_documents(
+        contents: list[Content],
+    ) -> tuple[list[Content], list[Content]]:
+        images: list[Content] = []
+        files: list[Content] = []
+        for c in contents:
+            if is_file_content(filename=c.key):
+                files.append(c)
+            if is_image_content(filename=c.key):
+                images.append(c)
+        return images, files
+
     def download_chat_images_and_documents(self) -> tuple[list[Content], list[Content]]:
         """Return images and documents from the content-scope chat (e.g. parent chat when subagent).
 
@@ -1995,19 +2008,34 @@ class ChatService(ChatServiceDeprecated):
         Returns:
             tuple[list[Content], list[Content]]: (images, documents) from the content-scope chat.
         """
-        images: list[Content] = []
-        files: list[Content] = []
-        for c in search_contents(
+        contents = search_contents(
             user_id=self._user_id,
             company_id=self._company_id,
             chat_id=self._content_scope_chat_id,
             where={"ownerId": {"equals": self._content_scope_chat_id}},
-        ):
-            if is_file_content(filename=c.key):
-                files.append(c)
-            if is_image_content(filename=c.key):
-                images.append(c)
-        return images, files
+        )
+        return self._filter_images_and_documents(contents)
+
+    async def download_chat_images_and_documents_async(
+        self,
+    ) -> tuple[list[Content], list[Content]]:
+        """Async version of download_chat_images_and_documents.
+
+        Return images and documents from the content-scope chat (e.g. parent chat when subagent).
+
+        Uses the service's content-scope chat id, so when running as a subagent
+        with correlation, this accesses files from the primary chat session.
+
+        Returns:
+            tuple[list[Content], list[Content]]: (images, documents) from the content-scope chat.
+        """
+        contents = await search_contents_async(
+            user_id=self._user_id,
+            company_id=self._company_id,
+            chat_id=self._content_scope_chat_id,
+            where={"ownerId": {"equals": self._content_scope_chat_id}},
+        )
+        return self._filter_images_and_documents(contents)
 
     async def search_content_chunks_async(
         self,