unique_internal_search 2026.32.0.dev0__tar.gz → 2026.32.0.dev1__tar.gz

{unique_internal_search-2026.32.0.dev0 → unique_internal_search-2026.32.0.dev1}/PKG-INFO

@@ -1,13 +1,13 @@
 Metadata-Version: 2.3
 Name: unique-internal-search
-Version: 2026.32.0.dev0
+Version: 2026.32.0.dev1
 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>
 License: Proprietary
 Requires-Dist: pydantic>=2.8.2,<3
 Requires-Dist: typing-extensions>=4.9.0,<5
-Requires-Dist: unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: unique-toolkit>=2026.32.0.dev4,<2026.32.0rc0
 Requires-Python: >=3.12, <4
 Description-Content-Type: text/markdown
 

{unique_internal_search-2026.32.0.dev0 → unique_internal_search-2026.32.0.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_internal_search"
-version = "2026.32.0.dev0"
+version = "2026.32.0.dev1"
 description = ""
 readme = "README.md"
 license = { text = "Proprietary" }
@@ -14,7 +14,7 @@ requires-python = ">=3.12,<4"
 dependencies = [
     "pydantic>=2.8.2,<3",
     "typing-extensions>=4.9.0,<5",
-    "unique-toolkit>=2026.32.0.dev0,<2026.32.0rc0",
+    "unique-toolkit>=2026.32.0.dev4,<2026.32.0rc0",
 ]
 
 [build-system]

unique_internal_search-2026.32.0.dev1/unique_internal_search/invocation_stats.py

@@ -0,0 +1,33 @@
+"""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)

{unique_internal_search-2026.32.0.dev0 → unique_internal_search-2026.32.0.dev1}/unique_internal_search/service.py

@@ -37,6 +37,10 @@ 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.services.message_log import (
     InternalSearchMessageLogger,
     InternalSearchMessageLoggerNoop,
@@ -311,6 +315,13 @@ class InternalSearchService:
                 chunks=found_chunks,
                 config=self.config.chunk_relevancy_sort_config,
             )
+            if isinstance(chunk_relevancy_sorter_result.relevancies, list):
+                record_invocation_stats(
+                    invocation
+                    for relevancy in chunk_relevancy_sorter_result.relevancies
+                    if relevancy.relevancy is not None
+                    for invocation in relevancy.relevancy.invocation_stats
+                )
             found_chunks = chunk_relevancy_sorter_result.content_chunks
         except ChunkRelevancySorterException as e:
             self.logger.warning(f"Error while sorting chunks: {e.error_message}")
@@ -474,6 +485,29 @@ 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:
+            try:
+                response = await self._run(tool_call)
+            except Exception as e:
+                # _run() has no top-level try/except of its own, so any error
+                # after self.search() (which may already have spent relevancy-
+                # sorting tokens) would otherwise escape straight to
+                # SafeTaskExecutor, which builds a fresh, stats-less error
+                # response one level up -- silently dropping tokens already
+                # spent before the failure.
+                # logger.exception() already attaches the active exception's
+                # type and traceback; don't stringify it into the message.
+                self.logger.exception("InternalSearch tool run failed")
+                return ToolCallResponse(
+                    id=tool_call.id,  # type: ignore
+                    name=self.name,
+                    error_message=str(e),
+                    invocation_stats=invocation_stats,
+                )
+        response.invocation_stats.extend(invocation_stats)
+        return response
+
+    async def _run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
         """
         Perform a search in the Vector DB based on the user's message and generate a response.
         """