unique_toolkit 2026.44.0.dev4__tar.gz → 2026.44.0.dev6__tar.gz

{unique_toolkit-2026.44.0.dev4 → unique_toolkit-2026.44.0.dev6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 2026.44.0.dev4
+Version: 2026.44.0.dev6
 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>
@@ -16,7 +16,7 @@ Requires-Dist: sseclient>=0.0.27,<0.0.28
 Requires-Dist: openai>=1.105.0,<3
 Requires-Dist: platformdirs>=4.0.0,<5
 Requires-Dist: pillow>=12.1.1
-Requires-Dist: unique-sdk>=2026.44.0.dev7,<2026.44.0rc0
+Requires-Dist: unique-sdk>=2026.44.0.dev8,<2026.44.0rc0
 Requires-Dist: jambo>=0.1.2,<0.2
 Requires-Dist: jinja2>=3.1.6,<4
 Requires-Dist: docxtpl>=0.20.1,<0.21

{unique_toolkit-2026.44.0.dev4 → unique_toolkit-2026.44.0.dev6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "2026.44.0.dev4"
+version = "2026.44.0.dev6"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12"
@@ -28,7 +28,7 @@ dependencies = [
     "openai>=1.105.0,<3",
     "platformdirs>=4.0.0,<5",
     "pillow>=12.1.1",
-    "unique-sdk>=2026.44.0.dev7,<2026.44.0rc0",
+    "unique-sdk>=2026.44.0.dev8,<2026.44.0rc0",
     "jambo>=0.1.2,<0.2",
     "jinja2>=3.1.6,<4",
     "docxtpl>=0.20.1,<0.21",

unique_toolkit-2026.44.0.dev6/unique_toolkit/agentic/history_manager/context_overflow.py

@@ -0,0 +1,259 @@
+"""Keep agent-loop prompts inside the provider's context window, verified by the provider's own token counter."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections.abc import Awaitable, Callable
+from logging import Logger
+from typing import Any, TypeVar
+
+from unique_toolkit.agentic.history_manager.history_manager import HistoryManager
+from unique_toolkit.agentic.loop_runner.base import _LoopIterationRunnerKwargs
+from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.language_model.settings import ModelFamily
+from unique_toolkit.protocols.support import SupportsTokenCount
+
+_T = TypeVar("_T")
+
+_CONTEXT_OVERFLOW_SIGNALS = (
+    "context_length_exceeded",
+    "contextwindowexceedederror",
+    "tokenlimitexceeded",
+    "maximum context length",
+    "prompt is too long",
+    "input is too long",
+    "too many input tokens",
+    "exceeds the maximum number of tokens",
+)
+_CONTEXT_WINDOW_QUALIFIERS = ("exceed", "too long", "maximum")
+
+CONTEXT_OVERFLOW_TOKEN_BUDGET_FACTOR = 0.7
+PROVIDER_COUNT_TARGET_RATIO = 0.95
+MAX_PROVIDER_OVERFLOW_REDUCTION_FACTOR = 0.9
+APPROXIMATE_TOKENIZER_TRIGGER_RATIO = 0.7
+NATIVE_TOKENIZER_TRIGGER_RATIO = 0.85
+LOCAL_LIMIT_TRIGGER_RATIO = 0.95
+DEFAULT_OUTPUT_TOKEN_RESERVE = 10_000
+PROVIDER_TOKEN_COUNT_TIMEOUT_SECONDS = 10
+
+_NATIVE_TOKENIZER_FAMILIES = {
+    ModelFamily.OPENAI,
+    ModelFamily.DEEPSEEK,
+    ModelFamily.QWEN,
+}
+
+WITHHELD_TOOL_RESULT_NOTICE = (
+    "The combined output of the tool calls made in this step was too large for "
+    "the model's context window, so this result was withheld. Do not repeat "
+    "these tool calls in parallel. Call them one at a time instead, and narrow "
+    "each call (for example a lower result limit, a shorter date range or a "
+    "more specific query) so that each result stays small."
+)
+
+
+def is_context_window_exceeded_error(exc: BaseException) -> bool:
+    """Return True if *exc* indicates the prompt exceeded the model's context window."""
+    error_text = f"{type(exc).__name__} {exc}".lower()
+    if any(signal in error_text for signal in _CONTEXT_OVERFLOW_SIGNALS):
+        return True
+    return "context window" in error_text and any(
+        qualifier in error_text for qualifier in _CONTEXT_WINDOW_QUALIFIERS
+    )
+
+
+def supports_provider_token_count(model: LanguageModelInfo) -> bool:
+    return model.family == ModelFamily.ANTHROPIC
+
+
+def provider_context_window(model: LanguageModelInfo) -> int:
+    return model.provider_context_window
+
+
+def provider_count_trigger_tokens(
+    model: LanguageModelInfo,
+    effective_local_limit: int,
+) -> int:
+    """Local count at which an authoritative provider count becomes worthwhile."""
+    ratio = (
+        NATIVE_TOKENIZER_TRIGGER_RATIO
+        if model.family in _NATIVE_TOKENIZER_FAMILIES
+        else APPROXIMATE_TOKENIZER_TRIGGER_RATIO
+    )
+    return min(
+        int(provider_context_window(model) * ratio),
+        int(effective_local_limit * LOCAL_LIMIT_TRIGGER_RATIO),
+    )
+
+
+def provider_input_target(
+    model: LanguageModelInfo,
+    other_options: dict[str, Any] | None,
+) -> int:
+    """Maximum provider-counted input while preserving output and safety room."""
+    options = other_options or {}
+    configured_output = options.get("max_completion_tokens", options.get("max_tokens"))
+    output_reserve = (
+        configured_output
+        if isinstance(configured_output, int) and configured_output >= 0
+        else min(
+            DEFAULT_OUTPUT_TOKEN_RESERVE,
+            model.token_limits.token_limit_output,
+        )
+    )
+    return max(
+        1,
+        int(provider_context_window(model) * PROVIDER_COUNT_TARGET_RATIO)
+        - output_reserve,
+    )
+
+
+def provider_overflow_reduction_factor(provider_count: int, target: int) -> float:
+    """Budget multiplier that brings an oversized provider count under *target*."""
+    return min(
+        MAX_PROVIDER_OVERFLOW_REDUCTION_FACTOR,
+        target / provider_count * PROVIDER_COUNT_TARGET_RATIO,
+    )
+
+
+class ContextOverflowGuard:
+    """Preflight provider token-budget verification and overflow retry policy."""
+
+    def __init__(
+        self,
+        *,
+        logger: Logger,
+        history_manager: HistoryManager,
+        token_counter: object,
+        model: LanguageModelInfo,
+    ) -> None:
+        self._logger = logger
+        self._history_manager = history_manager
+        self._token_counter = token_counter
+        self._model = model
+
+    async def preflight_token_budget(
+        self,
+        kwargs: _LoopIterationRunnerKwargs,
+        recompose: Callable[[], Awaitable[None]],
+    ) -> None:
+        """Verify near-limit payloads using the provider's tokenizer before the LLM call."""
+        if not supports_provider_token_count(self._model):
+            return
+        if not isinstance(self._token_counter, SupportsTokenCount):
+            return
+
+        local_count = self._history_manager.last_token_count
+        effective_limit = self._history_manager.effective_token_limit
+        trigger = provider_count_trigger_tokens(self._model, effective_limit)
+        if local_count < trigger:
+            return
+
+        target = provider_input_target(self._model, kwargs.get("other_options"))
+        provider_count = await self._try_provider_token_count(kwargs)
+        if provider_count is None or provider_count <= target:
+            return
+
+        self._logger.warning(
+            f"Provider token count {provider_count} exceeds the safe input target "
+            f"{target}. Reducing tool results before the LLM call."
+        )
+        reduction_factor = provider_overflow_reduction_factor(provider_count, target)
+        self._history_manager.tighten_token_budget(reduction_factor)
+        await recompose()
+
+        reduced_count = await self._try_provider_token_count(kwargs)
+        if reduced_count is None or reduced_count <= target:
+            return
+
+        if not self._history_manager.withhold_latest_tool_call_results(
+            WITHHELD_TOOL_RESULT_NOTICE
+        ):
+            self._logger.warning(
+                f"Provider token count remains above target ({reduced_count} > "
+                f"{target}), but there are no tool results to withhold."
+            )
+            return
+
+        self._logger.warning(
+            "Provider token count remains above target after reduction. "
+            "Withholding the latest tool-call batch."
+        )
+        await recompose()
+
+        withheld_count = await self._try_provider_token_count(kwargs)
+        if withheld_count is None or withheld_count <= target:
+            return
+
+        self._logger.warning(
+            f"Provider token count {withheld_count} still exceeds the safe input "
+            f"target {target} after withholding. Reducing history further."
+        )
+        self._history_manager.tighten_token_budget(
+            provider_overflow_reduction_factor(withheld_count, target)
+        )
+        await recompose()
+
+    async def retry_after_overflow(
+        self,
+        exc: Exception,
+        *,
+        recompose: Callable[[], Awaitable[None]],
+        run_iteration: Callable[[], Awaitable[_T]],
+    ) -> _T:
+        """Retry an oversized LLM call by tightening the budget or withholding the latest tool batch."""
+        self._logger.warning(
+            "LLM call exceeded the context window. Retrying with a tighter token budget.",
+            exc_info=exc,
+        )
+        self._history_manager.tighten_token_budget(CONTEXT_OVERFLOW_TOKEN_BUDGET_FACTOR)
+        await recompose()
+        try:
+            return await run_iteration()
+        except Exception as retry_exc:
+            if not is_context_window_exceeded_error(retry_exc):
+                raise
+            if not self._history_manager.withhold_latest_tool_call_results(
+                WITHHELD_TOOL_RESULT_NOTICE
+            ):
+                raise
+            self._logger.warning(
+                "LLM call still exceeded the context window. Retrying with the "
+                "latest tool call results withheld.",
+                exc_info=retry_exc,
+            )
+            await recompose()
+            return await run_iteration()
+
+    async def _try_provider_token_count(
+        self, kwargs: _LoopIterationRunnerKwargs
+    ) -> int | None:
+        counter = self._token_counter
+        if not isinstance(counter, SupportsTokenCount):
+            return None
+        started_at = time.perf_counter()
+        try:
+            count = await asyncio.wait_for(
+                counter.count_tokens_async(
+                    messages=kwargs["messages"],
+                    model_name=self._model.name,
+                    temperature=kwargs.get("temperature", 0.0),
+                    tools=kwargs.get("tools"),
+                    other_options=kwargs.get("other_options"),
+                ),
+                timeout=PROVIDER_TOKEN_COUNT_TIMEOUT_SECONDS,
+            )
+        except Exception:
+            elapsed_ms = (time.perf_counter() - started_at) * 1000
+            self._logger.warning(
+                f"Provider token counting failed after {elapsed_ms:.0f} ms. "
+                "Continuing with the local token budget.",
+                exc_info=True,
+            )
+            return None
+        elapsed_ms = (time.perf_counter() - started_at) * 1000
+        self._logger.info(
+            f"Provider token count for {self._model.name}: {count} tokens "
+            f"in {elapsed_ms:.0f} ms."
+        )
+        return count if count > 0 else None

{unique_toolkit-2026.44.0.dev4 → unique_toolkit-2026.44.0.dev6}/unique_toolkit/agentic/history_manager/history_manager.py

@@ -152,6 +152,7 @@ class HistoryManager:
         return len(self._loop_history) == 0
 
     def add_tool_call_results(self, tool_call_results: list[ToolCallResponse]):
+        self._collected_tool_response_image_urls = []
         for tool_response in tool_call_results:
             if tool_response.image_data_urls:
                 for url in tool_response.image_data_urls:
@@ -220,6 +221,83 @@ class HistoryManager:
     def add_assistant_message(self, message: LanguageModelAssistantMessage) -> None:
         self._loop_history.append(message)
 
+    def tighten_token_budget(self, factor: float) -> None:
+        self._token_reducer.tighten_effective_token_limit(factor)
+
+    @property
+    def effective_token_limit(self) -> int:
+        return self._token_reducer.effective_token_limit
+
+    @property
+    def last_token_count(self) -> int:
+        return self._token_reducer.last_token_count
+
+    def withhold_latest_tool_call_results(self, notice: str) -> int:
+        """Replace every result of the most recent tool-call batch with *notice*.
+
+        The withheld results' chunks are dropped from the reference manager. They
+        are the newest chunks, so source numbers of earlier results stay valid.
+        Returns the number of newly withheld results; results that already hold
+        *notice* are not counted.
+        """
+        batch_start: int | None = None
+        for index in range(len(self._loop_history) - 1, -1, -1):
+            candidate = self._loop_history[index]
+            if (
+                isinstance(candidate, LanguageModelAssistantMessage)
+                and candidate.tool_calls
+            ):
+                batch_start = index
+                break
+        if batch_start is None:
+            return 0
+
+        withheld_tool_call_ids: list[str] = []
+        for index in range(batch_start + 1, len(self._loop_history)):
+            message = self._loop_history[index]
+            if not isinstance(message, LanguageModelToolMessage):
+                continue
+            if message.content == notice:
+                continue
+            self._loop_history[index] = LanguageModelToolMessage(
+                content=notice,
+                tool_call_id=message.tool_call_id,
+                name=message.name,
+            )
+            withheld_tool_call_ids.append(message.tool_call_id)
+        if not withheld_tool_call_ids:
+            return 0
+
+        withheld_chunks = {
+            id(chunk)
+            for tool_call_id in withheld_tool_call_ids
+            for chunk in self._reference_manager.get_chunks_of_tool(tool_call_id)
+        }
+        if withheld_chunks:
+            self._reference_manager.replace(
+                chunks=[
+                    chunk
+                    for chunk in self._reference_manager.get_chunks()
+                    if id(chunk) not in withheld_chunks
+                ]
+            )
+            for tool_call_id in withheld_tool_call_ids:
+                self._reference_manager.replace_chunks_of_tool(tool_call_id, [])
+
+        if self._collected_tool_response_image_urls:
+            withheld_ids = set(withheld_tool_call_ids)
+            self._collected_tool_response_image_urls = [
+                (url, tool_call_id)
+                for url, tool_call_id in self._collected_tool_response_image_urls
+                if tool_call_id not in withheld_ids
+            ]
+
+        self._logger.warning(
+            f"Withheld {len(withheld_tool_call_ids)} tool call result(s) that "
+            "exceeded the context window."
+        )
+        return len(withheld_tool_call_ids)
+
     async def get_history_for_model_call(
         self,
         original_user_message: str,
@@ -238,7 +316,6 @@ class HistoryManager:
             remove_from_text=remove_from_text,
             image_data_urls_from_tools=image_data_from_tools,
         )
-        self._collected_tool_response_image_urls = []
 
         if not self._source_offset_initialized:
             offset = max(0, self._token_reducer.max_db_source_number + 1)

{unique_toolkit-2026.44.0.dev4 → unique_toolkit-2026.44.0.dev6}/unique_toolkit/agentic/history_manager/loop_token_reducer.py

@@ -38,7 +38,6 @@ MAX_INPUT_TOKENS_SAFETY_PERCENTAGE = env_settings.input_correction_reduction_fac
 SAFETY_MARGIN_FOR_AGGRESSIVE_REDUCTION = env_settings.input_aggressive_reduction_factor
 
 PLAIN_TEXT_TRUNCATION_MARKER = "\n\n[... tool response too long and was truncated]"
-# TODO: Make this configurable
 MIN_PLAIN_TEXT_CHARS_TO_KEEP = 1000
 
 
@@ -106,6 +105,7 @@ class LoopTokenReducer:
             self._language_model.token_limits.token_limit_input
             * (1 - MAX_INPUT_TOKENS_SAFETY_PERCENTAGE)
         )
+        self._last_token_count = 0
         self._max_db_source_number: int = -1
         self._db_source_map: dict[int, ContentChunk] = {}
         self._enable_tool_call_persistence = enable_tool_call_persistence
@@ -131,6 +131,23 @@ class LoopTokenReducer:
     def db_source_map(self) -> dict[int, ContentChunk]:
         return self._db_source_map
 
+    @property
+    def effective_token_limit(self) -> int:
+        return self._effective_token_limit
+
+    @property
+    def last_token_count(self) -> int:
+        return self._last_token_count
+
+    def tighten_effective_token_limit(self, factor: float) -> None:
+        """Shrink the input-token budget for the rest of the run after the provider counted more tokens than the local encoder."""
+        previous_limit = self._effective_token_limit
+        self._effective_token_limit = int(previous_limit * factor)
+        self._logger.warning(
+            f"Tightened effective token limit from {previous_limit} to "
+            f"{self._effective_token_limit} tokens."
+        )
+
     def _get_encoder(self, language_model: LMI) -> Callable[[str], list[int]]:
         return language_model.get_encoder()
 
@@ -179,6 +196,7 @@ class LoopTokenReducer:
                 break
 
         token_count = self._count_message_tokens(messages)
+        self._last_token_count = token_count
         self._logger.info(
             f"Final token count after reduction: {token_count} of model_capacity {self._language_model.token_limits.token_limit_input}"
         )
@@ -191,8 +209,6 @@ class LoopTokenReducer:
         This is purely a size check. Whether anything can actually be reduced to
         get back under the limit is a separate concern (:meth:`_can_reduce_history`).
         """
-        # TODO: This is not fully correct at the moment as the token_count
-        # include system_prompt and user question already
         return token_count > self._effective_token_limit
 
     def _can_reduce_history(self, loop_history: list[LanguageModelMessage]) -> bool:

{unique_toolkit-2026.44.0.dev4 → unique_toolkit-2026.44.0.dev6}/unique_toolkit/language_model/functions.py

@@ -21,6 +21,9 @@ from unique_toolkit.language_model.infos import (
     LanguageModelName,
     ModelCapabilities,
 )
+from unique_toolkit.language_model.provider_token_count import (
+    build_provider_token_count_payload,
+)
 from unique_toolkit.language_model.reference import (
     add_references_to_message,
 )
@@ -748,3 +751,35 @@ async def stream_complete_with_references_openai(
     except Exception as e:
         logger.error("Error streaming completion (model=%s): %s", model_name_, e)
         raise
+
+
+async def count_tokens_with_provider_async(
+    *,
+    company_id: str,
+    user_id: str,
+    messages: LanguageModelMessages | list[ChatCompletionMessageParam],
+    model_name: LanguageModelName | str,
+    temperature: float = DEFAULT_COMPLETE_TEMPERATURE,
+    tools: Sequence[LanguageModelTool | LanguageModelToolDescription] | None = None,
+    other_options: dict[str, Any] | None = None,
+) -> int:
+    """Count the completion payload with LiteLLM's provider-backed counter."""
+    options, model, messages_dict, _ = _prepare_all_completions_params_util(
+        messages=messages,
+        model_name=model_name,
+        temperature=temperature,
+        tools=tools,
+        other_options=other_options,
+    )
+    raw_tools = options.get("tools")
+    payload = build_provider_token_count_payload(
+        model=model,
+        messages=messages_dict,
+        tools=raw_tools if isinstance(raw_tools, list) else None,
+    )
+    response = await unique_sdk.Integrated.count_tokens_async(
+        user_id=user_id,
+        company_id=company_id,
+        **payload,
+    )
+    return response["input_tokens"]

{unique_toolkit-2026.44.0.dev4 → unique_toolkit-2026.44.0.dev6}/unique_toolkit/language_model/infos.py

@@ -335,6 +335,28 @@ def _load_custom_decoder(tokenizer_name: str) -> TypeDecoder:
     return tokenizer.decode
 
 
+_CLAUDE_CONTEXT_WINDOWS = {
+    "claude-3-7-sonnet": 200_000,
+    "claude-3-7-sonnet-thinking": 200_000,
+    "claude-haiku-4-5": 200_000,
+    "claude-opus-4": 200_000,
+    "claude-opus-4-1": 200_000,
+    "claude-opus-4-5": 200_000,
+    "claude-opus-4-6": 200_000,
+    "claude-sonnet-4": 200_000,
+    "claude-sonnet-4-5": 200_000,
+    "claude-sonnet-4-6": 200_000,
+    "claude-fable-5": 1_000_000,
+    "claude-haiku-5-5": 1_000_000,
+    "claude-opus-4-7": 1_000_000,
+    "claude-opus-4-8": 1_000_000,
+    "claude-opus-5": 1_000_000,
+    "claude-opus-5-5": 1_000_000,
+    "claude-sonnet-5": 1_000_000,
+    "claude-sonnet-5-5": 1_000_000,
+}
+
+
 class LanguageModelInfo(BaseModel):
     model_config = get_configuration_dict()
     name: (
@@ -3731,6 +3753,29 @@ class LanguageModelInfo(BaseModel):
         else:
             return self.name
 
+    @property
+    def provider_context_window(self) -> int:
+        """Return the raw provider input context window before local-tokenizer adjustments."""
+        if self.family == ModelFamily.ANTHROPIC:
+            known_window = _CLAUDE_CONTEXT_WINDOWS.get(self.version)
+            if known_window is None:
+                model_name = str(self.name)
+                known_window = next(
+                    (
+                        window
+                        for version, window in sorted(
+                            _CLAUDE_CONTEXT_WINDOWS.items(),
+                            key=lambda item: len(item[0]),
+                            reverse=True,
+                        )
+                        if version in model_name
+                    ),
+                    None,
+                )
+            if known_window is not None:
+                return known_window
+        return self.token_limits.token_limit_input
+
 
 @deprecated(
     """