unique_toolkit 2026.44.0.dev4__py3-none-any.whl → 2026.44.0.dev5__py3-none-any.whl

unique_toolkit/agentic/history_manager/context_overflow.py

@@ -0,0 +1,264 @@
+"""Keep agent-loop prompts inside the provider's context window.
+
+The local reducer estimates tokens with an OpenAI tokenizer, which can undercount
+other providers by 2x or more. ``ContextOverflowGuard`` verifies near-limit
+payloads with the provider's own counter before the LLM call and recovers when the
+provider still rejects a prompt as too long.
+"""
+
+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."
+            )
+            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/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,78 @@ 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 withheld results.
+        """
+        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
+            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)
+
+        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 +311,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/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,27 @@ 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.
+
+        Used after the provider rejected a payload this reducer counted as within
+        budget, i.e. the encoder underestimates the model's real token count.
+        """
+        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 +200,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 +213,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/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/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(
     """

unique_toolkit/language_model/provider_token_count.py

@@ -0,0 +1,170 @@
+"""Build request bodies for LiteLLM's Anthropic-compatible ``count_tokens`` endpoint.
+
+Agent-loop messages are in OpenAI chat format. The endpoint expects the Anthropic
+Messages format, so system messages, tool calls, tool results and images are
+converted before counting.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+
+def build_provider_token_count_payload(
+    *,
+    model: str,
+    messages: Sequence[Mapping[str, Any]],
+    tools: Sequence[Mapping[str, Any]] | None,
+) -> dict[str, Any]:
+    """Convert an OpenAI chat payload to LiteLLM's Anthropic-compatible counter."""
+    system_parts: list[str] = []
+    provider_messages: list[dict[str, Any]] = []
+
+    for message in messages:
+        role = message.get("role")
+        if role == "system":
+            content = message.get("content")
+            if isinstance(content, str) and content:
+                system_parts.append(content)
+            continue
+
+        if role == "assistant":
+            blocks = _assistant_content_blocks(message)
+            _append_message(provider_messages, "assistant", blocks)
+        elif role == "tool":
+            _append_message(
+                provider_messages,
+                "user",
+                [
+                    {
+                        "type": "tool_result",
+                        "tool_use_id": (
+                            message.get("tool_call_id")
+                            or message.get("toolCallId")
+                            or ""
+                        ),
+                        "content": _content_as_text(message.get("content")),
+                    }
+                ],
+            )
+        elif role == "user":
+            _append_message(
+                provider_messages,
+                "user",
+                _user_content_blocks(message.get("content")),
+            )
+
+    payload: dict[str, Any] = {
+        "model": model.removeprefix("litellm:"),
+        "messages": provider_messages,
+    }
+    if system_parts:
+        payload["system"] = "\n\n".join(system_parts)
+    if tools:
+        payload["tools"] = [_anthropic_tool(tool) for tool in tools]
+    return payload
+
+
+def _assistant_content_blocks(
+    message: Mapping[str, Any],
+) -> list[dict[str, Any]]:
+    blocks: list[dict[str, Any]] = []
+    content = message.get("content")
+    if isinstance(content, str) and content:
+        blocks.append({"type": "text", "text": content})
+
+    tool_calls = message.get("tool_calls") or message.get("toolCalls")
+    if isinstance(tool_calls, list):
+        for tool_call in tool_calls:
+            function = tool_call.get("function", {})
+            if not isinstance(function, dict):
+                continue
+            blocks.append(
+                {
+                    "type": "tool_use",
+                    "id": tool_call.get("id", ""),
+                    "name": function.get("name", ""),
+                    "input": _tool_arguments(function.get("arguments")),
+                }
+            )
+    return blocks
+
+
+def _user_content_blocks(content: object) -> list[dict[str, Any]]:
+    if isinstance(content, str):
+        return [{"type": "text", "text": content}]
+    if not isinstance(content, list):
+        return [{"type": "text", "text": ""}]
+
+    blocks: list[dict[str, Any]] = []
+    for part in content:
+        if not isinstance(part, dict):
+            continue
+        if part.get("type") == "text":
+            blocks.append({"type": "text", "text": str(part.get("text", ""))})
+        elif part.get("type") == "image_url":
+            image_url = part.get("image_url") or part.get("imageUrl")
+            if isinstance(image_url, dict):
+                image_block = _base64_image_block(image_url.get("url"))
+                if image_block is not None:
+                    blocks.append(image_block)
+    return blocks or [{"type": "text", "text": ""}]
+
+
+def _base64_image_block(url: object) -> dict[str, Any] | None:
+    if not isinstance(url, str) or not url.startswith("data:") or ";base64," not in url:
+        return None
+    metadata, data = url.split(";base64,", 1)
+    return {
+        "type": "image",
+        "source": {
+            "type": "base64",
+            "media_type": metadata.removeprefix("data:"),
+            "data": data,
+        },
+    }
+
+
+def _append_message(
+    messages: list[dict[str, Any]],
+    role: str,
+    blocks: list[dict[str, Any]],
+) -> None:
+    if messages and messages[-1]["role"] == role:
+        messages[-1]["content"].extend(blocks)
+        return
+    messages.append({"role": role, "content": blocks})
+
+
+def _anthropic_tool(tool: Mapping[str, Any]) -> dict[str, Any]:
+    function = tool.get("function", tool)
+    if not isinstance(function, dict):
+        function = {}
+    return {
+        "name": function.get("name", ""),
+        "description": function.get("description", ""),
+        "input_schema": function.get(
+            "parameters", {"type": "object", "properties": {}}
+        ),
+    }
+
+
+def _tool_arguments(arguments: object) -> dict[str, Any]:
+    if isinstance(arguments, dict):
+        return arguments
+    if isinstance(arguments, str):
+        try:
+            parsed = json.loads(arguments)
+        except json.JSONDecodeError:
+            return {"raw": arguments}
+        if isinstance(parsed, dict):
+            return parsed
+    return {}
+
+
+def _content_as_text(content: object) -> str:
+    if isinstance(content, str):
+        return content
+    return json.dumps(content)

unique_toolkit/protocols/support.py

@@ -1,4 +1,4 @@
-from typing import Any, Protocol, Sequence
+from typing import Any, Protocol, Sequence, runtime_checkable
 
 from openai.types.chat import (
     ChatCompletionMessageParam,
@@ -55,6 +55,19 @@ class SupportsComplete(Protocol):
     ) -> LanguageModelResponse: ...
 
 
+@runtime_checkable
+class SupportsTokenCount(Protocol):
+    async def count_tokens_async(
+        self,
+        *,
+        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: ...
+
+
 class SupportCompleteWithReferences(Protocol):
     def complete_with_references(
         self,

unique_toolkit/services/chat_service.py

@@ -105,6 +105,9 @@ from unique_toolkit.language_model.constants import (
     DEFAULT_COMPLETE_TEMPERATURE,
     DEFAULT_COMPLETE_TIMEOUT,
 )
+from unique_toolkit.language_model.functions import (
+    count_tokens_with_provider_async,
+)
 from unique_toolkit.language_model.infos import (
     LanguageModelName,
 )
@@ -1720,6 +1723,25 @@ class ChatService(ChatServiceDeprecated):
             other_options=other_options,
         )
 
+    async def count_tokens_async(
+        self,
+        *,
+        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:
+        return await count_tokens_with_provider_async(
+            company_id=self._company_id,
+            user_id=self._user_id,
+            messages=messages,
+            model_name=model_name,
+            temperature=temperature,
+            tools=tools,
+            other_options=other_options,
+        )
+
     async def complete_async(
         self,
         messages: LanguageModelMessages | list[ChatCompletionMessageParam],

{unique_toolkit-2026.44.0.dev4.dist-info → unique_toolkit-2026.44.0.dev5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 2026.44.0.dev4
+Version: 2026.44.0.dev5
 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.dist-info → unique_toolkit-2026.44.0.dev5.dist-info}/RECORD

@@ -113,9 +113,10 @@ unique_toolkit/agentic/evaluation/tests/test_prompt_loaders.py,sha256=zBREdlKf5t
 unique_toolkit/agentic/evaluation/utils.py,sha256=HmyPaDV8wdW-_gOjjW-wDaMKgdrsP5-SHP7OqTmGI_A,264
 unique_toolkit/agentic/feature_flags/__init__.py,sha256=fBR9FaLrc9e8sdYDe1LwP9aNJt5aLLtAq3nIphVWkIk,191
 unique_toolkit/agentic/feature_flags/feature_flags.py,sha256=LlFfcgGV4lRQ_mgOGsbZzjKeTXxjCCOl7agCxvJUhvQ,7133
+unique_toolkit/agentic/history_manager/context_overflow.py,sha256=1GrGcf44rKhyT4dMJI7W2l4JvXyPFqoLUy0zp2n4yYw,9677
 unique_toolkit/agentic/history_manager/history_construction_with_contents.py,sha256=bJH5RiLM__fg-i_ycact--9xtV5A321xe_cktT__bzA,26284
-unique_toolkit/agentic/history_manager/history_manager.py,sha256=PqSa87yzNIDRXqjh1KN7moeIpp_jWxj9qEPig4wv790,16062
-unique_toolkit/agentic/history_manager/loop_token_reducer.py,sha256=lgo5zAq6gYLo7FFTNWe5sc6o80vnadKsr-cgEDP6hM8,30754
+unique_toolkit/agentic/history_manager/history_manager.py,sha256=gDrbnOIJlnmh3KvKRtoYKWVAymEgboW61-6GJVWuZ4c,18885
+unique_toolkit/agentic/history_manager/loop_token_reducer.py,sha256=Z9Hr3Lb_2t5qqqRcUEYSoV8tp7UUO-uf4yyro3jCyQE,31458
 unique_toolkit/agentic/history_manager/settings.py,sha256=Lrx3a8yXeU0MXrWlV6RZSHsgYWF9Zrs8LyR96OpqQ_k,1216
 unique_toolkit/agentic/history_manager/utils.py,sha256=DiBip8xtah3I9-G-_Ul3OjlYKunrXR-5RJtzREHWV_0,6579
 unique_toolkit/agentic/loop_runner/__init__.py,sha256=bhwOIloPn2MZyCw60FkHy330twCCRovg3bgye9OjZRE,1496
@@ -396,11 +397,12 @@ unique_toolkit/language_model/_responses_api_utils.py,sha256=w2ZyleOQcWLpOWGJNNb
 unique_toolkit/language_model/builder.py,sha256=KwBzsD6-fvQnB4p2I7M-5Tk6j4tmQ973MYwU1ulXFb4,4526
 unique_toolkit/language_model/constants.py,sha256=B-topqW0r83dkC_25DeQfnPk3n53qzIHUCBS7YJ0-1U,119
 unique_toolkit/language_model/default_language_model.py,sha256=mcyuFs_V6tp52C8xJSHiEHO4kFGj6xed1gvozvCRD7k,1181
-unique_toolkit/language_model/functions.py,sha256=W6JC4LUh-LZ9waDgph--xVmKeQ7Ls6NRAqDx9DwQdeY,26691
-unique_toolkit/language_model/infos.py,sha256=f1xQQ8BvLG9jPsEt2vdSYzxVvaKIkem-Y6T4_n4tT_c,182327
+unique_toolkit/language_model/functions.py,sha256=sjuzKe6xAq6NFwwgIRvgAuU8t6apq1tcdGSkcYCl8EQ,27905
+unique_toolkit/language_model/infos.py,sha256=_fbEwGA9AjBm7TF1FTsihtKzi4SybTwVN7QyHRacMrQ,183899
 unique_toolkit/language_model/invocation_stats.py,sha256=CQbpUz394IsVd8eNTB1gPSMt5f0UXmkN1amNsqtAW6U,5857
 unique_toolkit/language_model/model_costs.py,sha256=ShuDAzMew1fUkJ3tzPgU5Zqjp6ivUZaG97BFh93tEn4,6358
 unique_toolkit/language_model/prompt.py,sha256=JSawaLjQg3VR-E2fK8engFyJnNdk21zaO8pPIodzN4Q,3991
+unique_toolkit/language_model/provider_token_count.py,sha256=nXgyL-figP1Zv88zlHXFLvopZAL3hL8lUvF_lusP98g,5499
 unique_toolkit/language_model/reference.py,sha256=T3wr2jP_uRg5wKVIbKqLbxJlcgxfbZQGq3XdjbQ2kSo,6984
 unique_toolkit/language_model/schemas.py,sha256=o8zGkfoYayO1gL85G-ulV3Akp_BwjgEWGGUb89j1OgU,35440
 unique_toolkit/language_model/service.py,sha256=lgpzWhqxLBOeFacCafQi4e7BqAxJbOlr8GnCYcKeTdM,13274
@@ -413,9 +415,9 @@ unique_toolkit/monitoring/memory.py,sha256=KKnP5NEY63RXR9mmgXOEl5S9fgJSDDr2nE60y
 unique_toolkit/monitoring/middleware.py,sha256=MNsUuKgt74mSnUwttXTB_AxX546Pz2zKoAP9OewVrwQ,4531
 unique_toolkit/monitoring/registry.py,sha256=De2Bx5QBBRJQY2MM3xtntiJ2JU2FoKkVT7rNMG7ewgA,5874
 unique_toolkit/monitoring/tracing.py,sha256=0lmLuQjyyC3iFpNTw8rOyCpEZ6PTvp16nVbzcO9d2sQ,10043
-unique_toolkit/protocols/support.py,sha256=x5edzbBmKAUuvm3oLFh52gkgQsqR5LU7M0t8RHK0DZk,5161
+unique_toolkit/protocols/support.py,sha256=vlXZEC9_Gupk2x9wZln5FJQlM6srWv4hIZY6O5WXhbk,5636
 unique_toolkit/services/__init__.py,sha256=b9Wl7Wvz1JmZDbDexIZzBoXZ0wscXuczGAH8LoL-Qak,572
-unique_toolkit/services/chat_service.py,sha256=o4gjFjtfoAryxpGhBpdnZomvqlgxR_m-U6UYyS72cEA,101002
+unique_toolkit/services/chat_service.py,sha256=UasMB15wQnHp8-yMY9kN1zK8UZPhFPCi4vgkxrqBrEM,101802
 unique_toolkit/services/factory.py,sha256=7gNeanyHkgX23Iz7RMCQTCkE9gI5asCI11nXxM0jjx4,9832
 unique_toolkit/services/knowledge_base.py,sha256=qvq9JhRUNciyvTn7PaJEuPGFHAyZbM8LzxCXfBo1Mqg,45625
 unique_toolkit/short_term_memory/__init__.py,sha256=2mI3AUrffgH7Yt-xS57EGqnHf7jnn6xquoKEhJqk3Wg,185
@@ -426,6 +428,6 @@ unique_toolkit/short_term_memory/service.py,sha256=2JMBEBn6GoyK-a0xvdcWpPft-3C74
 unique_toolkit/smart_rules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unique_toolkit/smart_rules/compile.py,sha256=vRe0aMJWuVIK-VsqomgxQImzke7MDmwgSxxJfYCa7zw,1386
 unique_toolkit/test_utilities/events.py,sha256=FLrhMh0A_mQutAUjoD0FdMcVQ4Gq31r35OVGzRjvLro,6760
-unique_toolkit-2026.44.0.dev4.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
-unique_toolkit-2026.44.0.dev4.dist-info/METADATA,sha256=UwPoIX-zlOyvOVUJDR8KMsPSckfeOq2IsAM-2rrogmI,7821
-unique_toolkit-2026.44.0.dev4.dist-info/RECORD,,
+unique_toolkit-2026.44.0.dev5.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
+unique_toolkit-2026.44.0.dev5.dist-info/METADATA,sha256=eIr1xETlMh66pzXJR__P64AzIJa1TPhYmfn8bL_v8Xw,7821
+unique_toolkit-2026.44.0.dev5.dist-info/RECORD,,