unique_internal_search 2026.32.0.dev0__tar.gz → 2026.32.0.dev2__tar.gz

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

@@ -1,13 +1,13 @@
 Metadata-Version: 2.3
 Name: unique-internal-search
-Version: 2026.32.0.dev0
+Version: 2026.32.0.dev2
 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.dev8,<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.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_internal_search"
-version = "2026.32.0.dev0"
+version = "2026.32.0.dev2"
 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.dev8,<2026.32.0rc0",
 ]
 
 [build-system]

{unique_internal_search-2026.32.0.dev0 → unique_internal_search-2026.32.0.dev2}/unique_internal_search/config.py

@@ -33,8 +33,7 @@ from unique_internal_search.prompts import (
 )
 from unique_internal_search.validators import get_string_field_with_pattern_validation
 
-DEFAULT_LIMIT_CHUNK_RELEVANCY_SORT_ENABLED = 200
-DEFAULT_LIMIT_CHUNK_RELEVANCY_SORT_DISABLED = 1000
+DEFAULT_SEARCH_LIMIT = 200
 
 
 class ToolResponseSystemReminderConfig(BaseModel):
@@ -72,14 +71,6 @@ class ExperimentalFeatures(FeatureExtendedSourceSerialization):
     )
 
 
-def _search_limit_factory(data: dict[str, Any]) -> int:
-    return (
-        DEFAULT_LIMIT_CHUNK_RELEVANCY_SORT_ENABLED
-        if data["chunk_relevancy_sort_config"].enabled
-        else DEFAULT_LIMIT_CHUNK_RELEVANCY_SORT_DISABLED
-    )
-
-
 _FIELD_ALIASES: dict[str, str] = {
     "ftsSearchLanguage": "searchLanguage",
 }
@@ -145,8 +136,8 @@ class InternalSearchConfig(BaseToolConfig):
         description="The chunk relevancy sort config to use for the search.",
     )
     limit: int = Field(
-        default_factory=_search_limit_factory,
-        description="The maximum limit of chunks returned by the search.",
+        default=DEFAULT_SEARCH_LIMIT,
+        description="Maximum number of chunks returned in total across all search strings, capped by the token budget. Also used as the per-search-string fetch ceiling.",
     )
     chat_only: bool = Field(
         default=False,

unique_internal_search-2026.32.0.dev2/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.dev2}/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,
@@ -224,6 +228,13 @@ class InternalSearchService:
             for result in found_chunks_per_search_string
             for chunk in result.chunks
         ]
+        if len(found_chunks) > self.config.limit:
+            self.logger.info(
+                "Truncating %s chunks to the configured total limit of %s",
+                len(found_chunks),
+                self.config.limit,
+            )
+            found_chunks = found_chunks[: self.config.limit]
         selected_chunks = pick_content_chunks_for_token_window(
             found_chunks,
             self._get_max_tokens(),
@@ -255,13 +266,12 @@ class InternalSearchService:
         content_ids: list[str] | None = None,
     ) -> SearchStringResult:
         try:
-            capped_limit = self._cap_limit_to_token_budget()
             found_chunks: list[
                 ContentChunk
             ] = await self.content_service.search_content_chunks_async(
                 search_string=search_string,  # type: ignore
                 search_type=self.config.search_type,
-                limit=capped_limit,
+                limit=self._cap_limit_to_token_budget(),
                 reranker_config=self.config.reranker_config,
                 search_language=self.config.search_language,
                 scope_ids=self.config.scope_ids,
@@ -311,6 +321,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}")
@@ -348,10 +365,12 @@ class InternalSearchService:
         capped_limit = min(self.config.limit, token_based_limit)
         if capped_limit < self.config.limit:
             self.logger.info(
-                f"Search limit capped from {self.config.limit} to {capped_limit} (token budget)"
+                "Search limit capped from %s to %s (token budget)",
+                self.config.limit,
+                capped_limit,
             )
         else:
-            self.logger.info(f"Search limit: {capped_limit} (within token budget)")
+            self.logger.info("Search limit: %s (within token budget)", capped_limit)
         return capped_limit
 
 
@@ -474,6 +493,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.
         """