diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 597d1cfb863..37ca989b8d3 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -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, diff --git a/tests/test_litellm/rerank_api/test_main.py b/tests/test_litellm/rerank_api/test_main.py index 62149c742d6..2b6cfeda2c2 100644 --- a/tests/test_litellm/rerank_api/test_main.py +++ b/tests/test_litellm/rerank_api/test_main.py @@ -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."""