unique_toolkit 1.52.0__tar.gz → 1.53.1__tar.gz

{unique_toolkit-1.52.0 → unique_toolkit-1.53.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 1.52.0
+Version: 1.53.1
 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>
@@ -29,6 +29,7 @@ Requires-Dist: pandas>=2.3.0,<3
 Requires-Dist: click>=8.1.0,<9
 Requires-Dist: deepdiff>=8.6.1,<9
 Requires-Dist: pypandoc>=1.16.2,<2
+Requires-Dist: tenacity>=9.1.4
 Requires-Dist: fastapi>=0.121.0,<0.122 ; extra == 'fastapi'
 Requires-Dist: starlette>=0.49.1 ; extra == 'fastapi'
 Requires-Dist: uvicorn>=0.37.0,<0.38 ; extra == 'fastapi'

{unique_toolkit-1.52.0 → unique_toolkit-1.53.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "1.52.0"
+version = "1.53.1"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12"
@@ -41,6 +41,7 @@ dependencies = [
     "click>=8.1.0,<9",
     "deepdiff>=8.6.1,<9",
     "pypandoc>=1.16.2,<2",
+    "tenacity>=9.1.4",
 ]
 
 [project.optional-dependencies]

{unique_toolkit-1.52.0 → unique_toolkit-1.53.1}/unique_toolkit/_common/pydantic/rjsf_tags.py

@@ -13,6 +13,7 @@ Key components:
 
 from __future__ import annotations
 
+import types
 from collections.abc import Callable
 from enum import StrEnum
 from typing import Annotated, Any, Union, get_args, get_origin
@@ -20,6 +21,8 @@ from typing import Annotated, Any, Union, get_args, get_origin
 from pydantic import BaseModel
 from typing_extensions import get_type_hints
 
+_UNION_ORIGINS = {Union, types.UnionType}
+
 
 class CustomWidgetName(StrEnum):
     """Mirrors TypeScript CustomWidgetName. Keep in sync when extending."""
@@ -774,7 +777,7 @@ def _unwrap_optional(ann: Any) -> Any:
         The unwrapped type if it was Optional/Union with None, otherwise
         the original type unchanged
     """
-    if get_origin(ann) is Union:
+    if get_origin(ann) in _UNION_ORIGINS:
         args = [a for a in get_args(ann) if a not in _NONE_TYPES]
         if len(args) == 1:
             return args[0]
@@ -921,7 +924,7 @@ def ui_schema_for_model(
             node["additionalProperties"] = val_node
 
         # Union -> anyOf branches
-        elif origin is Union:
+        elif origin in _UNION_ORIGINS:
             branches = []
             for alt in get_args(base):
                 if alt in _NONE_TYPES:

{unique_toolkit-1.52.0 → unique_toolkit-1.53.1}/unique_toolkit/chat/responses_api.py

@@ -1,5 +1,7 @@
 import json
 import logging
+import random
+from collections.abc import Awaitable, Callable
 from typing import Any, NamedTuple, Sequence
 
 import unique_sdk
@@ -13,7 +15,14 @@ from openai.types.responses import (
     response_create_params,
 )
 from openai.types.shared_params import Metadata, Reasoning
-from pydantic import BaseModel, TypeAdapter, ValidationError
+from pydantic import BaseModel, Field, TypeAdapter, ValidationError
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from tenacity import (
+    AsyncRetrying,
+    RetryCallState,
+    retry_if_exception,
+    stop_after_attempt,
+)
 
 from unique_toolkit._common.execution import (
     failsafe,
@@ -329,6 +338,114 @@ def _prepare_responses_args(
     return options
 
 
+_RATE_LIMIT_KEYWORDS = ("too_many_requests", "too many requests", "ratelimitreached")
+
+
+class RateLimitRetryConfig(BaseSettings):
+    """Config for Responses API rate-limit retry. Set via env vars with prefix RATE_LIMIT_RETRY_."""
+
+    initial_delay_seconds: float = Field(
+        default=30.0,
+        description="First wait after SDK retries exhausted. Backoff then 2x (30→60→120s).",
+    )
+    max_attempts: int = Field(
+        default=2,
+        ge=1,
+        le=10,
+        description="Total attempts (1 initial + retries). Default 2 = 1 retry (~30s wait). Set to 1 to disable.",
+    )
+    log_message_on_retry: bool = Field(
+        default=True,
+        description="Write a message-log entry when a retry is about to sleep (e.g. 'Retrying in 30s').",
+    )
+
+    model_config = SettingsConfigDict(
+        env_prefix="RATE_LIMIT_RETRY_",
+        env_file=".env",
+        env_file_encoding="utf-8",
+        case_sensitive=False,
+        extra="ignore",
+    )
+
+
+rate_limit_retry_config = RateLimitRetryConfig()
+
+
+def _rate_limit_wait(retry_state: RetryCallState) -> float:
+    """Exponential backoff with up to 10% jitter: initial_delay → 2x → 4x."""
+    n = retry_state.attempt_number - 1  # 0-indexed
+    base = rate_limit_retry_config.initial_delay_seconds * (2.0**n)
+    return base + random.uniform(0.0, base * 0.1)
+
+
+async def _responses_stream_with_rate_limit_retry(
+    responses_args: dict[str, Any],
+    model_name: str,
+    on_rate_limit_retry: Callable[[int, float], Awaitable[None]] | None = None,
+) -> "unique_sdk.Integrated.ResponsesStreamResult":
+    """Wrap responses_stream_async with exponential backoff for rate-limit errors.
+
+    The SDK already retries 3× (1s/2s/4s) before raising. This layer adds further
+    retries with longer delays, which are needed for models with tighter RPM limits
+    (e.g. GPT-4o) when large-token requests (web search, code execution) are in the loop.
+
+    Args:
+        responses_args: Keyword arguments forwarded to responses_stream_async.
+        model_name: Model name used in log messages.
+        on_rate_limit_retry: Optional async callback invoked before each retry sleep.
+            Receives (attempt_1based, wait_seconds). Useful for writing a message-log
+            entry so the user sees progress during long waits.
+    """
+
+    def _is_rate_limit(exc: BaseException) -> bool:
+        return isinstance(exc, unique_sdk.APIError) and any(
+            kw in str(exc).lower() for kw in _RATE_LIMIT_KEYWORDS
+        )
+
+    max_attempts = rate_limit_retry_config.max_attempts
+    max_retries = max_attempts - 1
+
+    def _log_attempt(retry_state: RetryCallState) -> None:
+        exc = retry_state.outcome.exception() if retry_state.outcome else None
+        if exc is None:
+            return
+        logger.error(
+            "responses_stream_async error on toolkit attempt %d/%d "
+            "model=%s code=%s http_status=%s is_rate_limit=%s original_error=%r",
+            retry_state.attempt_number,
+            max_attempts,
+            model_name,
+            getattr(exc, "code", None),
+            getattr(exc, "http_status", None),
+            _is_rate_limit(exc),
+            getattr(exc, "original_error", None),
+        )
+
+    async def _before_sleep(retry_state: RetryCallState) -> None:
+        wait_secs = retry_state.upcoming_sleep
+        logger.warning(
+            "Rate limit hit for model=%s. Toolkit-level retry %d/%d in %.1fs",
+            model_name,
+            retry_state.attempt_number,
+            max_retries,
+            wait_secs,
+        )
+        if on_rate_limit_retry is not None:
+            await on_rate_limit_retry(retry_state.attempt_number, wait_secs)
+
+    async def _call() -> "unique_sdk.Integrated.ResponsesStreamResult":
+        return await unique_sdk.Integrated.responses_stream_async(**responses_args)
+
+    return await AsyncRetrying(
+        retry=retry_if_exception(_is_rate_limit),
+        stop=stop_after_attempt(max_attempts),
+        wait=_rate_limit_wait,
+        after=_log_attempt,
+        before_sleep=_before_sleep,
+        reraise=True,
+    )(_call)
+
+
 def stream_responses_with_references(
     *,
     company_id: str,
@@ -426,6 +543,7 @@ async def stream_responses_with_references_async(
     top_p: float | None = None,
     reasoning: Reasoning | None = None,
     other_options: dict | None = None,
+    on_rate_limit_retry: Callable[[int, float], Awaitable[None]] | None = None,
 ) -> ResponsesLanguageModelStreamResponse:
     responses_params = _prepare_responses_params_util(
         model_name=model_name,
@@ -458,8 +576,14 @@ async def stream_responses_with_references_async(
         other_options=other_options,
     )
 
+    logger.info(
+        "Calling responses_stream_async model=%s",
+        responses_params.model_name,
+    )
     return ResponsesLanguageModelStreamResponse.model_validate(
-        await unique_sdk.Integrated.responses_stream_async(
-            **responses_args,
+        await _responses_stream_with_rate_limit_retry(
+            responses_args=responses_args,
+            model_name=responses_params.model_name,
+            on_rate_limit_retry=on_rate_limit_retry,
         )
     )

{unique_toolkit-1.52.0 → unique_toolkit-1.53.1}/unique_toolkit/services/chat_service.py

@@ -17,6 +17,7 @@ from pydantic import BaseModel
 from typing_extensions import deprecated
 
 from unique_toolkit._common.utils.files import is_file_content, is_image_content
+from unique_toolkit.agentic.feature_flags import feature_flags
 from unique_toolkit.chat.cancellation import CancellationWatcher
 from unique_toolkit.chat.constants import (
     DEFAULT_MAX_MESSAGES,
@@ -50,6 +51,7 @@ from unique_toolkit.chat.functions import (
     update_message_log_async,
 )
 from unique_toolkit.chat.responses_api import (
+    rate_limit_retry_config,
     stream_responses_with_references,
     stream_responses_with_references_async,
 )
@@ -1637,6 +1639,33 @@ class ChatService(ChatServiceDeprecated):
         reasoning: Reasoning | None = None,
         other_options: dict | None = None,
     ) -> ResponsesLanguageModelStreamResponse:
+        async def _on_rate_limit_retry(attempt: int, wait_secs: float) -> None:
+            if not rate_limit_retry_config.log_message_on_retry:
+                return
+            if not feature_flags.enable_new_answers_ui_un_14411.is_enabled(
+                self._company_id
+            ):
+                return
+            from unique_toolkit.agentic.message_log_manager.service import (
+                _request_counters,
+            )
+
+            msg_id = self._assistant_message_id
+            _request_counters[msg_id] += 1
+            order = _request_counters[msg_id]
+            try:
+                await self.create_message_log_async(
+                    message_id=msg_id,
+                    text=f"Rate limit reached; retrying in {wait_secs:.0f}s (attempt {attempt})",
+                    status=MessageLogStatus.RUNNING,
+                    order=order,
+                )
+            except Exception:
+                logger.warning(
+                    "Failed to write rate-limit retry message log",
+                    exc_info=True,
+                )
+
         return await stream_responses_with_references_async(
             company_id=self._company_id,
             user_id=self._user_id,
@@ -1661,6 +1690,7 @@ class ChatService(ChatServiceDeprecated):
             top_p=top_p,
             reasoning=reasoning,
             other_options=other_options,
+            on_rate_limit_retry=_on_rate_limit_retry,
         )
 
     # Chat Content Methods