unique_toolkit 1.68.10__tar.gz → 1.68.12__tar.gz

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

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

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

{unique_toolkit-1.68.10 → unique_toolkit-1.68.12}/unique_toolkit/agentic/short_term_memory_manager/persistent_short_term_memory_manager.py

@@ -85,7 +85,17 @@ class PersistentShortMemoryManager(Generic[TSchema]):
             self._short_term_memory_name,
         )
 
-        self._log_not_found() if not result.success else self._log_found()
+        if not result.success:
+            self._log_not_found()
+            if result.exception is not None:
+                logger.warning(
+                    "Short term memory lookup failed with %s: %s",
+                    type(result.exception).__name__,
+                    result.exception,
+                    exc_info=result.exception,
+                )
+        else:
+            self._log_found()
 
         return result.unpack(default=None)
 
@@ -95,7 +105,17 @@ class PersistentShortMemoryManager(Generic[TSchema]):
             self._short_term_memory_name,
         )
 
-        self._log_not_found() if not result.success else self._log_found()
+        if not result.success:
+            self._log_not_found()
+            if result.exception is not None:
+                logger.warning(
+                    "Short term memory lookup failed with %s: %s",
+                    type(result.exception).__name__,
+                    result.exception,
+                    exc_info=result.exception,
+                )
+        else:
+            self._log_found()
 
         return result.unpack(default=None)
 

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

@@ -2,9 +2,12 @@ import asyncio
 import json
 import logging
 import re
+import time
+from dataclasses import dataclass
 from mimetypes import guess_type
 from typing import NamedTuple, override
 
+import httpx
 from openai import AsyncOpenAI
 from openai.types.responses import ResponseCodeInterpreterToolCall
 from openai.types.responses.response_output_text import AnnotationContainerFileCitation
@@ -76,6 +79,20 @@ 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.",
+    )
+    download_read_timeout: float = Field(
+        default=120.0,
+        description="HTTP read timeout in seconds for container file downloads. "
+        "Applies per SDK attempt. The OpenAI SDK default of 600s is too generous "
+        "for small generated files; a shorter timeout lets us retry sooner.",
+    )
 
 
 class _ContentInfo(NamedTuple):
@@ -103,6 +120,157 @@ 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 tick_elapsed(self, filename: str, elapsed_seconds: float) -> None:
+        """Publish an elapsed-time update without changing percent or retry state.
+
+        Used by the background ticker to show time-based progress while the
+        OpenAI API hasn't started returning bytes yet.
+        """
+        async with self._lock:
+            state = self._states.get(filename)
+            if state is None or state.phase != "downloading":
+                return
+            state.elapsed_seconds = elapsed_seconds
+            now = time.monotonic()
+            if now - self._last_publish_time >= self._min_interval:
+                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,
 ):
@@ -158,6 +326,7 @@ class DisplayCodeInterpreterFilesPostProcessor(
 
     @override
     async def run(self, loop_response: ResponsesLanguageModelStreamResponse) -> None:
+        run_t0 = time.monotonic()
         self._log.info("run() started — fetching and uploading code interpreter files")
 
         container_files = loop_response.container_files
@@ -167,6 +336,7 @@ class DisplayCodeInterpreterFilesPostProcessor(
             [cf.filename for cf in container_files],
         )
 
+        phase_t0 = time.monotonic()
         try:
             self._content_map: dict[str, str | None] = await self._load_previous_files()  # type: ignore
             if self._content_map:
@@ -181,20 +351,36 @@ class DisplayCodeInterpreterFilesPostProcessor(
                 "proceeding without previous file context."
             )
             self._content_map = {}
+        load_ms = (time.monotonic() - phase_t0) * 1000
+
+        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()
 
+        phase_t0 = time.monotonic()
         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)
+        download_upload_ms = (time.monotonic() - phase_t0) * 1000
 
         for citation, result in zip(container_files, results):
             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]
@@ -207,6 +393,7 @@ class DisplayCodeInterpreterFilesPostProcessor(
             failed,
         )
 
+        phase_t0 = time.monotonic()
         try:
             await self._save_generated_files(
                 [
@@ -220,7 +407,9 @@ class DisplayCodeInterpreterFilesPostProcessor(
                 "run() failed to save generated files to short-term memory; "
                 "file replacement will still proceed."
             )
+        save_ms = (time.monotonic() - phase_t0) * 1000
 
+        phase_t0 = time.monotonic()
         try:
             if feature_flags.enable_code_execution_fence_un_17972.is_enabled(
                 self._company_id
@@ -241,18 +430,28 @@ class DisplayCodeInterpreterFilesPostProcessor(
                 "file replacement will still proceed."
             )
             self._orphan_code_blocks = []
+        orphan_ms = (time.monotonic() - phase_t0) * 1000
 
+        total_ms = (time.monotonic() - run_t0) * 1000
         self._log.info(
-            "run() finished — content_map has %d entries (%d with content_id, %d failed)",
+            "run() finished — content_map has %d entries (%d ok, %d failed). "
+            "Timing: total=%.0fms, load_stm=%.0fms, download_upload=%.0fms, "
+            "save_stm=%.0fms, orphan=%.0fms",
             len(self._content_map),
             len(succeeded),
             len(failed),
+            total_ms,
+            load_ms,
+            download_upload_ms,
+            save_ms,
+            orphan_ms,
         )
 
     @override
     def apply_postprocessing_to_response(
         self, loop_response: ResponsesLanguageModelStreamResponse
     ) -> bool:
+        apply_t0 = time.monotonic()
         self._log.info(
             "apply_postprocessing started — %d file(s) in content_map, fence_ff=%s",
             len(self._content_map),
@@ -390,14 +589,16 @@ class DisplayCodeInterpreterFilesPostProcessor(
                 "apply_postprocessing found and replaced dangling sandbox links"
             )
 
+        apply_ms = (time.monotonic() - apply_t0) * 1000
         self._log.info(
             "apply_postprocessing finished — changed=%s, replaced=%d, missed=%d, "
-            "error=%d, dangling_cleaned=%s",
+            "error=%d, dangling_cleaned=%s (%.0fms)",
             changed,
             len(replaced_files),
             len(missed_files),
             len(error_files),
             dangling_replaced,
+            apply_ms,
         )
 
         return changed
@@ -411,50 +612,215 @@ class DisplayCodeInterpreterFilesPostProcessor(
         self,
         container_file: AnnotationContainerFileCitation,
         semaphore: asyncio.Semaphore,
+        tracker: _FileProgressTracker | None = None,
     ) -> _ContentInfo | None:
         async with semaphore:
+            pipeline_t0 = time.monotonic()
             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,
+
+            download_t0 = time.monotonic()
+            file_bytes = await self._download_file_bytes_with_progress(
+                container_file, tracker
             )
+            download_ms = (time.monotonic() - download_t0) * 1000
+
             self._log.info(
-                "Downloaded container file '%s' (%d bytes)",
+                "Downloaded container file '%s' (%d bytes, %.0fms)",
                 container_file.filename,
-                len(file_content.content),
+                len(file_bytes),
+                download_ms,
             )
 
+            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)",
+                "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_content.content,
+                content=file_bytes,
                 content_name=container_file.filename,
                 mime_type=mime,
                 skip_ingestion=True,
                 hide_in_chat=True,
             )
+            upload_ms = (time.monotonic() - upload_t0) * 1000
+
+            if tracker:
+                await tracker.update(
+                    container_file.filename, "done", force_publish=True
+                )
 
+            pipeline_ms = (time.monotonic() - pipeline_t0) * 1000
             self._log.info(
-                "Uploaded '%s' — content_id=%s",
+                "Uploaded '%s' — content_id=%s. "
+                "Timing: download=%.0fms, upload=%.0fms, pipeline=%.0fms",
                 container_file.filename,
                 content.id,
+                download_ms,
+                upload_ms,
+                pipeline_ms,
             )
             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.
+
+        A background ticker task publishes elapsed-time updates at regular
+        intervals even when the OpenAI API is slow to return the first byte.
+        """
+        max_attempts = 1 + self._config.max_download_retries
+        download_start = time.monotonic()
+        last_exception: Exception | None = None
+
+        ticker_task: asyncio.Task[None] | None = None
+        if tracker:
+
+            async def _ticker() -> None:
+                while True:
+                    await asyncio.sleep(self._config.progress_update_interval)
+                    elapsed = time.monotonic() - download_start
+                    self._log.info(
+                        "Still downloading '%s' (%.0fs elapsed)",
+                        container_file.filename,
+                        elapsed,
+                    )
+                    await tracker.tick_elapsed(container_file.filename, elapsed)
+
+            ticker_task = asyncio.create_task(_ticker())
+
+        try:
+            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 - 1)
+                        )
+                        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
+        finally:
+            if ticker_task is not None:
+                ticker_task.cancel()
+                try:
+                    await ticker_task
+                except BaseException:
+                    pass
+
+    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.
+
+        SDK-level retries are disabled (``max_retries=0``) so that only our
+        manual retry loop retries, with full visibility in logs and progress
+        tracker.  A shorter read timeout avoids waiting 10 minutes on a
+        stalled connection.
+        """
+        t0 = time.monotonic()
+        async with self._client.with_options(
+            max_retries=0,
+            timeout=httpx.Timeout(5.0, read=self._config.download_read_timeout),
+        ).containers.files.content.with_streaming_response.retrieve(
+            file_id=container_file.file_id,
+            container_id=container_file.container_id,
+        ) as response:
+            first_byte_ms = (time.monotonic() - t0) * 1000
+            self._log.info(
+                "Stream opened for '%s' — first-byte latency=%.0fms, "
+                "content-length=%s, status=%s",
+                container_file.filename,
+                first_byte_ms,
+                response.headers.get("content-length", "unknown"),
+                response.status_code,
+            )
+            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,
+                    )
+            transfer_ms = (time.monotonic() - t0) * 1000 - first_byte_ms
+            self._log.info(
+                "Stream complete for '%s' — transfer=%.0fms, total=%d bytes",
+                container_file.filename,
+                transfer_ms,
+                downloaded,
+            )
+            return b"".join(chunks)
+
     async def _load_previous_files(self) -> dict[str, str]:
         if self._short_term_memory_manager is None:
             return {}
@@ -462,16 +828,22 @@ class DisplayCodeInterpreterFilesPostProcessor(
         self._log.info(
             "Loading previously generated code interpreter files from short term memory"
         )
+        t0 = time.monotonic()
         memory = await self._short_term_memory_manager.load_async()
+        elapsed_ms = (time.monotonic() - t0) * 1000
 
         if memory is None:
             self._log.info(
-                "No previously generated code interpreter files found in short term memory"
+                "No previously generated code interpreter files found "
+                "in short term memory (%.0fms)",
+                elapsed_ms,
             )
             return {}
 
         self._log.info(
-            "Found %s previously generated code interpreter files", len(memory.root)
+            "Found %s previously generated code interpreter files (%.0fms)",
+            len(memory.root),
+            elapsed_ms,
         )
 
         return {content.filename: content.content_id for content in memory.root}
@@ -480,9 +852,15 @@ class DisplayCodeInterpreterFilesPostProcessor(
         if self._short_term_memory_manager is None or len(content_infos) == 0:
             return
 
+        t0 = time.monotonic()
         await self._short_term_memory_manager.save_async(
             _DisplayedFilesShortTermMemorySchema(root=content_infos)
         )
+        self._log.info(
+            "Saved %d generated file(s) to short term memory (%.0fms)",
+            len(content_infos),
+            (time.monotonic() - t0) * 1000,
+        )
 
     async def _upload_orphan_code_as_txt(
         self, loop_response: ResponsesLanguageModelStreamResponse