unique_internal_search 2026.36.0.dev2__tar.gz → 2026.36.0.dev3__tar.gz

{unique_internal_search-2026.36.0.dev2 → unique_internal_search-2026.36.0.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-internal-search
-Version: 2026.36.0.dev2
+Version: 2026.36.0.dev3
 Summary: 
 Author: Martin Fadler, Sadique Sheik, Fabian Schläpfer, Pascal Hauri
 Author-email: Martin Fadler <martin.fadler@unique.ch>, Sadique Sheik <sadique@unique.ch>, Fabian Schläpfer <fabian@unique.ch>, Pascal Hauri <pascal@unique.ch>

{unique_internal_search-2026.36.0.dev2 → unique_internal_search-2026.36.0.dev3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_internal_search"
-version = "2026.36.0.dev2"
+version = "2026.36.0.dev3"
 description = ""
 readme = "README.md"
 license = { text = "Proprietary" }

unique_internal_search-2026.36.0.dev3/unique_internal_search/invocation_stats.py

@@ -0,0 +1,5 @@
+"""Run-scoped collection of language-model invocation usage."""
+
+from unique_toolkit.language_model.invocation_stats import InvocationStatsCollector
+
+collector = InvocationStatsCollector("internal_search")

{unique_internal_search-2026.36.0.dev2 → unique_internal_search-2026.36.0.dev3}/unique_internal_search/service.py

@@ -37,10 +37,7 @@ from unique_toolkit.language_model.schemas import (
 )
 
 from unique_internal_search.config import InternalSearchConfig
-from unique_internal_search.invocation_stats import (
-    invocation_stats_scope,
-    record_invocation_stats,
-)
+from unique_internal_search.invocation_stats import collector
 from unique_internal_search.services.message_log import (
     InternalSearchMessageLogger,
     InternalSearchMessageLoggerNoop,
@@ -322,7 +319,7 @@ class InternalSearchService:
                 config=self.config.chunk_relevancy_sort_config,
             )
             if isinstance(chunk_relevancy_sorter_result.relevancies, list):
-                record_invocation_stats(
+                collector.record_invocation_stats(
                     invocation
                     for relevancy in chunk_relevancy_sorter_result.relevancies
                     if relevancy.relevancy is not None
@@ -493,7 +490,7 @@ class InternalSearchTool(Tool[InternalSearchConfig], InternalSearchService):
     # TODO: find a solution for tracking
     # @track(name="internal_search_tool_run")
     async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
-        with invocation_stats_scope() as invocation_stats:
+        with collector.scope() as invocation_stats:
             try:
                 response = await self._run(tool_call)
             except Exception as e:

unique_internal_search-2026.36.0.dev2/unique_internal_search/invocation_stats.py

@@ -1,33 +0,0 @@
-"""Run-scoped collection of language-model invocation usage."""
-
-from collections.abc import Generator, Iterable
-from contextlib import contextmanager
-from contextvars import ContextVar, Token
-
-from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
-
-_CURRENT_INVOCATION_STATS: ContextVar[list[LanguageModelInvocationStats] | None] = (
-    ContextVar("internal_search_invocation_stats", default=None)
-)
-
-
-@contextmanager
-def invocation_stats_scope() -> Generator[list[LanguageModelInvocationStats]]:
-    """Collect stats for one internal-search run and isolate concurrent runs."""
-    invocation_stats: list[LanguageModelInvocationStats] = []
-    token: Token[list[LanguageModelInvocationStats] | None] = (
-        _CURRENT_INVOCATION_STATS.set(invocation_stats)
-    )
-    try:
-        yield invocation_stats
-    finally:
-        _CURRENT_INVOCATION_STATS.reset(token)
-
-
-def record_invocation_stats(
-    invocation_stats: Iterable[LanguageModelInvocationStats],
-) -> None:
-    """Record usage when an internal-search run scope is active."""
-    current_invocation_stats = _CURRENT_INVOCATION_STATS.get()
-    if current_invocation_stats is not None:
-        current_invocation_stats.extend(invocation_stats)