mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
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:
parent
848a3edbe8
commit
d59fcda8af
2 changed files with 45 additions and 6 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue