unique-toolkit 2026.32.0.dev3__tar.gz → 2026.32.0.dev5__tar.gz

{unique_toolkit-2026.32.0.dev3 → unique_toolkit-2026.32.0.dev5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 2026.32.0.dev3
+Version: 2026.32.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.32.0.dev2,<2026.32.0rc0
+Requires-Dist: unique-sdk>=2026.32.0.dev8,<2026.32.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.32.0.dev3 → unique_toolkit-2026.32.0.dev5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "2026.32.0.dev3"
+version = "2026.32.0.dev5"
 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.32.0.dev2,<2026.32.0rc0",
+    "unique-sdk>=2026.32.0.dev8,<2026.32.0rc0",
     "jambo>=0.1.2,<0.2",
     "jinja2>=3.1.6,<4",
     "docxtpl>=0.20.1,<0.21",

{unique_toolkit-2026.32.0.dev3 → unique_toolkit-2026.32.0.dev5}/unique_toolkit/agentic/debug_info_manager/debug_info_manager.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from typing import Any, Required, TypedDict
 
 from unique_toolkit.agentic.tools.openai_builtin import OpenAICodeInterpreterTool
@@ -7,8 +8,10 @@ from unique_toolkit.agentic.tools.openai_builtin.code_interpreter.postprocessors
 )
 from unique_toolkit.agentic.tools.schemas import ToolCallResponse
 from unique_toolkit.agentic.tools.tool_manager import ToolManager
+from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
 from unique_toolkit.language_model.schemas import (
     LanguageModelStreamResponse,
+    LanguageModelTokenUsage,
     ResponsesLanguageModelStreamResponse,
 )
 
@@ -40,6 +43,16 @@ class AnalyticsToolName(TypedDict):
     display_name: str
 
 
+class AnalyticsTokenUsage(TypedDict):
+    model_name: str
+    completion_tokens: int | None
+    prompt_tokens: int | None
+    total_tokens: int | None
+    reasoning_tokens: int | None
+    cached_tokens: int | None
+    cache_write_tokens: int | None
+
+
 class Analytics(TypedDict):
     answer_length: int
     artifacts_created_count: int | None
@@ -54,6 +67,7 @@ class Analytics(TypedDict):
     references_count: int
     skills_used: list[AnalyticsSkill]
     subagent_names_used: list[AnalyticsToolName]
+    tokens: list[AnalyticsTokenUsage]
     tool_call_count: int
     tools_used: list[AnalyticsTool]
     total_time_to_answer_ms: int | None
@@ -115,6 +129,7 @@ class DebugInfoManager:
         total_time_to_answer_ms: int | None = None,
         artifacts: ArtifactsDebugInfo | None = None,
         context_memory_updated: bool | None = None,
+        invocations: list[LanguageModelInvocationStats] | None = None,
     ) -> None:
         """Add a stable 'analytics' snapshot for downstream ROI/usage reporting.
 
@@ -152,6 +167,7 @@ class DebugInfoManager:
             loop_iteration_count=loop_iteration_count,
             subagent_names_used=_unique_tool_names(analytics_tools, "is_sub_agent"),
             mcp_tool_names_used=_unique_tool_names(analytics_tools, "is_mcp"),
+            tokens=_aggregate_tokens_by_model(invocations or []),
             total_time_to_answer_ms=total_time_to_answer_ms,
             artifacts_created_count=artifacts["count"] if artifacts else None,
             artifacts_created_filetype=artifacts["filetypes"] if artifacts else None,
@@ -163,6 +179,32 @@ class DebugInfoManager:
         self.add("analytics", analytics)
 
 
+def _aggregate_tokens_by_model(
+    invocations: list[LanguageModelInvocationStats],
+) -> list[AnalyticsTokenUsage]:
+    usages_by_model: dict[str, list[LanguageModelTokenUsage]] = defaultdict(list)
+    for invocation in invocations:
+        usages_by_model[str(invocation.model_name)].append(invocation.token_usage)
+
+    totals: list[AnalyticsTokenUsage] = []
+    for model_name in sorted(usages_by_model):
+        usage = LanguageModelTokenUsage.sum_usages(usages_by_model[model_name])
+        if usage is None:
+            continue
+        totals.append(
+            AnalyticsTokenUsage(
+                model_name=model_name,
+                completion_tokens=usage.completion_tokens,
+                prompt_tokens=usage.prompt_tokens,
+                total_tokens=usage.total_tokens,
+                reasoning_tokens=usage.reasoning_tokens,
+                cached_tokens=usage.cached_tokens,
+                cache_write_tokens=usage.cache_write_tokens,
+            )
+        )
+    return totals
+
+
 def _unique_tool_names(
     analytics_tools: list[AnalyticsTool], flag: str
 ) -> list[AnalyticsToolName]:

{unique_toolkit-2026.32.0.dev3 → unique_toolkit-2026.32.0.dev5}/unique_toolkit/agentic/postprocessor/postprocessor_manager.py

@@ -41,6 +41,17 @@ class Postprocessor(ABC):
         """Override in subclasses whose `run()` makes its own LLM call(s)."""
         return []
 
+    def take_pending_invocation_stats(self) -> list[LanguageModelInvocationStats]:
+        """Pop usage incurred before this turn's `run()` was reached.
+
+        Called unconditionally at the start of every turn so usage from work
+        done ahead of postprocessing (e.g. a load-time LLM call) is still
+        reported on turns that exit before `run()` executes. Override in
+        subclasses that do such work; must clear the pending stats so `run()`
+        does not report them a second time.
+        """
+        return []
+
 
 class ResponsesApiPostprocessor(ABC):
     def __init__(self, name: str):
@@ -74,6 +85,13 @@ class ResponsesApiPostprocessor(ABC):
         """
         return []
 
+    def take_pending_invocation_stats(self) -> list[LanguageModelInvocationStats]:
+        """Pop usage incurred before this turn's `run()` was reached.
+
+        See `Postprocessor.take_pending_invocation_stats` for why this exists.
+        """
+        return []
+
 
 class PostprocessorManager:
     """
@@ -196,6 +214,19 @@ class PostprocessorManager:
         """Per-LLM-call stats from every postprocessor that made its own LLM call(s)."""
         return list(self._invocation_stats)
 
+    def take_pending_invocation_stats(self) -> list[LanguageModelInvocationStats]:
+        """Pop pre-`run()` usage (e.g. load-time calls) from every postprocessor.
+
+        Call this unconditionally at the start of a turn, before it's known
+        whether `run_postprocessors` will even be reached this turn, so that
+        usage is never silently dropped on turns that exit early.
+        """
+        return [
+            invocation
+            for postprocessor in self._get_all_postprocessors()
+            for invocation in postprocessor.take_pending_invocation_stats()
+        ]
+
     async def execute_postprocessors(
         self,
         loop_response: LanguageModelStreamResponse,

{unique_toolkit-2026.32.0.dev3 → unique_toolkit-2026.32.0.dev5}/unique_toolkit/agentic/tools/a2a/tool/service.py

@@ -4,12 +4,15 @@ import json
 import logging
 import re
 from datetime import datetime
-from typing import override
+from typing import Any, override
 
 import unique_sdk
 from pydantic import Field, TypeAdapter, create_model
 from unique_sdk.api_resources._space import Space
-from unique_sdk.utils.chat_in_space import send_message_and_wait_for_completion
+from unique_sdk.utils.chat_in_space import (
+    get_message_invocations,
+    send_message_and_wait_for_completion,
+)
 
 from unique_toolkit._common.execution import SafeTaskExecutor
 from unique_toolkit._common.referencing import (
@@ -55,6 +58,7 @@ from unique_toolkit.language_model import (
     LanguageModelService,
     LanguageModelToolDescription,
 )
+from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
 
 logger = logging.getLogger(__name__)
 
@@ -296,6 +300,13 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
                 if self.config.passthrough_config.enabled:
                     await self._display_sub_agent_response(response)
 
+                sub_agent_invocations = _prefix_sub_agent_invocation_sources(
+                    invocations=_parse_sub_agent_invocation_stats(
+                        get_message_invocations(response)
+                    ),
+                    tool_name=self.name,
+                    tool_call_id=tool_call.id,
+                )
                 return ToolCallResponse(
                     id=tool_call.id,
                     name=tool_call.name,
@@ -307,9 +318,12 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
                     content_chunks=content_chunks,
                     debug_info={
                         "chat_id": response["chatId"],
+                        "user_message_id": response.get("triggeringUserMessageId"),
+                        "assistant_message_id": response.get("id"),
                         "assistant_id": self.config.assistant_id,
                         "display_name": self._display_name,
                     },
+                    invocation_stats=sub_agent_invocations,
                 )
         except TimeoutError as e:
             raise e
@@ -495,6 +509,7 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
                     parentAssistantId=self.config.assistant_id,
                 ),
                 on_message_update=on_message_update,
+                wait_for_invocations=True,
             )
         except TimeoutError as e:
             await self._notify_progress(
@@ -513,6 +528,42 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             ) from e
 
 
+def _parse_sub_agent_invocation_stats(
+    raw_invocations: list[dict[str, Any]],
+) -> list[LanguageModelInvocationStats]:
+    """Deserialize sub-agent invocation stats, ignoring malformed entries.
+
+    A single unparseable entry (e.g. a stale model name) must not fail the
+    whole tool call and drop an otherwise-successful sub-agent reply -- see
+    `_load_invocation_stats_from_debug_info` in unique_orchestrator, which
+    applies the same per-entry tolerance for the same reason.
+    """
+    invocations: list[LanguageModelInvocationStats] = []
+    for raw_invocation in raw_invocations:
+        try:
+            invocations.append(
+                LanguageModelInvocationStats.model_validate(raw_invocation)
+            )
+        except Exception:
+            logger.warning(
+                "Ignoring malformed sub-agent LLM invocation stats",
+                exc_info=True,
+            )
+    return invocations
+
+
+def _prefix_sub_agent_invocation_sources(
+    invocations: list[LanguageModelInvocationStats],
+    tool_name: str,
+    tool_call_id: str,
+) -> list[LanguageModelInvocationStats]:
+    prefix = f"subagent.{tool_name}.{tool_call_id}"
+    return [
+        invocation.model_copy(update={"source": f"{prefix}.{invocation.source}"})
+        for invocation in invocations
+    ]
+
+
 def _format_response(tool_name: str, text: str, system_reminders: list[str]) -> str:
     if len(system_reminders) == 0:
         return text

{unique_toolkit-2026.32.0.dev3 → unique_toolkit-2026.32.0.dev5}/unique_toolkit/agentic/tools/a2a/tool/test/test_service_utils.py

@@ -17,6 +17,7 @@ from unique_toolkit.agentic.tools.a2a.tool.config import (
 from unique_toolkit.agentic.tools.a2a.tool.service import (
     _format_response,
     _get_sub_agent_system_reminders,
+    _parse_sub_agent_invocation_stats,
     _prepare_sub_agent_response_refs,
     _remove_extra_refs,
     _response_has_refs,
@@ -59,6 +60,78 @@ def refs_list() -> list[dict[str, Any]]:
     ]
 
 
+# Tests for _parse_sub_agent_invocation_stats
+
+
+@pytest.mark.ai
+def test_parse_sub_agent_invocation_stats__skips_malformed_entry__keeps_valid_ones() -> (
+    None
+):
+    """
+    Purpose: Verify one malformed invocation entry doesn't discard the rest.
+    Why this matters: A single bad entry must not fail the whole sub-agent tool
+        call and drop an otherwise-successful reply.
+    Setup summary: Mix a valid entry with a malformed one (missing required
+        fields), verify only the valid entry survives.
+    """
+    # Arrange
+    raw_invocations = [
+        {
+            "modelName": "gpt-4",
+            "tokenUsage": {"totalTokens": 10},
+            "source": "main_loop",
+        },
+        {"modelName": "gpt-4"},  # missing required tokenUsage/source
+    ]
+
+    # Act
+    result = _parse_sub_agent_invocation_stats(raw_invocations)
+
+    # Assert
+    assert len(result) == 1
+    assert result[0].source == "main_loop"
+
+
+@pytest.mark.ai
+def test_parse_sub_agent_invocation_stats__returns_empty__when_all_malformed() -> None:
+    """
+    Purpose: Verify an all-malformed list returns empty instead of raising.
+    Why this matters: The sub-agent tool call must still succeed and return the
+        answer even if every persisted stats entry is unparseable.
+    Setup summary: Only malformed entries, assert empty list with no exception.
+    """
+    # Arrange
+    raw_invocations = [{"modelName": "gpt-4"}, {"source": "main_loop"}]
+
+    # Act
+    result = _parse_sub_agent_invocation_stats(raw_invocations)
+
+    # Assert
+    assert result == []
+
+
+@pytest.mark.ai
+def test_parse_sub_agent_invocation_stats__parses_all__when_all_valid() -> None:
+    """
+    Purpose: Verify parsing is unaffected by the malformed-entry tolerance
+        when every entry is well-formed.
+    Why this matters: The fix must not silently drop good data too.
+    Setup summary: Two valid entries, assert both are returned.
+    """
+    # Arrange
+    raw_invocations = [
+        {"modelName": "gpt-4", "tokenUsage": {"totalTokens": 3}, "source": "a"},
+        {"modelName": "gpt-4", "tokenUsage": {"totalTokens": 5}, "source": "b"},
+    ]
+
+    # Act
+    result = _parse_sub_agent_invocation_stats(raw_invocations)
+
+    # Assert
+    assert len(result) == 2
+    assert [inv.source for inv in result] == ["a", "b"]
+
+
 # Tests for _format_response
 
 

{unique_toolkit-2026.32.0.dev3 → unique_toolkit-2026.32.0.dev5}/unique_toolkit/monitoring/middleware.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import time
+from collections.abc import Sequence
 
 from prometheus_client import (  # pyright: ignore[reportMissingImports]
     Counter,
@@ -16,12 +17,6 @@ _http_requests_total = Counter(
     ["method", "path", "status_code"],
     registry=REGISTRY,
 )
-_http_request_duration_seconds = Histogram(
-    "python_http_request_duration_seconds",
-    "HTTP request duration",
-    ["method", "path"],
-    registry=REGISTRY,
-)
 _http_requests_in_progress = Gauge(
     "python_http_requests_in_progress",
     "In-flight HTTP requests",
@@ -30,6 +25,35 @@ _http_requests_in_progress = Gauge(
     multiprocess_mode="livesum",
 )
 
+# Created lazily on first MetricsMiddleware construction so services can choose the
+# bucket layout. Prometheus allows a metric name to be registered only once per
+# registry, so the buckets passed to the first middleware win for the process.
+_http_request_duration_seconds: Histogram | None = None
+
+
+def _get_or_create_duration_histogram(
+    buckets: Sequence[float] | None,
+) -> Histogram:
+    global _http_request_duration_seconds
+    if _http_request_duration_seconds is None:
+        if buckets is None:
+            _http_request_duration_seconds = Histogram(
+                "python_http_request_duration_seconds",
+                "HTTP request duration",
+                ["method", "path"],
+                registry=REGISTRY,
+            )
+        else:
+            _http_request_duration_seconds = Histogram(
+                "python_http_request_duration_seconds",
+                "HTTP request duration",
+                ["method", "path"],
+                registry=REGISTRY,
+                buckets=tuple(buckets),
+            )
+    return _http_request_duration_seconds
+
+
 _DEFAULT_EXCLUDED_PATHS = frozenset({"/health", "/metrics", "/"})
 
 
@@ -51,13 +75,29 @@ class MetricsMiddleware:
 
         # Quart
         app.asgi_app = MetricsMiddleware(app.asgi_app)
+
+    ``duration_buckets`` overrides the latency histogram buckets (seconds). Without
+    it, prometheus_client's defaults apply and their top finite bucket is 10s —
+    histogram_quantile cannot resolve past the largest finite bucket, so services
+    with slower endpoints should pass buckets covering their full latency range::
+
+        app.add_middleware(MetricsMiddleware, duration_buckets=(1, 5, 30, 60, 120))
+
+    The histogram is registered once per process; the buckets of the first
+    constructed middleware apply.
     """
 
-    def __init__(self, app, excluded_paths: set[str] | None = None):
+    def __init__(
+        self,
+        app,
+        excluded_paths: set[str] | None = None,
+        duration_buckets: Sequence[float] | None = None,
+    ):
         self.app = app
         self._excluded_paths = (
             excluded_paths if excluded_paths is not None else _DEFAULT_EXCLUDED_PATHS
         )
+        self._duration = _get_or_create_duration_histogram(duration_buckets)
 
     async def __call__(self, scope, receive, send):
         # Pass through non-HTTP (WebSocket, lifespan)
@@ -87,9 +127,7 @@ class MetricsMiddleware:
         finally:
             duration = time.perf_counter() - start
             _http_requests_in_progress.labels(method=method).dec()
-            _http_request_duration_seconds.labels(method=method, path=path).observe(
-                duration
-            )
+            self._duration.labels(method=method, path=path).observe(duration)
             _http_requests_total.labels(
                 method=method, path=path, status_code=status_code
             ).inc()