unique-web-search 2026.36.0.dev2__tar.gz → 2026.36.0.dev4__tar.gz

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-web-search
-Version: 2026.36.0.dev2
+Version: 2026.36.0.dev4
 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>=2.3,<3
 Requires-Dist: azure-identity>=1.25.0,<2
-Requires-Dist: unique-toolkit[monitoring]>=2026.36.0.dev6,<2026.36.0rc0
+Requires-Dist: unique-toolkit[monitoring]>=2026.36.0.dev10,<2026.36.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.36.0.dev2,<2026.36.0rc0
+Requires-Dist: unique-search-proxy-core>=2026.36.0.dev3,<2026.36.0rc0
 Requires-Dist: unique-search-proxy-sdk>=2026.36.0.dev1,<2026.36.0rc0
 Requires-Dist: openai>=1.105.0,<3
 Requires-Python: >=3.12, <4.0

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_web_search"
-version = "2026.36.0.dev2"
+version = "2026.36.0.dev4"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12,<4.0"
@@ -22,7 +22,7 @@ dependencies = [
     "unidecode>=1.4.0,<2",
     "azure-ai-projects>=2.3,<3",
     "azure-identity>=1.25.0,<2",
-    "unique-toolkit[monitoring]>=2026.36.0.dev6,<2026.36.0rc0",
+    "unique-toolkit[monitoring]>=2026.36.0.dev10,<2026.36.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.36.0.dev2,<2026.36.0rc0",
+    "unique-search-proxy-core>=2026.36.0.dev3,<2026.36.0rc0",
     "unique-search-proxy-sdk>=2026.36.0.dev1,<2026.36.0rc0",
     "openai>=1.105.0,<3",
 ]

unique_web_search-2026.36.0.dev4/src/unique_web_search/invocation_stats.py

@@ -0,0 +1,31 @@
+"""Run-scoped collection of language-model invocation usage."""
+
+from typing import Any
+
+from unique_toolkit.language_model.invocation_stats import InvocationStatsCollector
+
+collector = InvocationStatsCollector("web_search")
+
+
+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
+
+    collector.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.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/service.py

@@ -8,14 +8,13 @@ from unique_search_proxy_core.param_policy.exposed_params import ExposedParams
 from unique_search_proxy_core.search_engines.base import BaseSearchEngineConfig
 from unique_toolkit._common.chunk_relevancy_sorter.service import ChunkRelevancySorter
 from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
-from unique_toolkit.agentic.feature_flags.feature_flags import (
-    feature_flags,
-)
+from unique_toolkit.agentic.feature_flags import FeatureFlagNames
 from unique_toolkit.agentic.tools.factory import ToolFactory
 from unique_toolkit.agentic.tools.schemas import ToolCallResponse
 from unique_toolkit.agentic.tools.tool import (
     Tool,
 )
+from unique_toolkit.experimental.resources.feature_flags import is_flag_enabled
 from unique_toolkit.language_model.infos import LanguageModelInfo
 from unique_toolkit.language_model.schemas import (
     LanguageModelFunction,
@@ -24,7 +23,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.invocation_stats import collector
 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 (
@@ -152,7 +151,7 @@ class WebSearchTool(Tool[WebSearchConfig]):
 
     @override
     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:
@@ -311,8 +310,9 @@ class WebSearchTool(Tool[WebSearchConfig]):
     async def _get_argument_screening_service_if_ff_enabled(
         self,
     ) -> ArgumentScreeningService | None:
-        if not feature_flags.enable_web_search_argument_screening_un_18741.is_enabled(
-            self.company_id
+        if not await is_flag_enabled(
+            FeatureFlagNames.enable_web_search_argument_screening_un_18741,
+            company_id=self.company_id,
         ):
             return None
 

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/services/content_processing/processing_strategies/llm_guard_judge.py

@@ -6,7 +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.invocation_stats import collector
 from unique_web_search.services.content_processing.processing_strategies.base import (
     ProcessingStrategyKwargs,
     WebSearchResult,
@@ -317,7 +317,7 @@ class LLMGuardJudge:
                 "web_search.content_processing.guard_judge_and_sanitize"
             ),
         }
-        record_language_model_response(
+        collector.record_language_model_response(
             model_name=self._config.language_model.name,
             response=response,
             source=source_by_response_model[response_model],

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/services/content_processing/processing_strategies/llm_keyword_redact.py

@@ -8,7 +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.invocation_stats import collector
 from unique_web_search.services.content_processing.processing_strategies.base import (
     ProcessingStrategyKwargs,
     WebSearchResult,
@@ -181,7 +181,7 @@ class LLMKeywordRedact:
             structured_output_model=KeywordRedactResponse,
             structured_output_enforce_schema=True,
         )
-        record_language_model_response(
+        collector.record_language_model_response(
             model_name=self._config.language_model.name,
             response=response,
             source="web_search.content_processing.keyword_redact",

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/services/content_processing/processing_strategies/llm_process.py

@@ -13,7 +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.invocation_stats import collector
 from unique_web_search.services.content_processing.processing_strategies.base import (
     ProcessingStrategyKwargs,
     WebSearchResult,
@@ -317,7 +317,7 @@ class LLMProcess:
             structured_output_model=LLMProcessorResponse,
             structured_output_enforce_schema=True,
         )
-        record_language_model_response(
+        collector.record_language_model_response(
             model_name=self._config.language_model.name,
             response=response,
             source="web_search.content_processing.llm_process",

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/services/executors/base_executor.py

@@ -12,7 +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.invocation_stats import collector
 from unique_web_search.metrics import llm_duration, llm_errors
 from unique_web_search.schema import (
     StepDebugInfo,
@@ -165,7 +165,7 @@ class BaseWebSearchExecutor(ABC, Generic[T]):
             # 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(
+            collector.record_invocation_stats(
                 invocation
                 for relevancy in sorted_chunks.relevancies
                 if relevancy.relevancy is not None

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/services/executors/v1/executor.py

@@ -12,7 +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.invocation_stats import collector
 from unique_web_search.metrics import (
     crawl_duration,
     crawl_errors,
@@ -125,7 +125,7 @@ async def query_generation_agent(
         structured_output_model=structured_output_model,
         structured_output_enforce_schema=True,
     )
-    record_language_model_response(
+    collector.record_language_model_response(
         model_name=language_model.name,
         response=response,
         source="web_search.query_refinement",

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/services/search_engine/README.md

@@ -161,6 +161,12 @@ AZURE_IDENTITY_CREDENTIAL_TYPE=workload  # or default
 
 Requires GCP ADC / workload identity. Fields include `vertexai_model_name`, `generation_instructions`, `enable_entreprise_search`, `enable_redirect_resolution`, `requires_scraping`.
 
+```bash
+VERTEXAI_AGENT_DEFAULT_VERTEX_AI_MODEL_NAME=gemini-3-flash-preview
+# ZDR tenants: pin "Enable Enterprise Search" on and disable the checkbox
+VERTEXAI_AGENT_FORCE_ACTIVATE_ENTERPRISE_SEARCH=false
+```
+
 ---
 
 ## Factory

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/services/search_engine/utils/grounding/bing/runner.py

@@ -24,7 +24,7 @@ from openai.types.responses.response_output_text import (
 )
 from openai.types.responses.response_text_delta_event import ResponseTextDeltaEvent
 
-from unique_web_search.invocation_stats import record_token_usage
+from unique_web_search.invocation_stats import collector
 from unique_web_search.services.search_engine.schema import (
     WebSearchResult,
 )
@@ -284,7 +284,7 @@ async def _run_responses_agent(
                 if output_text and not emitted_text:
                     emitted_text = True
                     answer_parts.append(output_text)
-                record_token_usage(
+                collector.record_token_usage(
                     model_name=model,
                     usage=event.response.usage,
                     source="web_search.grounding.bing",

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/services/search_engine/utils/grounding/response_parsing.py

@@ -6,7 +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.invocation_stats import collector
 from unique_web_search.services.search_engine.schema import WebSearchResult
 from unique_web_search.services.search_engine.utils.grounding.models import (
     GroundingSearchResults,
@@ -93,7 +93,7 @@ class LLMParserStrategy(ResponseParser):
             structured_output_model=GroundingSearchResults,
             structured_output_enforce_schema=True,
         )
-        record_language_model_response(
+        collector.record_language_model_response(
             model_name=self.llm.name,
             response=llm_response,
             source="web_search.grounding.response_parser",

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/services/snippet_judge/service.py

@@ -9,7 +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.invocation_stats import collector
 from unique_web_search.services.helpers import extract_registered_domain
 from unique_web_search.services.search_engine.schema import (
     WebSearchResult,
@@ -100,7 +100,7 @@ async def score_and_explain(
         structured_output_model=SnippetJudgeResponse,
         structured_output_enforce_schema=True,
     )
-    record_language_model_response(
+    collector.record_language_model_response(
         model_name=language_model.name,
         response=response,
         source="web_search.snippet_judge",

{unique_web_search-2026.36.0.dev2 → unique_web_search-2026.36.0.dev4}/src/unique_web_search/services/structured_llm.py

@@ -7,7 +7,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.invocation_stats import collector
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -40,7 +40,7 @@ async def complete_structured_llm(
         structured_output_model=response_model,
         structured_output_enforce_schema=structured_output_enforce_schema,
     )
-    record_language_model_response(
+    collector.record_language_model_response(
         model_name=language_model.name,
         response=response,
         source=invocation_source,

unique_web_search-2026.36.0.dev2/src/unique_web_search/invocation_stats.py

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