scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

fix(search): harden bing_grounding auth, result cap, status, and cost

Browse source 
- 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
This commit is contained in:
mateo-berri 2026-08-24 12:26:07 -07:00
parent 671a454baa
commit a0c319878b
5 changed files with 190 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

130
litellm/llms/azure/search/transformation.py
View file

@ -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,

2
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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(

34
litellm/proxy/example_config_yaml/bing_grounding_websearch_config.yaml
View file

@ -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://<account>.services.ai.azure.com/api/projects/<project>
# 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"]

14
tests/search_tests/test_bing_grounding_search.py
View file

@ -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

63
tests/test_litellm/llms/azure/search/test_bing_grounding_search_transformation.py
View file

@ -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