unique_web_search 2026.32.0.dev0__tar.gz → 2026.32.0.dev2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-web-search
-Version: 2026.32.0.dev0
+Version: 2026.32.0.dev2
 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>
@@ -14,9 +14,9 @@ Requires-Dist: fake-useragent>=2.2.0,<3
 Requires-Dist: firecrawl>=3.3.2,<4
 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-ai-projects>=2.3,<3
 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,8 +25,9 @@ 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.dev3,<2026.32.0rc0
 Requires-Dist: unique-search-proxy-sdk>=2026.32.0.dev0,<2026.32.0rc0
+Requires-Dist: openai>=1.105.0,<3
 Requires-Python: >=3.12, <4.0
 Description-Content-Type: text/markdown
 

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

@@ -1,6 +1,6 @@
 [project]
 name = "unique_web_search"
-version = "2026.32.0.dev0"
+version = "2026.32.0.dev2"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12,<4.0"
@@ -20,9 +20,9 @@ dependencies = [
     "firecrawl>=3.3.2,<4",
     "tavily-python>=0.7.11,<1",
     "unidecode>=1.4.0,<2",
-    "azure-ai-projects>=1.0.0,<2",
+    "azure-ai-projects>=2.3,<3",
     "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,8 +31,9 @@ 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.dev3,<2026.32.0rc0",
     "unique-search-proxy-sdk>=2026.32.0.dev0,<2026.32.0rc0",
+    "openai>=1.105.0,<3",
 ]
 
 [build-system]
@@ -91,4 +92,4 @@ extend-select = ["I"]
 [tool.uv.sources]
 unique-toolkit = { workspace = true }
 unique-search-proxy-sdk = { workspace = true }
-unique-search-proxy-core = { workspace = true }
+unique-search-proxy-core = { workspace = true }

unique_web_search-2026.32.0.dev2/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.dev2}/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.dev2}/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.dev2}/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.dev2}/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.dev2}/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.dev2}/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.dev2}/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.dev2}/src/unique_web_search/services/search_engine/bing.py

@@ -100,13 +100,14 @@ class BingSearch(SearchEngine[BingSearchConfig]):
         del params
         agent_client = get_project_client(self.credentials, self.config.endpoint)
 
-        search_results = await create_and_process_run(
-            agent_client,
-            agent_id=self.config.agent_id,
-            query=query,
-            fetch_size=self.config.fetch_size,
-            response_parsers_strategies=self.response_parsers,
-            generation_instructions=self.config.generation_instructions,
-        )
+        async with agent_client:
+            search_results = await create_and_process_run(
+                agent_client,
+                agent_id=self.config.agent_id,
+                query=query,
+                fetch_size=self.config.fetch_size,
+                response_parsers_strategies=self.response_parsers,
+                generation_instructions=self.config.generation_instructions,
+            )
 
         return search_results

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

@@ -1,14 +1,17 @@
 import logging
 
 import certifi
+import httpx
 from azure.ai.projects.aio import AIProjectClient
 from azure.core.credentials_async import AsyncTokenCredential
 from azure.core.pipeline.transport import AsyncioRequestsTransport
 from azure.identity.aio import DefaultAzureCredential, WorkloadIdentityCredential
+from openai import AsyncOpenAI
 
 from unique_web_search.settings import env_settings
 
 _LOGGER = logging.getLogger(__name__)
+_private_endpoint_http_client: httpx.AsyncClient | None = None
 
 
 def _get_workload_identity_credentials(
@@ -16,14 +19,11 @@ def _get_workload_identity_credentials(
 ) -> WorkloadIdentityCredential:
     if with_request_transport:
         transport = AsyncioRequestsTransport(connection_verify=certifi.where())
-        credentials = WorkloadIdentityCredential(transport=transport)
-    else:
-        credentials = WorkloadIdentityCredential()
-
-    return credentials
+        return WorkloadIdentityCredential(transport=transport)
+    return WorkloadIdentityCredential()
 
 
-def get_credentials():
+def get_credentials() -> AsyncTokenCredential:
     match env_settings.azure_identity_credential_type:
         case "workload":
             return _get_workload_identity_credentials(
@@ -31,10 +31,8 @@ def get_credentials():
             )
         case "default":
             return DefaultAzureCredential()
-        case _:
-            raise ValueError(
-                f"Invalid Azure identity credential type: {env_settings.azure_identity_credential_type}"
-            )
+        case other:
+            raise ValueError(f"Invalid Azure identity credential type: {other}")
 
 
 async def credentials_are_valid(credentials: AsyncTokenCredential) -> bool:
@@ -72,3 +70,36 @@ def get_project_client(
             credential=credentials,
             endpoint=endpoint,
         )
+
+
+def _get_private_endpoint_http_client() -> httpx.AsyncClient:
+    """Return a process-wide certifi-backed httpx client (created once)."""
+    global _private_endpoint_http_client
+    if _private_endpoint_http_client is None or _private_endpoint_http_client.is_closed:
+        _private_endpoint_http_client = httpx.AsyncClient(verify=certifi.where())
+    return _private_endpoint_http_client
+
+
+async def aclose_private_endpoint_http_client() -> None:
+    """Close the shared private-endpoint httpx client if it was created."""
+    global _private_endpoint_http_client
+    if (
+        _private_endpoint_http_client is not None
+        and not _private_endpoint_http_client.is_closed
+    ):
+        await _private_endpoint_http_client.aclose()
+    _private_endpoint_http_client = None
+
+
+def get_openai_client(project_client: AIProjectClient) -> AsyncOpenAI:
+    """Return an authenticated AsyncOpenAI client from the Foundry project client.
+
+    When private-endpoint transport is enabled, reuse a shared certifi-backed
+    ``httpx.AsyncClient`` so TLS verification matches the AIProjectClient path
+    without leaking a new client per request.
+    """
+    if env_settings.use_unique_private_endpoint_transport:
+        return project_client.get_openai_client(
+            http_client=_get_private_endpoint_http_client(),
+        )
+    return project_client.get_openai_client()