mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
290 lines
11 KiB
Python
290 lines
11 KiB
Python
import logging
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
|
|
|
|
import litellm
|
|
|
|
MARKER_QUERY = "MARKER_QUERY_do_not_log_at_info"
|
|
MARKER_DOC = "MARKER_DOC_sensitive_customer_text"
|
|
|
|
|
|
def _mock_cohere_response() -> MagicMock:
|
|
mock_response = MagicMock()
|
|
|
|
def return_val():
|
|
return {
|
|
"id": "cmpl-mockid",
|
|
"results": [{"index": 0, "relevance_score": 0.95}],
|
|
"meta": {
|
|
"api_version": {"version": "1.0"},
|
|
"billed_units": {"search_units": 1},
|
|
},
|
|
}
|
|
|
|
mock_response.json = return_val
|
|
mock_response.headers = {"key": "value"}
|
|
mock_response.status_code = 200
|
|
return mock_response
|
|
|
|
|
|
def test_rerank_does_not_log_request_content_at_info(caplog):
|
|
"""Regression for #32525: rerank must not emit query/documents to logs at INFO.
|
|
|
|
The mapped ``optional_rerank_params`` (which always contains ``query`` and
|
|
``documents``) bypasses ``turn_off_message_logging`` / ``redact_messages``,
|
|
so logging it at INFO leaks raw request content into stdout and any log sink.
|
|
"""
|
|
litellm.cohere_key = "test_api_key"
|
|
caplog.set_level(logging.DEBUG, logger="LiteLLM")
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post",
|
|
return_value=_mock_cohere_response(),
|
|
):
|
|
litellm.rerank(
|
|
model="cohere/rerank-english-v3.0",
|
|
query=MARKER_QUERY,
|
|
documents=[MARKER_DOC, "unrelated"],
|
|
top_n=2,
|
|
)
|
|
|
|
litellm_records = [r for r in caplog.records if r.name == "LiteLLM"]
|
|
|
|
info_or_above = [
|
|
r.getMessage()
|
|
for r in litellm_records
|
|
if r.levelno >= logging.INFO and (MARKER_QUERY in r.getMessage() or MARKER_DOC in r.getMessage())
|
|
]
|
|
assert not info_or_above, f"rerank leaked request content at INFO+: {info_or_above}"
|
|
|
|
optional_params_logs = [r for r in litellm_records if "optional_rerank_params" in r.getMessage()]
|
|
assert optional_params_logs, "expected the optional_rerank_params line to be logged"
|
|
assert all(
|
|
r.levelno == logging.DEBUG for r in optional_params_logs
|
|
), "optional_rerank_params must be logged at DEBUG, not INFO"
|
|
|
|
|
|
TOGETHER_RERANK_BODY = {
|
|
"id": "rerank-mock-id",
|
|
"results": [{"index": 0, "relevance_score": 0.95}],
|
|
"usage": {"prompt_tokens": 10, "total_tokens": 10},
|
|
}
|
|
|
|
|
|
def test_together_rerank_defaults_to_together_ai_host(respx_mock: respx.MockRouter, monkeypatch):
|
|
"""Regression for the Together host migration: rerank used to hardcode
|
|
https://api.together.xyz/v1/rerank. The default must now be api.together.ai."""
|
|
monkeypatch.delenv("TOGETHER_AI_API_BASE", raising=False)
|
|
|
|
mock_route = respx_mock.post("https://api.together.ai/v1/rerank")
|
|
mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
|
|
|
|
response = litellm.rerank(
|
|
model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
|
|
query=MARKER_QUERY,
|
|
documents=[MARKER_DOC],
|
|
api_key="fake-together-key",
|
|
)
|
|
|
|
assert mock_route.called
|
|
assert response.results[0]["relevance_score"] == 0.95
|
|
|
|
|
|
def test_together_rerank_honors_api_base(respx_mock: respx.MockRouter):
|
|
"""Regression: a custom api_base was silently ignored by the Together rerank handler."""
|
|
mock_route = respx_mock.post("https://custom-together.example/v1/rerank")
|
|
mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
|
|
|
|
litellm.rerank(
|
|
model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
|
|
query=MARKER_QUERY,
|
|
documents=[MARKER_DOC],
|
|
api_key="fake-together-key",
|
|
api_base="https://custom-together.example/v1",
|
|
)
|
|
|
|
assert mock_route.called
|
|
assert mock_route.calls[0].request.headers["authorization"] == "Bearer fake-together-key"
|
|
|
|
|
|
DASHSCOPE_RERANK_BODY = {
|
|
"object": "list",
|
|
"results": [{"index": 0, "relevance_score": 0.95}],
|
|
"model": "qwen3-rerank",
|
|
"id": "rerank-mock-id",
|
|
"usage": {"total_tokens": 10},
|
|
}
|
|
|
|
|
|
def test_dashscope_rerank_defaults_to_live_rerank_route(respx_mock: respx.MockRouter, monkeypatch):
|
|
"""Regression for the dead default endpoint: get_llm_provider always returns the
|
|
chat base for dashscope, which used to hijack rerank onto the dead
|
|
/compatible-mode/v1/reranks route."""
|
|
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
|
|
monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False)
|
|
|
|
mock_route = respx_mock.post("https://dashscope.aliyuncs.com/compatible-api/v1/reranks")
|
|
mock_route.return_value = httpx.Response(200, json=DASHSCOPE_RERANK_BODY)
|
|
|
|
response = litellm.rerank(
|
|
model="dashscope/qwen3-rerank",
|
|
query=MARKER_QUERY,
|
|
documents=[MARKER_DOC],
|
|
api_key="fake-dashscope-key",
|
|
)
|
|
|
|
assert mock_route.called
|
|
assert response.results[0]["relevance_score"] == 0.95
|
|
|
|
|
|
def test_dashscope_rerank_chat_env_base_keeps_host_and_rerank_route(respx_mock: respx.MockRouter, monkeypatch):
|
|
"""Regression: a chat-style DASHSCOPE_API_BASE must not hijack rerank onto the
|
|
chat path, while its host (the region) is preserved."""
|
|
monkeypatch.setenv("DASHSCOPE_API_BASE", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1")
|
|
monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False)
|
|
|
|
mock_route = respx_mock.post("https://dashscope-intl.aliyuncs.com/compatible-api/v1/reranks")
|
|
mock_route.return_value = httpx.Response(200, json=DASHSCOPE_RERANK_BODY)
|
|
|
|
litellm.rerank(
|
|
model="dashscope/qwen3-rerank",
|
|
query=MARKER_QUERY,
|
|
documents=[MARKER_DOC],
|
|
api_key="fake-dashscope-key",
|
|
)
|
|
|
|
assert mock_route.called
|
|
|
|
|
|
def test_dashscope_rerank_explicit_api_base_wins(respx_mock: respx.MockRouter, monkeypatch):
|
|
monkeypatch.setenv("DASHSCOPE_API_BASE", "https://dashscope-intl.aliyuncs.com/compatible-mode/v1")
|
|
|
|
mock_route = respx_mock.post("https://custom-rerank.example/v1/reranks")
|
|
mock_route.return_value = httpx.Response(200, json=DASHSCOPE_RERANK_BODY)
|
|
|
|
litellm.rerank(
|
|
model="dashscope/qwen3-rerank",
|
|
query=MARKER_QUERY,
|
|
documents=[MARKER_DOC],
|
|
api_key="fake-dashscope-key",
|
|
api_base="https://custom-rerank.example/v1",
|
|
)
|
|
|
|
assert mock_route.called
|
|
|
|
|
|
DASHSCOPE_404_BODY = {
|
|
"error": {
|
|
"message": "The model `does-not-exist` does not exist or you do not have access to it.",
|
|
"type": "invalid_request_error",
|
|
"param": None,
|
|
"code": "model_not_found",
|
|
},
|
|
"request_id": "mock-request-id",
|
|
}
|
|
|
|
|
|
def test_rerank_error_names_provider_and_keeps_body(respx_mock: respx.MockRouter, monkeypatch):
|
|
"""Regression for the rerank error path mapping with the unresolved provider param:
|
|
a provider 404 surfaced as 'None - ' instead of naming the provider and its error body."""
|
|
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
|
|
monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False)
|
|
|
|
mock_route = respx_mock.post("https://dashscope.example/v1/reranks")
|
|
mock_route.return_value = httpx.Response(404, json=DASHSCOPE_404_BODY)
|
|
|
|
with pytest.raises(litellm.NotFoundError) as exc_info:
|
|
litellm.rerank(
|
|
model="dashscope/does-not-exist",
|
|
query=MARKER_QUERY,
|
|
documents=[MARKER_DOC],
|
|
api_key="fake-dashscope-key",
|
|
api_base="https://dashscope.example/v1",
|
|
)
|
|
|
|
assert mock_route.called
|
|
assert "DashscopeException" in str(exc_info.value)
|
|
assert "does not exist or you do not have access to it" in str(exc_info.value)
|
|
assert "None - " not in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_arerank_error_is_mapped_to_litellm_exception(respx_mock: respx.MockRouter, monkeypatch):
|
|
"""Regression for arerank's bare re-raise: provider errors escaped as raw
|
|
provider exception classes instead of the mapped litellm exception contract."""
|
|
monkeypatch.delenv("DASHSCOPE_API_BASE", raising=False)
|
|
monkeypatch.delenv("DASHSCOPE_API_BASE_RERANK", raising=False)
|
|
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
|
|
|
mock_route = respx_mock.post("https://dashscope.example/v1/reranks")
|
|
mock_route.return_value = httpx.Response(404, json=DASHSCOPE_404_BODY)
|
|
|
|
with pytest.raises(litellm.NotFoundError) as exc_info:
|
|
await litellm.arerank(
|
|
model="dashscope/does-not-exist",
|
|
query=MARKER_QUERY,
|
|
documents=[MARKER_DOC],
|
|
api_key="fake-dashscope-key",
|
|
api_base="https://dashscope.example/v1",
|
|
)
|
|
|
|
assert mock_route.called
|
|
assert "DashscopeException" in str(exc_info.value)
|
|
assert "does not exist or you do not have access to it" in str(exc_info.value)
|
|
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."""
|
|
monkeypatch.setenv("TOGETHER_AI_API_BASE", "https://env-together.example/v1")
|
|
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
|
|
|
mock_route = respx_mock.post("https://env-together.example/v1/rerank")
|
|
mock_route.return_value = httpx.Response(200, json=TOGETHER_RERANK_BODY)
|
|
|
|
response = await litellm.arerank(
|
|
model="together_ai/mixedbread-ai/mxbai-rerank-large-v2",
|
|
query=MARKER_QUERY,
|
|
documents=[MARKER_DOC],
|
|
api_key="fake-together-key",
|
|
)
|
|
|
|
assert mock_route.called
|
|
assert response.results[0]["relevance_score"] == 0.95
|