unique_toolkit 1.51.0__tar.gz → 1.53.0__tar.gz

{unique_toolkit-1.51.0 → unique_toolkit-1.53.0}/PKG-INFO

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

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "1.51.0"
+version = "1.53.0"
 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.51.0 → unique_toolkit-1.53.0}/unique_toolkit/agentic/loop_runner/__init__.py

@@ -18,9 +18,9 @@ from unique_toolkit.agentic.loop_runner.runners import (
     QWEN_MAX_LOOP_ITERATIONS,
     BasicLoopIterationRunner,
     BasicLoopIterationRunnerConfig,
+    MistralLoopIterationRunner,
     QwenLoopIterationRunner,
     ResponsesBasicLoopIterationRunner,
-    is_qwen_model,
 )
 
 __all__ = [
@@ -28,10 +28,10 @@ __all__ = [
     "PlanningConfig",
     "PlanningMiddleware",
     "PlanningSchemaConfig",
-    "QwenLoopIterationRunner",
-    "is_qwen_model",
     "BasicLoopIterationRunnerConfig",
     "BasicLoopIterationRunner",
+    "MistralLoopIterationRunner",
+    "QwenLoopIterationRunner",
     "handle_forced_tools_iteration",
     "handle_last_iteration",
     "handle_normal_iteration",

{unique_toolkit-1.51.0 → unique_toolkit-1.53.0}/unique_toolkit/agentic/loop_runner/_iteration_handler_utils.py

@@ -47,13 +47,17 @@ async def run_forced_tools_iteration(
     *,
     loop_runner_kwargs: _LoopIterationRunnerKwargs,
     prepare_loop_runner_kwargs: PrepareForcedToolIterationKwargs | None = None,
+    tool_choice_override: str | None = None,
 ) -> LanguageModelStreamResponse:
     """
     Execute a "forced tools" iteration by running one stream_response per tool choice,
     then merging tool calls and references into a single response.
 
-    Some models (e.g. Qwen) need per-tool-choice message rewriting; this can be done
-    via prepare_loop_runner_kwargs (called once per tool choice, on a copy of kwargs).
+    Args:
+        prepare_loop_runner_kwargs: Optional callback to transform kwargs per tool choice
+            (e.g. Qwen rewrites messages with tool-specific instructions).
+        tool_choice_override: If set, replaces the per-tool named dict with this value
+            (e.g. Mistral needs "any" instead of the named format).
     """
     assert "tool_choices" in loop_runner_kwargs
 
@@ -74,13 +78,16 @@ async def run_forced_tools_iteration(
         if prepare_loop_runner_kwargs:
             per_choice_kwargs = prepare_loop_runner_kwargs(func_name, per_choice_kwargs)
 
+        effective_tool_choice = tool_choice_override or opt
         limited_tool = available_tools.get(func_name) if func_name else None
-        stream_kwargs = {"loop_runner_kwargs": per_choice_kwargs, "tool_choice": opt}
+        stream_kwargs = {
+            "loop_runner_kwargs": per_choice_kwargs,
+            "tool_choice": effective_tool_choice,
+        }
         if limited_tool:
             stream_kwargs["tools"] = [limited_tool]
         responses.append(await stream_response(**stream_kwargs))
 
-    # Merge responses and refs:
     tool_calls = []
     references = []
     for r in responses:

{unique_toolkit-1.51.0 → unique_toolkit-1.53.0}/unique_toolkit/agentic/loop_runner/_stream_handler_utils.py

@@ -1,6 +1,6 @@
 from typing import Any, Required
 
-from openai.types.chat import ChatCompletionNamedToolChoiceParam
+from openai.types.chat import ChatCompletionToolChoiceOptionParam
 from typing_extensions import TypedDict
 
 from unique_toolkit import LanguageModelToolDescription
@@ -18,7 +18,7 @@ class _StreamingHandlerKwargs(TypedDict, total=False):
     start_text: str
     debug_info: dict[str, Any]
     temperature: float
-    tool_choice: ChatCompletionNamedToolChoiceParam
+    tool_choice: ChatCompletionToolChoiceOptionParam
     other_options: dict[str, Any]
 
 

{unique_toolkit-1.51.0 → unique_toolkit-1.53.0}/unique_toolkit/agentic/loop_runner/runners/__init__.py

@@ -3,21 +3,23 @@ from unique_toolkit.agentic.loop_runner.runners.basic import (
     BasicLoopIterationRunnerConfig,
     ResponsesBasicLoopIterationRunner,
 )
+from unique_toolkit.agentic.loop_runner.runners.mistral import (
+    MistralLoopIterationRunner,
+)
 from unique_toolkit.agentic.loop_runner.runners.qwen import (
     QWEN_FORCED_TOOL_CALL_INSTRUCTION,
     QWEN_LAST_ITERATION_INSTRUCTION,
     QWEN_MAX_LOOP_ITERATIONS,
     QwenLoopIterationRunner,
-    is_qwen_model,
 )
 
 __all__ = [
     "BasicLoopIterationRunnerConfig",
     "BasicLoopIterationRunner",
     "ResponsesBasicLoopIterationRunner",
+    "MistralLoopIterationRunner",
     "QwenLoopIterationRunner",
     "QWEN_FORCED_TOOL_CALL_INSTRUCTION",
     "QWEN_LAST_ITERATION_INSTRUCTION",
     "QWEN_MAX_LOOP_ITERATIONS",
-    "is_qwen_model",
 ]

{unique_toolkit-1.51.0 → unique_toolkit-1.53.0}/unique_toolkit/agentic/loop_runner/runners/basic.py

@@ -5,9 +5,9 @@ from pydantic import BaseModel
 
 from unique_toolkit._common.pydantic_helpers import get_configuration_dict
 from unique_toolkit.agentic.loop_runner._iteration_handler_utils import (
-    handle_forced_tools_iteration,
     handle_last_iteration,
     handle_normal_iteration,
+    run_forced_tools_iteration,
 )
 from unique_toolkit.agentic.loop_runner._responses_iteration_handler_utils import (
     handle_responses_forced_tools_iteration,
@@ -44,11 +44,29 @@ class BasicLoopIterationRunner(LoopIterationRunner):
         iteration_index = kwargs["iteration_index"]
 
         if len(tool_choices) > 0 and iteration_index == 0:
-            return await handle_forced_tools_iteration(**kwargs)
+            return await self._handle_forced_tools(**kwargs)
         elif iteration_index == self._config.max_loop_iterations - 1:
-            return await handle_last_iteration(**kwargs)
+            return await self._handle_last_iteration(**kwargs)
         else:
-            return await handle_normal_iteration(**kwargs)
+            return await self._handle_normal_iteration(**kwargs)
+
+    async def _handle_forced_tools(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse:
+        return await run_forced_tools_iteration(loop_runner_kwargs=kwargs)
+
+    async def _handle_last_iteration(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse:
+        return await handle_last_iteration(**kwargs)
+
+    async def _handle_normal_iteration(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse:
+        return await handle_normal_iteration(**kwargs)
 
 
 class ResponsesBasicLoopIterationRunner(ResponsesLoopIterationRunner):

unique_toolkit-1.53.0/unique_toolkit/agentic/loop_runner/runners/mistral/__init__.py

@@ -0,0 +1,7 @@
+from unique_toolkit.agentic.loop_runner.runners.mistral.mistral_runner import (
+    MistralLoopIterationRunner,
+)
+
+__all__ = [
+    "MistralLoopIterationRunner",
+]

unique_toolkit-1.53.0/unique_toolkit/agentic/loop_runner/runners/mistral/mistral_runner.py

@@ -0,0 +1,30 @@
+from typing import Unpack
+
+from unique_toolkit.agentic.loop_runner._iteration_handler_utils import (
+    run_forced_tools_iteration,
+)
+from unique_toolkit.agentic.loop_runner.base import _LoopIterationRunnerKwargs
+from unique_toolkit.agentic.loop_runner.runners.basic import (
+    BasicLoopIterationRunner,
+)
+from unique_toolkit.chat.functions import LanguageModelStreamResponse
+
+
+class MistralLoopIterationRunner(BasicLoopIterationRunner):
+    """Runner for Mistral models.
+
+    Mistral doesn't support named tool_choice (``{"type": "function", ...}``).
+    This runner overrides forced-tool handling to send ``"any"`` instead,
+    which forces the model to call a tool. Since each forced iteration already
+    restricts the tools list to a single tool, ``"any"`` effectively forces
+    the correct one.
+    """
+
+    async def _handle_forced_tools(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse:
+        return await run_forced_tools_iteration(
+            loop_runner_kwargs=kwargs,
+            tool_choice_override="any",
+        )

{unique_toolkit-1.51.0 → unique_toolkit-1.53.0}/unique_toolkit/agentic/loop_runner/runners/qwen/__init__.py

@@ -1,6 +1,3 @@
-from unique_toolkit.agentic.loop_runner.runners.qwen.helpers import (
-    is_qwen_model,
-)
 from unique_toolkit.agentic.loop_runner.runners.qwen.qwen_runner import (
     QWEN_FORCED_TOOL_CALL_INSTRUCTION,
     QWEN_LAST_ITERATION_INSTRUCTION,
@@ -10,7 +7,6 @@ from unique_toolkit.agentic.loop_runner.runners.qwen.qwen_runner import (
 
 __all__ = [
     "QwenLoopIterationRunner",
-    "is_qwen_model",
     "QWEN_FORCED_TOOL_CALL_INSTRUCTION",
     "QWEN_LAST_ITERATION_INSTRUCTION",
     "QWEN_MAX_LOOP_ITERATIONS",

{unique_toolkit-1.51.0 → unique_toolkit-1.53.0}/unique_toolkit/agentic/loop_runner/runners/qwen/helpers.py

@@ -1,4 +1,3 @@
-from unique_toolkit.language_model.infos import LanguageModelInfo
 from unique_toolkit.language_model.schemas import (
     LanguageModelAssistantMessage,
     LanguageModelMessageRole,
@@ -6,17 +5,6 @@ from unique_toolkit.language_model.schemas import (
 )
 
 
-def is_qwen_model(*, model: str | LanguageModelInfo | None) -> bool:
-    """Check if the model is a Qwen model."""
-    if isinstance(model, LanguageModelInfo):
-        name = model.name
-        # name is an Enum with a .value attribute
-        return "qwen" in str(getattr(name, "value", name)).lower()
-    elif isinstance(model, str):
-        return "qwen" in model.lower()
-    return False
-
-
 def append_qwen_forced_tool_call_instruction(
     *,
     messages: LanguageModelMessages,

{unique_toolkit-1.51.0 → unique_toolkit-1.53.0}/unique_toolkit/agentic/loop_runner/runners/qwen/qwen_runner.py

@@ -2,13 +2,12 @@ import logging
 from typing import Unpack
 
 from unique_toolkit.agentic.loop_runner._iteration_handler_utils import (
-    handle_last_iteration,
-    handle_normal_iteration,
     run_forced_tools_iteration,
 )
-from unique_toolkit.agentic.loop_runner.base import (
-    LoopIterationRunner,
-    _LoopIterationRunnerKwargs,
+from unique_toolkit.agentic.loop_runner.base import _LoopIterationRunnerKwargs
+from unique_toolkit.agentic.loop_runner.runners.basic import (
+    BasicLoopIterationRunner,
+    BasicLoopIterationRunnerConfig,
 )
 from unique_toolkit.agentic.loop_runner.runners.qwen.helpers import (
     append_qwen_forced_tool_call_instruction,
@@ -30,45 +29,30 @@ QWEN_LAST_ITERATION_INSTRUCTION = """The maximum number of loop iteration have b
 Based on the found information, an answer should be generated"""
 
 
-class QwenLoopIterationRunner(LoopIterationRunner):
+class QwenLoopIterationRunner(BasicLoopIterationRunner):
     def __init__(
         self,
         *,
-        qwen_forced_tool_call_instruction: str,
-        qwen_last_iteration_instruction: str,
-        max_loop_iterations: int,
+        config: BasicLoopIterationRunnerConfig,
+        forced_tool_call_instruction: str,
+        last_iteration_instruction: str,
         chat_service: ChatService,
     ) -> None:
-        self._qwen_forced_tool_call_instruction = qwen_forced_tool_call_instruction
-        self._qwen_last_iteration_instruction = qwen_last_iteration_instruction
-        self._max_loop_iterations = max_loop_iterations
+        super().__init__(config)
+        self._forced_tool_call_instruction = forced_tool_call_instruction
+        self._last_iteration_instruction = last_iteration_instruction
         self._chat_service = chat_service
 
-    async def __call__(
+    async def _handle_forced_tools(
         self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
     ) -> LanguageModelStreamResponse:
-        tool_choices = kwargs.get("tool_choices") or []
-        iteration_index = kwargs["iteration_index"]
-
-        if len(tool_choices) > 0 and iteration_index == 0:
-            return await self._qwen_handle_forced_tools_iteration(**kwargs)
-        elif iteration_index == self._max_loop_iterations - 1:
-            return await self._qwen_handle_last_iteration(**kwargs)
-        else:
-            return await self._qwen_handle_normal_iteration(**kwargs)
-
-    async def _qwen_handle_forced_tools_iteration(
-        self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
-    ) -> LanguageModelStreamResponse:
-        # For Qwen models, append tool call instruction to the last user message. These models ignore the parameter tool_choice.
-        # As the message has to be modified for each tool call instruction, the function from the basic runner cant be used.
         original_messages = kwargs["messages"].model_copy(deep=True)
 
         def _prepare(
             func_name: str | None,
             per_choice_kwargs: _LoopIterationRunnerKwargs,
         ) -> _LoopIterationRunnerKwargs:
-            prompt_instruction = self._qwen_forced_tool_call_instruction.format(
+            prompt_instruction = self._forced_tool_call_instruction.format(
                 TOOL_NAME=func_name or ""
             )
             per_choice_kwargs["messages"] = append_qwen_forced_tool_call_instruction(
@@ -83,34 +67,28 @@ class QwenLoopIterationRunner(LoopIterationRunner):
         )
         return self._process_response(response)
 
-    async def _qwen_handle_last_iteration(
+    async def _handle_last_iteration(
         self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
     ) -> LanguageModelStreamResponse:
-        # For Qwen models, append an assistant message with instructions to not call any tool in this iteration.
         _LOGGER.info(
             "Reached last iteration, removing tools. Appending assistant message with instructions to not call any tool in this iteration."
         )
         kwargs["messages"] = append_qwen_last_iteration_assistant_message(
             messages=kwargs["messages"],
-            last_iteration_instruction=self._qwen_last_iteration_instruction,
+            last_iteration_instruction=self._last_iteration_instruction,
         )
-
-        response = await handle_last_iteration(
-            **kwargs,
-        )
-
+        response = await super()._handle_last_iteration(**kwargs)
         return self._process_response(response)
 
-    async def _qwen_handle_normal_iteration(
+    async def _handle_normal_iteration(
         self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
     ) -> LanguageModelStreamResponse:
-        response = await handle_normal_iteration(**kwargs)
+        response = await super()._handle_normal_iteration(**kwargs)
         return self._process_response(response)
 
     def _process_response(
         self, response: LanguageModelStreamResponse
     ) -> LanguageModelStreamResponse:
-        # Check if content of response is </tool_call>
         if "</tool_call>" == response.message.text.strip():
             _LOGGER.warning(
                 "Response contains only <tool_call>. This is not allowed. Returning empty response."

{unique_toolkit-1.51.0 → unique_toolkit-1.53.0}/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.51.0 → unique_toolkit-1.53.0}/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