unique_toolkit 1.68.9__tar.gz → 1.68.11__tar.gz

{unique_toolkit-1.68.9 → unique_toolkit-1.68.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 1.68.9
+Version: 1.68.11
 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.9 → unique_toolkit-1.68.11}/pyproject.toml

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

{unique_toolkit-1.68.9 → unique_toolkit-1.68.11}/unique_toolkit/agentic/tools/openai_builtin/code_interpreter/postprocessors/generated_files.py

@@ -2,6 +2,8 @@ import asyncio
 import json
 import logging
 import re
+import time
+from dataclasses import dataclass
 from mimetypes import guess_type
 from typing import NamedTuple, override
 
@@ -41,6 +43,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 +63,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,
@@ -68,6 +78,14 @@ class DisplayCodeInterpreterFilesPostProcessorConfig(BaseModel):
         default=0.5,
         description="Base delay in seconds for exponential backoff between download/upload retries.",
     )
+    progress_update_interval: float = Field(
+        default=3.0,
+        description="Minimum seconds between progress message updates sent to the user.",
+    )
+    download_chunk_size: int = Field(
+        default=8192,
+        description="Chunk size in bytes for streaming container file downloads.",
+    )
 
 
 class _ContentInfo(NamedTuple):
@@ -95,6 +113,141 @@ def _init_short_term_memory_manager(
     )
 
 
+@dataclass
+class _FileState:
+    phase: str = "pending"
+    percent: int | None = None
+    elapsed_seconds: float = 0.0
+    retry_attempt: int = 0
+    max_retries: int = 0
+
+
+class _FileProgressTracker:
+    """Concurrency-safe coordinator that publishes file download/upload progress
+    to the user-visible assistant message in real time.
+
+    Replaces sandbox links inline with progress text and appends a summary block.
+    Throttles ``modify_assistant_message_async`` calls to avoid API spam.
+    """
+
+    def __init__(
+        self,
+        filenames: list[str],
+        original_text: str,
+        chat_service: ChatService,
+        log: logging.LoggerAdapter[logging.Logger],
+        min_publish_interval: float = 3.0,
+    ) -> None:
+        self._states: dict[str, _FileState] = {f: _FileState() for f in filenames}
+        self._original_text = original_text
+        self._chat_service = chat_service
+        self._log = log
+        self._lock = asyncio.Lock()
+        self._last_publish_time = 0.0
+        self._min_interval = min_publish_interval
+
+    async def update(
+        self,
+        filename: str,
+        phase: str,
+        *,
+        percent: int | None = None,
+        elapsed_seconds: float = 0.0,
+        retry_attempt: int = 0,
+        max_retries: int = 0,
+        force_publish: bool = False,
+    ) -> None:
+        async with self._lock:
+            state = self._states.get(filename)
+            if state is None:
+                return
+            state.phase = phase
+            state.percent = percent
+            state.elapsed_seconds = elapsed_seconds
+            state.retry_attempt = retry_attempt
+            state.max_retries = max_retries
+
+            now = time.monotonic()
+            should_publish = (
+                force_publish
+                or phase in ("done", "failed", "uploading")
+                or now - self._last_publish_time >= self._min_interval
+            )
+            if should_publish:
+                await self._publish()
+                self._last_publish_time = now
+
+    async def publish_initial(self) -> None:
+        """Send the first progress update (marks all files as pending)."""
+        async with self._lock:
+            for state in self._states.values():
+                state.phase = "downloading"
+            await self._publish()
+            self._last_publish_time = time.monotonic()
+
+    async def _publish(self) -> None:
+        text = self._build_progress_text()
+        try:
+            await self._chat_service.modify_assistant_message_async(content=text)
+        except Exception:
+            self._log.debug("Failed to publish progress update", exc_info=True)
+
+    def _build_progress_text(self) -> str:
+        text = self._original_text
+
+        for filename, state in self._states.items():
+            if state.phase == "pending":
+                continue
+            progress = self._format_inline(filename, state)
+            pattern = rf"!?\[.*?\]\(sandbox:/mnt/data/{re.escape(filename)}\)"
+            text = re.sub(pattern, progress, text)
+
+        active = {
+            f: s for f, s in self._states.items() if s.phase not in ("done", "pending")
+        }
+        if active:
+            lines = ["---", "Preparing files:"]
+            for filename, state in active.items():
+                lines.append(f"- {filename} — {self._format_summary(state)}")
+            text = text.rstrip() + "\n\n" + "\n".join(lines)
+
+        return text
+
+    @staticmethod
+    def _format_inline(filename: str, state: _FileState) -> str:
+        if state.phase == "downloading":
+            base = f"Downloading {filename}..."
+            if state.retry_attempt > 0:
+                base += f" retry {state.retry_attempt}/{state.max_retries}"
+            if state.percent is not None:
+                return f"{base} {state.percent}%"
+            if state.elapsed_seconds > 0:
+                return f"{base} ({int(state.elapsed_seconds)}s)"
+            return base
+        if state.phase == "uploading":
+            return f"Uploading {filename}..."
+        if state.phase == "failed":
+            return "File could not be generated. Please try again."
+        return filename
+
+    @staticmethod
+    def _format_summary(state: _FileState) -> str:
+        if state.phase == "downloading":
+            base = "Downloading"
+            if state.retry_attempt > 0:
+                base += f" (retry {state.retry_attempt}/{state.max_retries})"
+            if state.percent is not None:
+                return f"{base} {state.percent}%"
+            if state.elapsed_seconds > 0:
+                return f"{base} ({int(state.elapsed_seconds)}s)"
+            return f"{base}..."
+        if state.phase == "uploading":
+            return "Uploading..."
+        if state.phase == "failed":
+            return "Failed"
+        return "Pending"
+
+
 class DisplayCodeInterpreterFilesPostProcessor(
     ResponsesApiPostprocessor,
 ):
@@ -122,6 +275,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:
@@ -140,40 +297,73 @@ class DisplayCodeInterpreterFilesPostProcessor(
         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(logger, logging.WARNING),
+            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],
+        )
 
         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:
-            logger.exception(
-                "Failed to load previous code interpreter files from short-term "
-                "memory; proceeding without previous file context."
+            self._log.exception(
+                "run() failed to load previous files from short-term memory; "
+                "proceeding without previous file context."
             )
             self._content_map = {}
 
+        tracker: _FileProgressTracker | None = None
+        if container_files and self._chat_service is not None:
+            tracker = _FileProgressTracker(
+                filenames=[cf.filename for cf in container_files],
+                original_text=loop_response.message.text or "",
+                chat_service=self._chat_service,
+                log=self._log,
+                min_publish_interval=self._config.progress_update_interval,
+            )
+            await tracker.publish_initial()
+
         semaphore = asyncio.Semaphore(self._config.max_concurrent_file_downloads)
         tasks = [
             self._download_and_upload_container_files_to_knowledge_base(
-                citation, semaphore
+                citation, semaphore, tracker
             )
             for citation in container_files
         ]
         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
             )
+            if result is None and tracker:
+                await tracker.update(citation.filename, "failed", force_publish=True)
+
+        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,
+        )
 
         try:
             await self._save_generated_files(
@@ -184,9 +374,9 @@ class DisplayCodeInterpreterFilesPostProcessor(
                 ]
             )
         except Exception:
-            logger.exception(
-                "Failed to save generated code interpreter files to short-term "
-                "memory; file replacement will still proceed."
+            self._log.exception(
+                "run() failed to save generated files to short-term memory; "
+                "file replacement will still proceed."
             )
 
         try:
@@ -196,19 +386,39 @@ class DisplayCodeInterpreterFilesPostProcessor(
                 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:
-            logger.exception(
-                "Failed to process orphan code blocks; "
+            self._log.exception(
+                "run() failed to process orphan code blocks; "
                 "file replacement will still proceed."
             )
             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:
@@ -220,6 +430,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(
@@ -228,12 +442,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(
@@ -262,6 +478,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(
@@ -275,15 +502,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)
@@ -292,12 +534,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
 
@@ -310,48 +569,166 @@ class DisplayCodeInterpreterFilesPostProcessor(
         self,
         container_file: AnnotationContainerFileCitation,
         semaphore: asyncio.Semaphore,
+        tracker: _FileProgressTracker | None = None,
     ) -> _ContentInfo | None:
         async with semaphore:
-            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(
+                "Downloading container file '%s' (container=%s, file_id=%s)",
+                container_file.filename,
+                container_file.container_id,
+                container_file.file_id,
             )
 
-            logger.info(
-                "Uploading file content for %s to knowledge base",
+            file_bytes = await self._download_file_bytes_with_progress(
+                container_file, tracker
+            )
+
+            self._log.info(
+                "Downloaded container file '%s' (%d bytes)",
                 container_file.filename,
+                len(file_bytes),
+            )
+
+            if tracker:
+                await tracker.update(
+                    container_file.filename, "uploading", force_publish=True
+                )
+
+            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._build_retry()(
                 self._chat_service.upload_to_chat_from_bytes_async,
-                content=file_content.content,
+                content=file_bytes,
                 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,
             )
 
+            if tracker:
+                await tracker.update(
+                    container_file.filename, "done", force_publish=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 _download_file_bytes_with_progress(
+        self,
+        container_file: AnnotationContainerFileCitation,
+        tracker: _FileProgressTracker | None,
+    ) -> bytes:
+        """Download container file with streaming progress and manual retry loop.
+
+        Uses ``with_streaming_response`` so we can read ``content-length``
+        and report download percentage.  Falls back to elapsed-time display
+        when the header is absent.
+        """
+        max_attempts = 1 + self._config.max_download_retries
+        download_start = time.monotonic()
+        last_exception: Exception | None = None
+
+        for attempt_num in range(1, max_attempts + 1):
+            retry_num = attempt_num - 1
+            try:
+                if retry_num > 0:
+                    delay = self._config.download_retry_base_delay * (2**retry_num)
+                    self._log.warning(
+                        "Retrying download of '%s' (attempt %d/%d) after %.1fs",
+                        container_file.filename,
+                        attempt_num,
+                        max_attempts,
+                        delay,
+                    )
+                    if tracker:
+                        await tracker.update(
+                            container_file.filename,
+                            "downloading",
+                            elapsed_seconds=time.monotonic() - download_start,
+                            retry_attempt=retry_num,
+                            max_retries=self._config.max_download_retries,
+                            force_publish=True,
+                        )
+                    await asyncio.sleep(delay)
+
+                return await self._stream_download_bytes(
+                    container_file, tracker, retry_num, download_start
+                )
+            except Exception as exc:
+                last_exception = exc
+                self._log.warning(
+                    "Download attempt %d/%d failed for '%s': %s",
+                    attempt_num,
+                    max_attempts,
+                    container_file.filename,
+                    exc,
+                )
+
+        assert last_exception is not None
+        raise last_exception
+
+    async def _stream_download_bytes(
+        self,
+        container_file: AnnotationContainerFileCitation,
+        tracker: _FileProgressTracker | None,
+        retry_attempt: int,
+        download_start: float,
+    ) -> bytes:
+        """Stream-download a single container file, reporting chunk-level progress."""
+        async with (
+            self._client.containers.files.content.with_streaming_response.retrieve(
+                file_id=container_file.file_id,
+                container_id=container_file.container_id,
+            ) as response
+        ):
+            content_length = response.headers.get("content-length")
+            total = int(content_length) if content_length else None
+            chunks: list[bytes] = []
+            downloaded = 0
+            async for chunk in response.iter_bytes(
+                chunk_size=self._config.download_chunk_size
+            ):
+                chunks.append(chunk)
+                downloaded += len(chunk)
+                if tracker:
+                    pct = (downloaded * 100 // total) if total else None
+                    elapsed = time.monotonic() - download_start
+                    await tracker.update(
+                        container_file.filename,
+                        "downloading",
+                        percent=pct,
+                        elapsed_seconds=elapsed,
+                        retry_attempt=retry_attempt,
+                        max_retries=self._config.max_download_retries,
+                    )
+            return b"".join(chunks)
+
     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)
         )
 
@@ -414,7 +791,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,
@@ -430,7 +807,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,