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

unique_toolkit/language_model/invocation_stats.py

@@ -3,9 +3,18 @@
 Lives in its own module (not ``schemas.py``) because it needs
 ``LanguageModelName`` from ``infos.py``, which itself imports from
 ``schemas.py``.
+
+:class:`InvocationStatsCollector` collects those stats for one namespaced
+run. ``record_token_usage`` swallows all exceptions around parsing and
+:meth:`LanguageModelInvocationStats.from_usage`: usage is optional
+billing metadata and must not abort a tool run.
 """
 
-from typing import Annotated, Self
+import logging
+from collections.abc import Iterable, Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from typing import Annotated, Any, Self
 
 from humps import camelize
 from pydantic import BaseModel, BeforeValidator, ConfigDict
@@ -14,6 +23,8 @@ from unique_toolkit.language_model.infos import LanguageModelName
 from unique_toolkit.language_model.model_costs import calculate_invocation_cost_usd
 from unique_toolkit.language_model.schemas import LanguageModelTokenUsage
 
+_LOGGER = logging.getLogger(__name__)
+
 # `protected_namespaces=()` allows the `model_name` field name.
 model_config = ConfigDict(
     alias_generator=camelize,
@@ -75,3 +86,88 @@ class LanguageModelInvocationStats(BaseModel):
             source=source,
             cost_usd=calculate_invocation_cost_usd(str(model_name), token_usage),
         )
+
+
+class InvocationStatsCollector:
+    """Run-scoped collector with a namespaced ContextVar so nested tools don't mix.
+
+    Instantiate once per namespace and share that instance.
+    """
+
+    def __init__(self, namespace: str) -> None:
+        self._namespace = namespace
+        self._var: ContextVar[list[LanguageModelInvocationStats] | None] = ContextVar(
+            f"{namespace}_invocation_stats", default=None
+        )
+
+    @contextmanager
+    def scope(self) -> Iterator[list[LanguageModelInvocationStats]]:
+        """Collect stats for one run and isolate concurrent runs."""
+        invocation_stats: list[LanguageModelInvocationStats] = []
+        token: Token[list[LanguageModelInvocationStats] | None] = self._var.set(
+            invocation_stats
+        )
+        try:
+            yield invocation_stats
+        finally:
+            self._var.reset(token)
+
+    def record_invocation_stats(
+        self,
+        invocation_stats: Iterable[LanguageModelInvocationStats],
+    ) -> None:
+        """Record already-built stats when a run scope is active. No-op otherwise."""
+        current_invocation_stats = self._var.get()
+        if current_invocation_stats is not None:
+            current_invocation_stats.extend(invocation_stats)
+
+    def record_token_usage(
+        self,
+        *,
+        model_name: str,
+        usage: Any,
+        source: str,
+    ) -> None:
+        """Record provider usage when a run scope and usage are available.
+
+        Swallows all exceptions: this metadata is optional and must not abort
+        the calling tool.
+        """
+        invocation_stats = self._var.get()
+        if invocation_stats is None or usage is None:
+            return
+
+        try:
+            token_usage = LanguageModelTokenUsage.model_validate(usage)
+            invocation_stats.append(
+                LanguageModelInvocationStats.from_usage(
+                    model_name=model_name,
+                    token_usage=token_usage,
+                    source=source,
+                )
+            )
+        except Exception:
+            _LOGGER.warning(
+                "Unable to parse %s token usage for %s",
+                self._namespace,
+                source,
+                exc_info=True,
+            )
+
+    def record_language_model_response(
+        self,
+        *,
+        model_name: str,
+        response: Any,
+        source: str,
+    ) -> None:
+        """Record usage from ``.usage``, ``.usage_metadata``, or ``response_metadata``."""
+        usage = getattr(response, "usage", None)
+        if usage is None:
+            usage = getattr(response, "usage_metadata", None)
+        if usage is None:
+            response_metadata = getattr(response, "response_metadata", None)
+            if isinstance(response_metadata, dict):
+                usage = response_metadata.get("token_usage")
+
+        self.record_token_usage(model_name=model_name, usage=usage, source=source)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 2026.36.0.dev4
+Version: 2026.36.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>

{unique_toolkit-2026.36.0.dev4.dist-info → unique_toolkit-2026.36.0.dev5.dist-info}/RECORD

@@ -394,7 +394,7 @@ unique_toolkit/language_model/constants.py,sha256=B-topqW0r83dkC_25DeQfnPk3n53qz
 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=zs0UFhds4UkSRr3zxyTJ-jJ9u2O2x9VgZx3a9d4zbYU,146221
-unique_toolkit/language_model/invocation_stats.py,sha256=XAM_qckI-BoMBSgDpVcxFD1cFSdi1-GHZToHz_jLynM,2491
+unique_toolkit/language_model/invocation_stats.py,sha256=CQbpUz394IsVd8eNTB1gPSMt5f0UXmkN1amNsqtAW6U,5857
 unique_toolkit/language_model/model_costs.py,sha256=WM1u2fWAKpPB_creS-ew8YAzc1V8oiaCbvhpSJaO8B0,3789
 unique_toolkit/language_model/prompt.py,sha256=JSawaLjQg3VR-E2fK8engFyJnNdk21zaO8pPIodzN4Q,3991
 unique_toolkit/language_model/reference.py,sha256=T3wr2jP_uRg5wKVIbKqLbxJlcgxfbZQGq3XdjbQ2kSo,6984
@@ -421,6 +421,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.36.0.dev4.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
-unique_toolkit-2026.36.0.dev4.dist-info/METADATA,sha256=cFB1qZ_frQUp2w6_eMAL5auiBr8ELgaGhfLApCzNywI,7821
-unique_toolkit-2026.36.0.dev4.dist-info/RECORD,,
+unique_toolkit-2026.36.0.dev5.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
+unique_toolkit-2026.36.0.dev5.dist-info/METADATA,sha256=XnRiUVEB6mURHDnoA_QevRw7qH1IadNMN1XBKSUgiLg,7821
+unique_toolkit-2026.36.0.dev5.dist-info/RECORD,,