unique_web_search 2026.32.0.dev0__tar.gz → 2026.32.0.dev1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-web-search
-Version: 2026.32.0.dev0
+Version: 2026.32.0.dev1
 Summary: 
 Author: Andreas Hauri, Gustav Hartz, Rami Azouz
 Author-email: Andreas Hauri <andreas@unique.ch>, Gustav Hartz <gustav.hartz.ext@unique.ch>, Rami Azouz <rami.ext@unique.ch>
@@ -16,7 +16,7 @@ Requires-Dist: tavily-python>=0.7.11,<1
 Requires-Dist: unidecode>=1.4.0,<2
 Requires-Dist: azure-ai-projects>=1.0.0,<2
 Requires-Dist: azure-identity>=1.25.0,<2
-Requires-Dist: unique-toolkit[monitoring]>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: unique-toolkit[monitoring]>=2026.32.0.dev4,<2026.32.0rc0
 Requires-Dist: azure-core>=1.36.0,<2
 Requires-Dist: google-cloud-aiplatform>=1.128.0,<2
 Requires-Dist: google-auth>=2.43.0,<3
@@ -25,7 +25,7 @@ Requires-Dist: rapidfuzz>=3.0.0,<4
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: certifi>=2025.11.12
 Requires-Dist: jinja2>=3.1.6
-Requires-Dist: unique-search-proxy-core>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: unique-search-proxy-core>=2026.32.0.dev2,<2026.32.0rc0
 Requires-Dist: unique-search-proxy-sdk>=2026.32.0.dev0,<2026.32.0rc0
 Requires-Python: >=3.12, <4.0
 Description-Content-Type: text/markdown

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

@@ -1,6 +1,6 @@
 [project]
 name = "unique_web_search"
-version = "2026.32.0.dev0"
+version = "2026.32.0.dev1"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12,<4.0"
@@ -22,7 +22,7 @@ dependencies = [
     "unidecode>=1.4.0,<2",
     "azure-ai-projects>=1.0.0,<2",
     "azure-identity>=1.25.0,<2",
-    "unique-toolkit[monitoring]>=2026.32.0.dev0,<2026.32.0rc0",
+    "unique-toolkit[monitoring]>=2026.32.0.dev4,<2026.32.0rc0",
     "azure-core>=1.36.0,<2",
     "google-cloud-aiplatform>=1.128.0,<2",
     "google-auth>=2.43.0,<3",
@@ -31,7 +31,7 @@ dependencies = [
     "httpx>=0.28.1",
     "certifi>=2025.11.12",
     "jinja2>=3.1.6",
-    "unique-search-proxy-core>=2026.32.0.dev0,<2026.32.0rc0",
+    "unique-search-proxy-core>=2026.32.0.dev2,<2026.32.0rc0",
     "unique-search-proxy-sdk>=2026.32.0.dev0,<2026.32.0rc0",
 ]
 

unique_web_search-2026.32.0.dev1/src/unique_web_search/invocation_stats.py

@@ -0,0 +1,113 @@
+"""Run-scoped collection of language-model invocation usage."""
+
+import logging
+from collections.abc import Iterable, Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from typing import Any
+
+from pydantic import ValidationError
+from unique_toolkit.language_model.invocation_stats import LanguageModelInvocationStats
+from unique_toolkit.language_model.schemas import LanguageModelTokenUsage
+
+_LOGGER = logging.getLogger(__name__)
+
+_CURRENT_INVOCATION_STATS: ContextVar[list[LanguageModelInvocationStats] | None] = (
+    ContextVar("web_search_invocation_stats", default=None)
+)
+
+
+@contextmanager
+def invocation_stats_scope() -> Iterator[list[LanguageModelInvocationStats]]:
+    """Collect stats for one WebSearch 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 already-built invocation stats when a run scope is active.
+
+    Used for usage produced by nested dependencies that hand back
+    ``LanguageModelInvocationStats`` objects directly (e.g. the chunk relevancy
+    sorter attaching per-chunk stats on its result), rather than a raw provider
+    usage payload.
+    """
+    current_invocation_stats = _CURRENT_INVOCATION_STATS.get()
+    if current_invocation_stats is not None:
+        current_invocation_stats.extend(invocation_stats)
+
+
+def record_token_usage(
+    *,
+    model_name: str,
+    usage: Any,
+    source: str,
+) -> None:
+    """Record provider usage when both a run scope and usage are available."""
+    invocation_stats = _CURRENT_INVOCATION_STATS.get()
+    if invocation_stats is None or usage is None:
+        return
+
+    try:
+        token_usage = LanguageModelTokenUsage.model_validate(usage)
+    except ValidationError:
+        _LOGGER.warning(
+            "Unable to parse Web Search token usage for %s",
+            source,
+            exc_info=True,
+        )
+        return
+
+    invocation_stats.append(
+        LanguageModelInvocationStats.from_usage(
+            model_name=model_name,
+            token_usage=token_usage,
+            source=source,
+        )
+    )
+
+
+def record_language_model_response(
+    *,
+    model_name: str,
+    response: Any,
+    source: str,
+) -> None:
+    """Record usage from a Unique Toolkit language-model response."""
+    record_token_usage(
+        model_name=model_name,
+        usage=getattr(response, "usage", None),
+        source=source,
+    )
+
+
+def record_vertex_response(
+    *,
+    model_name: str,
+    response: Any,
+    source: str,
+) -> None:
+    """Record Google GenAI usage metadata using toolkit token names."""
+    usage = getattr(response, "usage_metadata", None)
+    if usage is None:
+        return
+
+    record_token_usage(
+        model_name=model_name,
+        usage={
+            "prompt_tokens": getattr(usage, "prompt_token_count", None),
+            "completion_tokens": getattr(usage, "candidates_token_count", None),
+            "total_tokens": getattr(usage, "total_token_count", None),
+            "reasoning_tokens": getattr(usage, "thoughts_token_count", None),
+            "cached_tokens": getattr(usage, "cached_content_token_count", None),
+        },
+        source=source,
+    )

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/service.py

@@ -25,6 +25,7 @@ from unique_toolkit.language_model.schemas import (
 from unique_toolkit.monitoring import metric_scope
 
 from unique_web_search.config import WebSearchConfig
+from unique_web_search.invocation_stats import invocation_stats_scope
 from unique_web_search.metrics import tool_duration, tool_empty_results, tool_errors
 from unique_web_search.schema import WebSearchDebugInfo
 from unique_web_search.services.argument_screening import (
@@ -104,7 +105,6 @@ class WebSearchTool(Tool[WebSearchConfig]):
         self._mode_strategy = get_mode_strategy(self.config.web_search_mode_config)
 
         def content_reducer(web_page_chunks: list[WebPageChunk]) -> list[WebPageChunk]:
-
             return reduce_sources_to_token_limit(
                 web_page_chunks,
                 self.config.language_model_max_input_tokens,
@@ -153,6 +153,31 @@ class WebSearchTool(Tool[WebSearchConfig]):
 
     @override
     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() already catches its own errors and returns a
+                # ToolCallResponse gracefully in the common case; this is a
+                # safety net for the rarer case where something raises while
+                # handling that failure (or otherwise escapes _run()
+                # uncaught). Without it, tokens already spent on screening/
+                # processing/relevancy before the failure would never reach
+                # analytics, since SafeTaskExecutor builds its own fresh,
+                # stats-less error response one level up.
+                # logger.exception() already attaches the active exception's
+                # type and traceback; don't stringify it into the message.
+                _LOGGER.exception("WebSearch 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:
         _LOGGER.info("Running the WebSearch tool")
         start_time = time()
         parameters = self.tool_parameter_calls.model_validate(

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/argument_screening/service.py

@@ -116,6 +116,7 @@ class ArgumentScreeningService:
                 system_message=self._config.system_prompt,
                 user_message=user_prompt,
                 response_model=ArgumentScreeningResult,
+                invocation_source="web_search.argument_screening",
             )
         except StructuredLlmUnparseableResponseError as e:
             raise ArgumentScreeningUnparseableResponseException(

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/content_processing/processing_strategies/llm_guard_judge.py

@@ -6,6 +6,7 @@ from unique_toolkit._common.utils.jinja.render import render_template
 from unique_toolkit.language_model import LanguageModelService
 from unique_toolkit.language_model.builder import MessagesBuilder
 
+from unique_web_search.invocation_stats import record_language_model_response
 from unique_web_search.services.content_processing.processing_strategies.base import (
     ProcessingStrategyKwargs,
     WebSearchResult,
@@ -309,6 +310,18 @@ class LLMGuardJudge:
             structured_output_model=response_model,
             structured_output_enforce_schema=True,
         )
+        source_by_response_model = {
+            LLMGuardResponse: "web_search.content_processing.guard_sanitize",
+            JudgeResponse: "web_search.content_processing.guard_judge",
+            JudgeAndSanitizeResponse: (
+                "web_search.content_processing.guard_judge_and_sanitize"
+            ),
+        }
+        record_language_model_response(
+            model_name=self._config.language_model.name,
+            response=response,
+            source=source_by_response_model[response_model],
+        )
 
         if response.choices[0].message.parsed is None:
             raise ValueError("Judge call returned no parsed response")

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/content_processing/processing_strategies/llm_keyword_redact.py

@@ -8,6 +8,7 @@ from unique_toolkit._common.utils.jinja.render import render_template
 from unique_toolkit.language_model import LanguageModelService
 from unique_toolkit.language_model.builder import MessagesBuilder
 
+from unique_web_search.invocation_stats import record_language_model_response
 from unique_web_search.services.content_processing.processing_strategies.base import (
     ProcessingStrategyKwargs,
     WebSearchResult,
@@ -180,6 +181,11 @@ class LLMKeywordRedact:
             structured_output_model=KeywordRedactResponse,
             structured_output_enforce_schema=True,
         )
+        record_language_model_response(
+            model_name=self._config.language_model.name,
+            response=response,
+            source="web_search.content_processing.keyword_redact",
+        )
 
         if response.choices[0].message.parsed is None:
             raise ValueError("Keyword extraction call returned no parsed response")

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/content_processing/processing_strategies/llm_process.py

@@ -13,6 +13,7 @@ from unique_toolkit.language_model import (
 )
 from unique_toolkit.language_model.builder import MessagesBuilder
 
+from unique_web_search.invocation_stats import record_language_model_response
 from unique_web_search.services.content_processing.processing_strategies.base import (
     ProcessingStrategyKwargs,
     WebSearchResult,
@@ -316,6 +317,11 @@ class LLMProcess:
             structured_output_model=LLMProcessorResponse,
             structured_output_enforce_schema=True,
         )
+        record_language_model_response(
+            model_name=self._config.language_model.name,
+            response=response,
+            source="web_search.content_processing.llm_process",
+        )
 
         if response.choices[0].message.parsed is None:
             raise ValueError("Failed to parse response")

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/executors/base_executor.py

@@ -12,6 +12,7 @@ from unique_toolkit.content import ContentChunk
 from unique_toolkit.language_model import LanguageModelFunction
 from unique_toolkit.monitoring import metric_scope
 
+from unique_web_search.invocation_stats import record_invocation_stats
 from unique_web_search.metrics import llm_duration, llm_errors
 from unique_web_search.schema import (
     StepDebugInfo,
@@ -160,6 +161,16 @@ class BaseWebSearchExecutor(ABC, Generic[T]):
                 chunks=content,
                 config=self.chunk_relevancy_sort_config,
             )
+            # The relevancy sorter makes one LLM call per chunk and attaches the
+            # usage to each ChunkRelevancy.relevancy.invocation_stats; merge it
+            # into the active WebSearch stats scope so it isn't dropped from
+            # token analytics (Internal Search does the same after its sort).
+            record_invocation_stats(
+                invocation
+                for relevancy in sorted_chunks.relevancies
+                if relevancy.relevancy is not None
+                for invocation in relevancy.relevancy.invocation_stats
+            )
             _LOGGER.info(f"Sorting chunks message: {sorted_chunks.user_message}")
             return sorted_chunks.content_chunks
         else:

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/executors/v1/executor.py

@@ -12,6 +12,7 @@ from unique_toolkit.language_model import LanguageModelFunction
 from unique_toolkit.language_model.builder import MessagesBuilder
 from unique_toolkit.monitoring import metric_scope
 
+from unique_web_search.invocation_stats import record_language_model_response
 from unique_web_search.metrics import (
     crawl_duration,
     crawl_errors,
@@ -124,6 +125,11 @@ async def query_generation_agent(
         structured_output_model=structured_output_model,
         structured_output_enforce_schema=True,
     )
+    record_language_model_response(
+        model_name=language_model.name,
+        response=response,
+        source="web_search.query_refinement",
+    )
 
     parsed_response = response.choices[0].message.parsed
     if parsed_response is None:

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/search_engine/utils/grounding/bing/runner.py

@@ -10,6 +10,7 @@ from azure.ai.agents.models import (
 )
 from azure.ai.projects.aio import AIProjectClient
 
+from unique_web_search.invocation_stats import record_token_usage
 from unique_web_search.services.search_engine.schema import (
     WebSearchResult,
 )
@@ -129,6 +130,12 @@ async def create_and_process_run(
         _LOGGER.warning(f"Using existing agent ID: {agent_id}")
         thread = await _create_agent_run_with_agent_id(agent_client, agent_id, query)
 
+    record_token_usage(
+        model_name=env_settings.azure_ai_bing_agent_model,
+        usage=getattr(thread, "usage", None),
+        source="web_search.grounding.bing",
+    )
+
     if thread.status in [RunStatus.FAILED, RunStatus.CANCELLED, RunStatus.EXPIRED]:
         raise Exception(f"Run failed: {thread.last_error}")
 

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/search_engine/utils/grounding/response_parsing.py

@@ -6,6 +6,7 @@ from unique_toolkit._common.validators import LMI
 from unique_toolkit.language_model.builder import MessagesBuilder
 from unique_toolkit.language_model.service import LanguageModelService
 
+from unique_web_search.invocation_stats import record_language_model_response
 from unique_web_search.services.search_engine.schema import WebSearchResult
 from unique_web_search.services.search_engine.utils.grounding.models import (
     GroundingSearchResults,
@@ -92,6 +93,11 @@ class LLMParserStrategy(ResponseParser):
             structured_output_model=GroundingSearchResults,
             structured_output_enforce_schema=True,
         )
+        record_language_model_response(
+            model_name=self.llm.name,
+            response=llm_response,
+            source="web_search.grounding.response_parser",
+        )
 
         if not llm_response.choices[0].message.parsed:
             raise ValueError("No JSON found in the response")

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/search_engine/vertexai.py

@@ -13,6 +13,7 @@ from unique_toolkit._common.default_language_model import DEFAULT_LANGUAGE_MODEL
 from unique_toolkit._common.validators import LMI, get_LMI_default_field
 from unique_toolkit.language_model import LanguageModelService
 
+from unique_web_search.invocation_stats import record_vertex_response
 from unique_web_search.services.proxy.bridge import (
     search_proxy_client_enabled,
 )
@@ -114,6 +115,11 @@ class VertexAI(SearchEngine[VertexAIConfig]):
             ),
             contents=query,
         )
+        record_vertex_response(
+            model_name=self.config.vertexai_model_name,
+            response=response,
+            source="web_search.grounding.vertexai",
+        )
 
         answer_with_citations = add_citations(response)
 

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/snippet_judge/service.py

@@ -9,6 +9,7 @@ from unique_toolkit._common.validators import LMI
 from unique_toolkit.language_model import LanguageModelService
 from unique_toolkit.language_model.builder import MessagesBuilder
 
+from unique_web_search.invocation_stats import record_language_model_response
 from unique_web_search.services.helpers import extract_registered_domain
 from unique_web_search.services.search_engine.schema import (
     WebSearchResult,
@@ -99,6 +100,11 @@ async def score_and_explain(
         structured_output_model=SnippetJudgeResponse,
         structured_output_enforce_schema=True,
     )
+    record_language_model_response(
+        model_name=language_model.name,
+        response=response,
+        source="web_search.snippet_judge",
+    )
 
     parsed = response.choices[0].message.parsed
     if parsed is None:

{unique_web_search-2026.32.0.dev0 → unique_web_search-2026.32.0.dev1}/src/unique_web_search/services/structured_llm.py

@@ -7,6 +7,8 @@ from unique_toolkit._common.validators import LMI
 from unique_toolkit.language_model import LanguageModelService
 from unique_toolkit.language_model.builder import MessagesBuilder
 
+from unique_web_search.invocation_stats import record_language_model_response
+
 T = TypeVar("T", bound=BaseModel)
 
 
@@ -21,6 +23,7 @@ async def complete_structured_llm(
     system_message: str,
     user_message: str,
     response_model: type[T],
+    invocation_source: str,
     structured_output_enforce_schema: bool = True,
 ) -> T:
     """Run a structured completion and return a validated model instance."""
@@ -37,6 +40,11 @@ async def complete_structured_llm(
         structured_output_model=response_model,
         structured_output_enforce_schema=structured_output_enforce_schema,
     )
+    record_language_model_response(
+        model_name=language_model.name,
+        response=response,
+        source=invocation_source,
+    )
 
     parsed = response.choices[0].message.parsed
     if parsed is None: