From a0c319878b3f0b10b287dc3d80ed9e36747ba513 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 24 Aug 2026 12:26:07 -0700 Subject: [PATCH] fix(search): harden bing_grounding auth, result cap, status, and cost - send a caller api_key via the Azure api-key header instead of Authorization: Bearer - cap web_search results to the requested max_results (the tool has no count knob) - surface a Foundry failed/incomplete response status as a 502 error - zero the per-query cost in web_search mode; keep the map price for connection mode - trim the example config to terse env-var pointers --- litellm/llms/azure/search/transformation.py | 130 ++++++++++++++---- litellm/llms/custom_httpx/llm_http_handler.py | 2 + .../bing_grounding_websearch_config.yaml | 34 ++--- .../test_bing_grounding_search.py | 14 +- ...st_bing_grounding_search_transformation.py | 63 ++++++++- 5 files changed, 190 insertions(+), 53 deletions(-) diff --git a/litellm/llms/azure/search/transformation.py b/litellm/llms/azure/search/transformation.py index 2caee9b50a0..7bc631e814d 100644 --- a/litellm/llms/azure/search/transformation.py +++ b/litellm/llms/azure/search/transformation.py @@ -12,10 +12,11 @@ Setup: 3. Optional: set BING_GROUNDING_CONNECTION_ID to a Grounding with Bing Search project connection id to use the `bing_grounding` tool; without it the project's built-in `web_search` tool is used - 4. Auth: pass api_key, or set BING_GROUNDING_TOKEN to an Entra bearer token for - scope https://ai.azure.com/.default, or configure azure-identity - (AZURE_CLIENT_ID / AZURE_CLIENT_SECRET / AZURE_TENANT_ID, managed identity, - or any DefaultAzureCredential source) and the token is minted automatically + 4. Auth: pass api_key (an Azure API key, sent in the api-key header), or set + BING_GROUNDING_TOKEN to an Entra bearer token for scope + https://ai.azure.com/.default, or configure azure-identity (AZURE_CLIENT_ID / + AZURE_CLIENT_SECRET / AZURE_TENANT_ID, managed identity, or any + DefaultAzureCredential source) and the token is minted automatically Usage: response = litellm.search( @@ -56,6 +57,8 @@ ENTRA_SCOPE: Final = "https://ai.azure.com/.default" _RESPONSES_PATH: Final = "/openai/v1/responses" _SNIPPET_FALLBACK_LENGTH: Final = 300 +_UPSTREAM_ERROR_STATUS: Final = 502 +_RESPONSE_COST_HEADER: Final = "llm_provider-x-litellm-response-cost" class _Annotation(BaseModel): @@ -83,21 +86,33 @@ class _OutputItem(BaseModel): content: tuple[_ContentPart, ...] = () -class _ResponsesEnvelope(BaseModel): - """A Foundry Responses API body. `output` is required: a body without it is not a - Responses API response and must not be reported as a successful empty search.""" - - model_config = ConfigDict(extra="ignore", frozen=True) - - output: tuple[_OutputItem, ...] - - class _ErrorBody(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) message: str | None = None +class _IncompleteDetails(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + reason: str | None = None + + +class _ResponsesEnvelope(BaseModel): + """A Foundry Responses API body. `output` is required: a body without it is not a + Responses API response and must not be reported as a successful empty search. + + A 200 body can still carry `status` `failed` or `incomplete`; those are surfaced as + errors rather than reported as a successful empty search.""" + + model_config = ConfigDict(extra="ignore", frozen=True) + + output: tuple[_OutputItem, ...] + status: str | None = None + error: _ErrorBody | None = None + incomplete_details: _IncompleteDetails | None = None + + class _ErrorEnvelope(BaseModel): model_config = ConfigDict(extra="ignore", frozen=True) @@ -163,6 +178,26 @@ def _citation_results(envelope: _ResponsesEnvelope) -> tuple[SearchResult, ...]: return tuple(first_by_url[url] for url in dict.fromkeys(result.url for result in cited)) +def _requested_max_results(response_kwargs: Mapping[str, object]) -> int | None: + """The unified `max_results` cap the caller asked for, if any. + + The built-in web_search tool has no server-side result-count knob, so the cap is + enforced here after the fact; connection mode also honors it as a hard ceiling on + top of the tool's `count` hint. + """ + optional_params: Final = response_kwargs.get("optional_params") + if not isinstance(optional_params, Mapping): + return None + max_results: Final = optional_params.get("max_results") + return ( + max_results if isinstance(max_results, int) and not isinstance(max_results, bool) and max_results > 0 else None + ) + + +def _capped(results: tuple[SearchResult, ...], max_results: int | None) -> tuple[SearchResult, ...]: + return results[:max_results] if max_results is not None else results + + class _SearchConfiguration(BaseModel): model_config = ConfigDict(frozen=True) @@ -247,18 +282,28 @@ class BingGroundingSearchConfig(BaseSearchConfig): Returns a new dict rather than mutating ``headers``: the http handler calls this a second time after ``litellm/search/main.py`` already did, so it has to be idempotent. """ - resolved_token: Final = self.resolve_server_api_key( - caller_api_key=api_key, + return { # mutable-ok: httpx requires a plain dict of headers + **headers, + **self._auth_header(api_key, api_base), + "Content-Type": "application/json", + } + + def _auth_header(self, api_key: str | None, api_base: str | None) -> Mapping[str, str]: + """ + A caller-supplied ``api_key`` is an Azure API key and rides the ``api-key`` header; + an Entra bearer token (``BING_GROUNDING_TOKEN`` or one minted via azure-identity) + rides ``Authorization: Bearer``. Foundry rejects the wrong scheme for each. + """ + if api_key: + return MappingProxyType({"api-key": api_key}) + token: Final = self.resolve_server_api_key( + caller_api_key=None, caller_api_base=api_base, key_env_vars=(TOKEN_ENV,), base_env_var=PROJECT_ENDPOINT_ENV, default_api_base=None, ) or self._mint_entra_token(api_base) - return { # mutable-ok: httpx requires a plain dict of headers - **headers, - "Authorization": f"Bearer {resolved_token}", - "Content-Type": "application/json", - } + return MappingProxyType({"Authorization": f"Bearer {token}"}) def _mint_entra_token(self, caller_api_base: str | None) -> str: self._assert_trusted_api_base_for_server_credential( @@ -303,8 +348,9 @@ class BingGroundingSearchConfig(BaseSearchConfig): Transform Search request to the Foundry Responses API format. The unified params map as far as the API allows: - - max_results -> the bing_grounding search configuration's `count` (the built-in - web_search tool has no result-count knob, so it is dropped in that mode) + - max_results -> the bing_grounding search configuration's `count`; the built-in + web_search tool has no result-count knob, so that mode instead caps the returned + results after the fact (see transform_search_response) - country -> web_search's approximate `user_location` (bing_grounding's `market` wants a full locale like en-US, which a bare country code cannot fill) - search_domain_filter, max_tokens_per_page -> no API equivalent, dropped @@ -336,8 +382,44 @@ class BingGroundingSearchConfig(BaseSearchConfig): status_code=raw_response.status_code, headers=dict(raw_response.headers), # mutable-ok: BaseSearchConfig.get_error_class signature ) - results: Final = list(_citation_results(parsed)) # mutable-ok: SearchResponse.results is list[SearchResult] - return SearchResponse(results=results, object="search") + if parsed.status == "failed": + detail: Final = ( + parsed.error.message if parsed.error and parsed.error.message else "the grounded search failed" + ) + raise self._upstream_error(detail, raw_response) + results: Final = _capped(_citation_results(parsed), _requested_max_results(kwargs)) + if not results and parsed.status == "incomplete": + reason: Final = ( + parsed.incomplete_details.reason + if parsed.incomplete_details and parsed.incomplete_details.reason + else "unknown reason" + ) + raise self._upstream_error(f"the grounded search was incomplete: {reason}", raw_response) + return self._priced(results) + + def _upstream_error(self, detail: str, raw_response: httpx.Response) -> Exception: + return self.get_error_class( + error_message=detail, + status_code=_UPSTREAM_ERROR_STATUS, + headers=dict(raw_response.headers), # mutable-ok: BaseSearchConfig.get_error_class signature + ) + + def _priced(self, results: tuple[SearchResult, ...]) -> SearchResponse: + """web_search mode runs no paid Grounding with Bing transaction, so it must not + inherit the connection-mode ``bing_grounding/search`` price; zero its per-query + cost while leaving connection mode to the cost map.""" + response: Final = SearchResponse( + results=list(results), # mutable-ok: SearchResponse.results is list[SearchResult] + object="search", + ) + if get_secret_str(CONNECTION_ID_ENV): + return response + response._hidden_params[ + "additional_headers" + ] = { # mutable-ok: response_cost_calculator writes into _hidden_params + _RESPONSE_COST_HEADER: 0.0 + } + return response def get_error_class( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ed079197513..8fa025d6e4c 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1844,6 +1844,7 @@ class BaseLLMHTTPHandler: return provider_config.transform_search_response( raw_response=response, logging_obj=logging_obj, + optional_params=optional_params, ) async def async_search( @@ -1942,6 +1943,7 @@ class BaseLLMHTTPHandler: return provider_config.transform_search_response( raw_response=response, logging_obj=logging_obj, + optional_params=optional_params, ) async def _async_post_anthropic_messages_with_http_error_retry( diff --git a/litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml b/litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml index 5ab18723f7f..ac8a3db2d21 100644 --- a/litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml +++ b/litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml @@ -1,24 +1,14 @@ -# Web search via Microsoft Foundry: Grounding with Bing Search / the built-in -# web_search tool, called through the Foundry Responses API. -# See litellm/llms/azure/search/transformation.py for details. +# Web search via Microsoft Foundry (Grounding with Bing Search / the built-in +# web_search tool), called through the Foundry Responses API. # -# Required environment variables (the search router forwards only -# search_provider / api_key / api_base from the litellm_params block, so -# provider configuration rides env vars): -# BING_GROUNDING_PROJECT_ENDPOINT: the Foundry project endpoint, e.g. -# https://.services.ai.azure.com/api/projects/ -# BING_GROUNDING_MODEL: a model deployment in that project (e.g. gpt-4.1); -# it runs the grounded search, its tokens are billed on that deployment -# Optional: -# BING_GROUNDING_CONNECTION_ID: a Grounding with Bing Search project -# connection id; set it to use the bing_grounding tool ($35 per 1,000 -# transactions on the G1 SKU). Without it the project's built-in -# web_search tool is used -# BING_GROUNDING_TOKEN: an Entra bearer token for scope -# https://ai.azure.com/.default. Without it (and without api_key below) -# the token is minted via azure-identity (AZURE_CLIENT_ID / -# AZURE_CLIENT_SECRET / AZURE_TENANT_ID, managed identity, or any other -# DefaultAzureCredential source) +# Configure the provider with env vars (setup and pricing are in the LiteLLM docs; +# the code lives in litellm/llms/azure/search/transformation.py): +# BING_GROUNDING_PROJECT_ENDPOINT (required) the Foundry project endpoint +# BING_GROUNDING_MODEL (required) a model deployment in that project +# BING_GROUNDING_CONNECTION_ID (optional) a Grounding with Bing connection id; +# without it the built-in web_search tool is used +# BING_GROUNDING_TOKEN (optional) an Entra bearer token; without it (and +# without api_key) azure-identity mints one model_list: - model_name: claude-sonnet @@ -30,8 +20,8 @@ search_tools: - search_tool_name: bing-grounding-search litellm_params: search_provider: bing_grounding - # Alternative to BING_GROUNDING_TOKEN / azure-identity: - # api_key: os.environ/BING_GROUNDING_TOKEN + # Optional: an Azure API key instead of BING_GROUNDING_TOKEN / azure-identity + # api_key: os.environ/AZURE_AI_API_KEY litellm_settings: callbacks: ["websearch_interception"] diff --git a/tests/search_tests/test_bing_grounding_search.py b/tests/search_tests/test_bing_grounding_search.py index 00e5e382eef..3d1737477a1 100644 --- a/tests/search_tests/test_bing_grounding_search.py +++ b/tests/search_tests/test_bing_grounding_search.py @@ -175,7 +175,19 @@ class TestBingGroundingSearchTransformation: assert mock_post.call_args.kwargs["json"]["tools"] == [{"type": "web_search"}] assert len(response.results) == 2 - def test_bing_grounding_search_tracks_cost(self, monkeypatch: pytest.MonkeyPatch): + def test_web_search_mode_is_not_billed_the_g1_price(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + with patch( # test-quality-ok: litellm.search has no client injection seam + "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", + return_value=_mock_response(), + ): + response = litellm.search(query="pricing check", search_provider="bing_grounding") + + assert response._hidden_params["response_cost"] == 0.0 + + def test_connection_mode_tracks_the_g1_cost(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_CONNECTION_ID", "conn-id") monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) with patch( # test-quality-ok: litellm.search has no client injection seam diff --git a/tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py b/tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py index 25dfa5fbe2b..d33c9f92df6 100644 --- a/tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py +++ b/tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py @@ -55,15 +55,18 @@ def test_ui_friendly_name(): assert _config().ui_friendly_name() == "Grounding with Bing Search" -def test_validate_environment_with_explicit_key(): - headers = _config().validate_environment({}, api_key="explicit-token") - assert headers["Authorization"] == "Bearer explicit-token" +def test_validate_environment_api_key_uses_api_key_header_not_bearer(): + headers = _config().validate_environment({}, api_key="azure-api-key") + assert headers["api-key"] == "azure-api-key" + assert "Authorization" not in headers assert headers["Content-Type"] == "application/json" def test_validate_environment_reads_env_token(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("BING_GROUNDING_TOKEN", "env-token") - assert _config().validate_environment({})["Authorization"] == "Bearer env-token" + headers = _config().validate_environment({}) + assert headers["Authorization"] == "Bearer env-token" + assert "api-key" not in headers def test_validate_environment_falls_back_to_entra_minter(): @@ -74,8 +77,9 @@ def test_validate_environment_falls_back_to_entra_minter(): def test_validate_environment_api_key_beats_env_token(monkeypatch: pytest.MonkeyPatch): monkeypatch.setenv("BING_GROUNDING_TOKEN", "env-token") minter = Mock(return_value="entra-token") - headers = _config(entra_token_minter=minter).validate_environment({}, api_key="explicit-token") - assert headers["Authorization"] == "Bearer explicit-token" + headers = _config(entra_token_minter=minter).validate_environment({}, api_key="azure-api-key") + assert headers["api-key"] == "azure-api-key" + assert "Authorization" not in headers minter.assert_not_called() @@ -265,6 +269,53 @@ def test_transform_search_response_malformed_body_raises_instead_of_reporting_em _config().transform_search_response(_resp(body, status_code=502), logging_obj=Mock()) +def test_transform_search_response_caps_results_to_max_results(): + annotations = [_citation(f"https://example.com/{i}", f"T{i}", 0, 5) for i in range(5)] + resp = _config().transform_search_response( + _resp(_message_response("claim", annotations)), logging_obj=Mock(), optional_params={"max_results": 2} + ) + assert [r.url for r in resp.results] == ["https://example.com/0", "https://example.com/1"] + + +def test_transform_search_response_without_max_results_returns_all_citations(): + annotations = [_citation(f"https://example.com/{i}", f"T{i}", 0, 5) for i in range(4)] + resp = _config().transform_search_response(_resp(_message_response("c", annotations)), logging_obj=Mock()) + assert len(resp.results) == 4 + + +def test_transform_search_response_failed_status_raises_with_error_message(): + payload = {"output": [], "status": "failed", "error": {"message": "content was filtered"}} + with pytest.raises(Exception, match="content was filtered") as excinfo: + _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert excinfo.value.status_code == 502 + + +def test_transform_search_response_incomplete_with_no_results_raises_with_reason(): + payload = {"output": [], "status": "incomplete", "incomplete_details": {"reason": "max_output_tokens"}} + with pytest.raises(Exception, match="incomplete: max_output_tokens"): + _config().transform_search_response(_resp(payload), logging_obj=Mock()) + + +def test_transform_search_response_incomplete_with_partial_results_returns_them(): + payload = _message_response("claim", [_citation("https://example.com", "Example", 0, 5)]) + payload["status"] = "incomplete" + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert [r.url for r in resp.results] == ["https://example.com"] + + +def test_transform_search_response_web_search_mode_zeroes_per_query_cost(): + payload = _message_response("claim", [_citation("https://example.com", "Example", 0, 5)]) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert resp._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.0 + + +def test_transform_search_response_connection_mode_leaves_price_to_cost_map(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("BING_GROUNDING_CONNECTION_ID", "conn-id") + payload = _message_response("claim", [_citation("https://example.com", "Example", 0, 5)]) + resp = _config().transform_search_response(_resp(payload), logging_obj=Mock()) + assert "additional_headers" not in resp._hidden_params + + def test_get_error_class_attributes_the_provider(): error = _config().get_error_class(error_message="quota exceeded", status_code=429, headers={}) assert error.status_code == 429