fix(search): validate search provider environment off the event loop

async_search called validate_environment on the event loop, so a Grounding with
Bing Search request that mints an Entra token blocked the loop for the token
endpoint round trip. It now runs in a worker thread like the Responses handlers
This commit is contained in:
mateo-berri 2026-10-19 10:05:35 -07:00
parent 7d3d655e20
commit 6dadd73bb8
2 changed files with 48 additions and 5 deletions

View file

@ -2084,8 +2084,8 @@ class BaseLLMHTTPHandler:
if provider_config is None:
raise ValueError(f"No provider config found for provider: {custom_llm_provider}")
# Validate environment and get headers
headers = provider_config.validate_environment(
validated_headers: Final = await asyncio.to_thread(
provider_config.validate_environment,
api_key=api_key,
api_base=api_base,
headers=headers or {},
@ -2097,7 +2097,7 @@ class BaseLLMHTTPHandler:
optional_params=optional_params,
api_key=api_key,
api_base=api_base,
headers=headers or {},
headers=validated_headers,
)
# Get complete URL (pass data for providers that need request body for URL construction)
@ -2111,7 +2111,7 @@ class BaseLLMHTTPHandler:
signed_headers, signed_json_body = await sign_request_off_loop_if_aws(
provider_config,
provider_config.sign_request,
headers=headers,
headers=validated_headers,
optional_params=optional_params,
request_data=data,
api_base=complete_url,
@ -2125,7 +2125,7 @@ class BaseLLMHTTPHandler:
additional_args={
"complete_input_dict": data,
"api_base": complete_url,
"headers": headers,
"headers": validated_headers,
},
)

View file

@ -24,6 +24,7 @@ from litellm.integrations.code_interpreter_interception.handler import (
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.websearch_interception.handler import WebSearchInterceptionLogger
from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig
from litellm.llms.azure.search.transformation import BingGroundingSearchConfig
from litellm.llms.base_llm.audio_transcription.transformation import (
AudioTranscriptionRequestData,
BaseAudioTranscriptionConfig,
@ -135,6 +136,48 @@ async def test_get_search_preserves_successful_results(is_async: bool, has_resul
assert response.results[0].snippet == "Example snippet"
@pytest.mark.asyncio
async def test_async_search_mints_the_bing_grounding_token_off_the_event_loop(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("BING_GROUNDING_PROJECT_ENDPOINT", "https://acct.services.ai.azure.com/api/projects/proj")
monkeypatch.setenv("BING_GROUNDING_MODEL", "gpt-4.1")
monkeypatch.delenv("BING_GROUNDING_TOKEN", raising=False)
minting_threads: Final[list[int]] = []
def mint() -> str:
minting_threads.append(threading.get_ident())
return "minted-entra-token"
authorizations: Final[list[str | None]] = []
def foundry(request: httpx.Request) -> httpx.Response:
authorizations.append(request.headers.get("authorization"))
return httpx.Response(200, json={"output": []})
client: Final = AsyncHTTPHandler()
await client.close()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(foundry))
response: Final = await BaseLLMHTTPHandler().search(
query="test",
optional_params={},
timeout=5,
logging_obj=Mock(),
api_key=None,
api_base=None,
custom_llm_provider="bing_grounding",
client=client,
asearch=True,
provider_config=BingGroundingSearchConfig(entra_token_minter=mint),
)
assert response.object == "search"
assert authorizations == ["Bearer minted-entra-token"]
assert len(minting_threads) == 1
assert threading.get_ident() not in minting_threads
@pytest.mark.asyncio
@pytest.mark.parametrize("is_async", (False, True))
async def test_get_search_preserves_tinyfish_http_error_formatting(is_async: bool) -> None: