unique_toolkit 1.68.8__tar.gz → 1.68.10__tar.gz

{unique_toolkit-1.68.8 → unique_toolkit-1.68.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 1.68.8
+Version: 1.68.10
 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>

{unique_toolkit-1.68.8 → unique_toolkit-1.68.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "1.68.8"
+version = "1.68.10"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12"

{unique_toolkit-1.68.8 → unique_toolkit-1.68.10}/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/postprocessors/generated_files.py

@@ -41,6 +41,14 @@ from unique_toolkit.short_term_memory.service import ShortTermMemoryService
 logger = logging.getLogger(__name__)
 
 
+class _ChatLoggerAdapter(logging.LoggerAdapter[logging.Logger]):
+    """LoggerAdapter that prefixes every message with ``[chat_id=…]``."""
+
+    def process(self, msg: str, kwargs):  # type: ignore[override]
+        extra = dict(self.extra) if self.extra else {}  # type: ignore[arg-type]
+        return f"[chat_id={extra.get('chat_id', 'n/a')}] {msg}", kwargs
+
+
 class DisplayCodeInterpreterFilesPostProcessorConfig(BaseModel):
     model_config = get_configuration_dict()
     upload_to_chat: SkipJsonSchema[bool] = Field(
@@ -53,8 +61,8 @@ class DisplayCodeInterpreterFilesPostProcessorConfig(BaseModel):
     )
 
     file_download_failed_message: str = Field(
-        default="⚠️ File download failed ...",
-        description="The message to display when a file download fails.",
+        default="⚠️ File could not be generated. Please try again.",
+        description="The message to display when a file download fails after all retry attempts.",
     )
     max_concurrent_file_downloads: int = Field(
         default=10,
@@ -122,6 +130,10 @@ class DisplayCodeInterpreterFilesPostProcessor(
         if self._chat_service is None:
             raise ValueError("ChatService is required if uploadToChat is True")
 
+        self._log: logging.LoggerAdapter[logging.Logger] = _ChatLoggerAdapter(
+            logger, {"chat_id": chat_id or "n/a"}
+        )
+
         self._orphan_code_blocks: list[CodeInterpreterBlock] = []
         self._short_term_memory_manager = None
         if chat_id is not None and user_id is not None and company_id is not None:
@@ -131,14 +143,44 @@ class DisplayCodeInterpreterFilesPostProcessor(
                 chat_id=chat_id,
             )
 
+    def _build_retry(self) -> AsyncRetrying:
+        """Build a tenacity retry policy from the current config.
+
+        Reused for both container-file downloads and knowledge-base uploads so
+        that every I/O call gets the same exponential-backoff behaviour.
+        """
+        return AsyncRetrying(
+            stop=stop_after_attempt(1 + self._config.max_download_retries),
+            wait=wait_exponential(multiplier=self._config.download_retry_base_delay),
+            before_sleep=before_sleep_log(self._log, logging.WARNING),  # type: ignore[arg-type]
+            reraise=True,
+        )
+
     @override
     async def run(self, loop_response: ResponsesLanguageModelStreamResponse) -> None:
-        logger.info("Fetching and adding code interpreter files to the response")
+        self._log.info("run() started — fetching and uploading code interpreter files")
 
         container_files = loop_response.container_files
-        logger.info("Found %s container files", len(container_files))
+        self._log.info(
+            "run() found %d container file annotation(s): %s",
+            len(container_files),
+            [cf.filename for cf in container_files],
+        )
 
-        self._content_map: dict[str, str | None] = await self._load_previous_files()  # type: ignore
+        try:
+            self._content_map: dict[str, str | None] = await self._load_previous_files()  # type: ignore
+            if self._content_map:
+                self._log.info(
+                    "run() loaded %d previous file(s) from short-term memory: %s",
+                    len(self._content_map),
+                    list(self._content_map.keys()),
+                )
+        except Exception:
+            self._log.exception(
+                "run() failed to load previous files from short-term memory; "
+                "proceeding without previous file context."
+            )
+            self._content_map = {}
 
         semaphore = asyncio.Semaphore(self._config.max_concurrent_file_downloads)
         tasks = [
@@ -150,32 +192,75 @@ class DisplayCodeInterpreterFilesPostProcessor(
         results = await asyncio.gather(*tasks)
 
         for citation, result in zip(container_files, results):
-            # Overwrite if file has been re-generated
             self._content_map[citation.filename] = (
                 result.content_id if result is not None else None
             )
 
-        await self._save_generated_files(
-            [
-                _ContentInfo(filename=filename, content_id=content_id)
-                for filename, content_id in self._content_map.items()
-                if content_id is not None
-            ]
+        succeeded = {f: cid for f, cid in self._content_map.items() if cid is not None}
+        failed = [f for f, cid in self._content_map.items() if cid is None]
+        self._log.info(
+            "run() download/upload complete — %d succeeded, %d failed. "
+            "succeeded=%s, failed=%s",
+            len(succeeded),
+            len(failed),
+            list(succeeded.keys()),
+            failed,
         )
 
-        if feature_flags.enable_code_execution_fence_un_17972.is_enabled(
-            self._company_id
-        ):
-            self._orphan_code_blocks = await self._upload_orphan_code_as_txt(
-                loop_response
+        try:
+            await self._save_generated_files(
+                [
+                    _ContentInfo(filename=filename, content_id=content_id)
+                    for filename, content_id in self._content_map.items()
+                    if content_id is not None
+                ]
+            )
+        except Exception:
+            self._log.exception(
+                "run() failed to save generated files to short-term memory; "
+                "file replacement will still proceed."
+            )
+
+        try:
+            if feature_flags.enable_code_execution_fence_un_17972.is_enabled(
+                self._company_id
+            ):
+                self._orphan_code_blocks = await self._upload_orphan_code_as_txt(
+                    loop_response
+                )
+                if self._orphan_code_blocks:
+                    self._log.info(
+                        "run() produced %d orphan code block(s)",
+                        len(self._orphan_code_blocks),
+                    )
+            else:
+                self._orphan_code_blocks = []
+        except Exception:
+            self._log.exception(
+                "run() failed to process orphan code blocks; "
+                "file replacement will still proceed."
             )
-        else:
             self._orphan_code_blocks = []
 
+        self._log.info(
+            "run() finished — content_map has %d entries (%d with content_id, %d failed)",
+            len(self._content_map),
+            len(succeeded),
+            len(failed),
+        )
+
     @override
     def apply_postprocessing_to_response(
         self, loop_response: ResponsesLanguageModelStreamResponse
     ) -> bool:
+        self._log.info(
+            "apply_postprocessing started — %d file(s) in content_map, fence_ff=%s",
+            len(self._content_map),
+            feature_flags.enable_code_execution_fence_un_17972.is_enabled(
+                self._company_id
+            ),
+        )
+
         if loop_response.message.references is None:
             loop_response.message.references = []
         if loop_response.message.text is None:
@@ -187,6 +272,10 @@ class DisplayCodeInterpreterFilesPostProcessor(
             self._company_id
         )
 
+        replaced_files: list[str] = []
+        missed_files: list[str] = []
+        error_files: list[str] = []
+
         for filename, content_id in self._content_map.items():
             if content_id is None:
                 loop_response.message.text, replaced = _replace_container_file_error(
@@ -195,12 +284,14 @@ class DisplayCodeInterpreterFilesPostProcessor(
                     error_message=self._config.file_download_failed_message,
                 )
                 changed |= replaced
+                error_files.append(filename)
                 continue
 
-            is_image = (guess_type(filename)[0] or "").startswith("image/")
-            is_html = (guess_type(filename)[0] or "") == "text/html"
+            mime = guess_type(filename)[0] or ""
+            is_image = mime.startswith("image/")
+            is_html = mime == "text/html"
+            file_type = "image" if is_image else ("html" if is_html else "document")
 
-            # Images
             if is_image:
                 loop_response.message.text, replaced = (
                     _replace_container_image_citation(
@@ -229,6 +320,17 @@ class DisplayCodeInterpreterFilesPostProcessor(
                 )
                 changed |= replaced
 
+            if replaced:
+                replaced_files.append(filename)
+                self._log.info(
+                    "Replaced sandbox link for '%s' (type=%s, content_id=%s)",
+                    filename,
+                    file_type,
+                    content_id,
+                )
+            else:
+                missed_files.append(filename)
+
             is_html_rendered = is_html
             if replaced and not (is_image or is_html_rendered or fence_ff_on):
                 loop_response.message.references.append(
@@ -242,15 +344,30 @@ class DisplayCodeInterpreterFilesPostProcessor(
                 )
                 ref_number += 1
 
+        self._log.info(
+            "Stage-1 replacement summary — replaced=%s, missed=%s, error=%s",
+            replaced_files,
+            missed_files,
+            error_files,
+        )
+
         if fence_ff_on:
             code_blocks = _build_code_blocks(loop_response, self._content_map)
+            self._log.info(
+                "Fence injection — %d code block(s), files: %s",
+                len(code_blocks),
+                [f.filename for b in code_blocks for f in b.files],
+            )
             _warn_unmatched_code_blocks(self._content_map, code_blocks)
             text_before = loop_response.message.text
             loop_response.message.text = _inject_code_execution_fences(
                 loop_response.message.text,
                 code_blocks,
             )
-            changed |= loop_response.message.text != text_before
+            fences_changed = loop_response.message.text != text_before
+            changed |= fences_changed
+            if fences_changed:
+                self._log.info("Fence injection modified the message text")
 
             if self._orphan_code_blocks:
                 fence_id = _get_next_fence_id(loop_response.message.text)
@@ -259,12 +376,29 @@ class DisplayCodeInterpreterFilesPostProcessor(
                     loop_response.message.text.rstrip() + "\n\n" + orphan_fences
                 )
                 changed = True
+                self._log.info(
+                    "Appended %d orphan fence(s)", len(self._orphan_code_blocks)
+                )
 
         _warn_missing_content_ids(loop_response.message.text, self._content_map)
         loop_response.message.text, dangling_replaced = _replace_dangling_sandbox_links(
             loop_response.message.text, self._config.file_download_failed_message
         )
         changed |= dangling_replaced
+        if dangling_replaced:
+            self._log.warning(
+                "apply_postprocessing found and replaced dangling sandbox links"
+            )
+
+        self._log.info(
+            "apply_postprocessing finished — changed=%s, replaced=%d, missed=%d, "
+            "error=%d, dangling_cleaned=%s",
+            changed,
+            len(replaced_files),
+            len(missed_files),
+            len(error_files),
+            dangling_replaced,
+        )
 
         return changed
 
@@ -279,52 +413,64 @@ class DisplayCodeInterpreterFilesPostProcessor(
         semaphore: asyncio.Semaphore,
     ) -> _ContentInfo | None:
         async with semaphore:
-            retry = AsyncRetrying(
-                stop=stop_after_attempt(1 + self._config.max_download_retries),
-                wait=wait_exponential(
-                    multiplier=self._config.download_retry_base_delay
-                ),
-                before_sleep=before_sleep_log(logger, logging.WARNING),
-                reraise=True,
+            self._log.info(
+                "Downloading container file '%s' (container=%s, file_id=%s)",
+                container_file.filename,
+                container_file.container_id,
+                container_file.file_id,
             )
+            retry = self._build_retry()
             file_content = await retry(
                 self._client.containers.files.content.retrieve,
                 container_id=container_file.container_id,
                 file_id=container_file.file_id,
             )
+            self._log.info(
+                "Downloaded container file '%s' (%d bytes)",
+                container_file.filename,
+                len(file_content.content),
+            )
 
-            logger.info(
-                "Uploading file content for %s to knowledge base",
+            mime = guess_type(container_file.filename)[0] or "text/plain"
+            self._log.info(
+                "Uploading '%s' to knowledge base (mime=%s)",
                 container_file.filename,
+                mime,
             )
 
             assert self._chat_service is not None  # Checked in __init__
-            content = await self._chat_service.upload_to_chat_from_bytes_async(
+            content = await self._build_retry()(
+                self._chat_service.upload_to_chat_from_bytes_async,
                 content=file_content.content,
                 content_name=container_file.filename,
-                mime_type=guess_type(container_file.filename)[0] or "text/plain",
+                mime_type=mime,
                 skip_ingestion=True,
                 hide_in_chat=True,
             )
 
+            self._log.info(
+                "Uploaded '%s' — content_id=%s",
+                container_file.filename,
+                content.id,
+            )
             return _ContentInfo(filename=container_file.filename, content_id=content.id)
 
     async def _load_previous_files(self) -> dict[str, str]:
         if self._short_term_memory_manager is None:
             return {}
 
-        logger.info(
+        self._log.info(
             "Loading previously generated code interpreter files from short term memory"
         )
         memory = await self._short_term_memory_manager.load_async()
 
         if memory is None:
-            logger.info(
+            self._log.info(
                 "No previously generated code interpreter files found in short term memory"
             )
             return {}
 
-        logger.info(
+        self._log.info(
             "Found %s previously generated code interpreter files", len(memory.root)
         )
 
@@ -378,7 +524,8 @@ class DisplayCodeInterpreterFilesPostProcessor(
             filename = f"code_{i + 1}.txt" if use_numeric_suffix else "code.txt"
 
             try:
-                content = await self._chat_service.upload_to_chat_from_bytes_async(
+                content = await self._build_retry()(
+                    self._chat_service.upload_to_chat_from_bytes_async,
                     content=txt_content.encode("utf-8"),
                     content_name=filename,
                     mime_type="text/plain",
@@ -386,7 +533,7 @@ class DisplayCodeInterpreterFilesPostProcessor(
                     hide_in_chat=True,
                 )
             except Exception:
-                logger.exception(
+                self._log.exception(
                     "Failed to upload orphan code block as txt for call %d ('%s'); "
                     "falling back to legacy code block display.",
                     i,
@@ -402,7 +549,7 @@ class DisplayCodeInterpreterFilesPostProcessor(
             orphan_blocks.append(
                 CodeInterpreterBlock(code=call.code, files=[orphan_file])
             )
-            logger.info(
+            self._log.info(
                 "Uploaded orphan code block as '%s' (content_id=%s)",
                 filename,
                 content.id,
@@ -938,7 +1085,7 @@ def _replace_container_file_citation(
     original pre-fence behaviour is restored: the link is replaced with <sup>N</sup>
     and the file remains accessible via the references panel.
     """
-    file_markdown = rf"\[.*?\]\(sandbox:/mnt/data/{re.escape(filename)}\)"
+    file_markdown = rf"!?\[.*?\]\(sandbox:/mnt/data/{re.escape(filename)}\)"
 
     if not re.search(file_markdown, text):
         logger.warning(