mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
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:
parent
7d3d655e20
commit
6dadd73bb8
2 changed files with 48 additions and 5 deletions
|
|
@ -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,
|
||||
},
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue