fix(rerank): adopt declared authenticating providers in arerank instead of resolving them

get_llm_provider runs the OAuth device flow for github_copilot and chatgpt,
so calling it on the event loop before the executor dispatch let an
authenticated caller block the loop for the length of the polling window.
Adopt the declared provider via declared_authenticating_provider, matching
the metadata callers in utils.py, and only resolve for everything else.
This commit is contained in:
mateo-berri 2026-09-01 14:47:59 -07:00
parent 848a3edbe8
commit d59fcda8af
2 changed files with 45 additions and 6 deletions

View file

@ -6,6 +6,7 @@ from typing import Any, Final, Literal
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler
@ -43,16 +44,22 @@ async def arerank(
"""
Async: Reranks a list of documents based on their relevance to the query
"""
_custom_llm_provider: str | None = None # rebind-ok: set by the get_llm_provider unpack; read in the except
_custom_llm_provider: str | None = (
None # rebind-ok: set by the declared-provider guard or the get_llm_provider unpack; read in the except
)
try:
loop: Final = asyncio.get_event_loop()
kwargs["arerank"] = True
_, _custom_llm_provider, _, _ = litellm.get_llm_provider( # rebind-ok: see pre-declaration above
model=model,
custom_llm_provider=custom_llm_provider,
api_base=kwargs.get("api_base", None),
)
declared_provider: Final = declared_authenticating_provider(model, custom_llm_provider)
if declared_provider is not None:
_custom_llm_provider = declared_provider # rebind-ok: see pre-declaration above
else:
_, _custom_llm_provider, _, _ = litellm.get_llm_provider( # rebind-ok: see pre-declaration above
model=model,
custom_llm_provider=custom_llm_provider,
api_base=kwargs.get("api_base", None),
)
func: Final = partial(
rerank,

View file

@ -172,6 +172,38 @@ async def test_arerank_error_is_mapped_to_litellm_exception(respx_mock: respx.Mo
assert "None - " not in str(exc_info.value)
@pytest.mark.asyncio
async def test_arerank_declared_authenticating_provider_skips_resolution(monkeypatch):
"""Regression for the event-loop hazard in arerank's provider pre-resolution:
get_llm_provider runs the blocking OAuth device flow for github_copilot/chatgpt,
so arerank must adopt the declared provider instead of resolving it, while the
except path still maps with that declared provider."""
from litellm.llms.base_llm.chat.transformation import BaseLLMException
resolution_calls = []
def record_resolution(*args, **kwargs):
resolution_calls.append((args, kwargs))
return "gpt-4o", "github_copilot", None, None
def rerank_raises_provider_error(*args, **kwargs):
raise BaseLLMException(status_code=401, message='{"error":"bad key"}')
monkeypatch.setattr(litellm, "get_llm_provider", record_resolution)
monkeypatch.setattr("litellm.rerank_api.main.rerank", rerank_raises_provider_error)
with pytest.raises(litellm.AuthenticationError) as exc_info:
await litellm.arerank(
model="github_copilot/gpt-4o",
query=MARKER_QUERY,
documents=[MARKER_DOC],
)
assert resolution_calls == []
assert "Github_copilotException" in str(exc_info.value)
assert "None - " not in str(exc_info.value)
@pytest.mark.asyncio
async def test_together_rerank_async_honors_env_api_base(respx_mock: respx.MockRouter, monkeypatch):
"""Regression: TOGETHER_AI_API_BASE was honored by chat but ignored by rerank."""