mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
e0398455d0
* test: migrate selected llm translation contracts * test: migrate anthropic completion translation tests * test: migrate bedrock and cohere translation tests * test: tighten bedrock translation migration * test: migrate azure and openai model translation tests * test: migrate openai and small-provider translation tests * test: migrate gemini and vertex translation tests * test: migrate optional params and rerank translation tests * test: migrate huggingface, json and openrouter translation tests * test: fix openrouter integration expectations * test: migrate databricks, xai, cometapi and fireworks translation tests * test: migrate ovhcloud, prompt and small provider translation tests * test: split reasoning effort grid into unit and e2e * test: migrate router and realtime guardrail translation tests * test: migrate google interactions translation tests * test: migrate a2a, evals, skills and hosted vllm translation tests * test: remove remaining llm translation suite * test: tighten migrated llm translation unit tests and land missing rows Restores rows the first pass dropped or weakened (bedrock agentcore and openwebui thinking stream, anthropic prefix, citations and tool_choice, azure agents handler, huggingface bare api_base, optional params num_retries, alternating roles, rerank response type), replaces truthiness and call-count checks with whole-value assertions, and moves tests to the mirror file of the code they cover * test: move cohere, databricks, fireworks and gemini rows to mirror files over mocked HTTP Restores three TestCohereV2Transform methods the first pass nested out of collection, replaces direct transform calls with respx-backed completion calls, and drops the databricks and fireworks files that sat outside their mirror paths * test: close lost-intent gaps found in review of the llm translation migration Points the hosted-mock guard at the triton test's new integration path, restores exact rerank result assertions for infinity, drops the reasoning e2e fallback that passed for any multi-line answer, and adds the mixtral invoke system prompt variants, eval deletion, and async pdf url download that the legacy suite exercised * test: assert provider acceptance in the reasoning effort e2e The live half of the split grid proves each route accepts reasoning_effort. claude-fable-5-1 on the direct route answers without surfacing reasoning content, so a reasoning-presence check either fails on head or needs a fallback that passes for any answer. Wire shape stays covered by the unit grid * test: cover async and stream-end paths the legacy translation suite exercised Adds async variants for bedrock converse file and thinking content, bedrock embedding timeout and region resolution, interaction stream termination and success logging, xai reasoning token folding, and bytes chunks with a prefix before data:. Also scopes the vertex safety settings global in test_vertex_safety_settings with monkeypatch so it stops leaking into gemini request bodies on the same worker * test: type the new llm translation e2e calls and mark moved fixture dirs as packages * test: name the reasoning effort e2e models as module constants * test: assert reasoning content, register the new gemini coverage ids and give the guardrail config e2e a unique name * test: block the prompt the shared guardrail denies and drop the reasoning presence check * test: drop the duplicate litellm import left by the main sync --------- Co-authored-by: yuneng <yuneng@berri.ai>
416 lines
14 KiB
Python
416 lines
14 KiB
Python
import asyncio
|
|
import importlib
|
|
import json
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_infinity_rerank():
|
|
mock_response = AsyncMock()
|
|
|
|
def return_val():
|
|
return {
|
|
"id": "cmpl-mockid",
|
|
"results": [{"index": 0, "relevance_score": 0.95}],
|
|
"usage": {"prompt_tokens": 100, "total_tokens": 150},
|
|
}
|
|
|
|
mock_response.json = return_val
|
|
mock_response.headers = {"key": "value"}
|
|
mock_response.status_code = 200
|
|
|
|
expected_payload = {
|
|
"model": "rerank-model",
|
|
"query": "hello",
|
|
"top_n": 3,
|
|
"documents": ["hello", "world"],
|
|
}
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
response = await litellm.arerank(
|
|
model="infinity/rerank-model",
|
|
query="hello",
|
|
documents=["hello", "world"],
|
|
top_n=3,
|
|
api_base="https://api.infinity.ai",
|
|
)
|
|
|
|
print("async re rank response: ", response)
|
|
|
|
# Assert
|
|
mock_post.assert_called_once()
|
|
print("call args", mock_post.call_args)
|
|
args_to_api = mock_post.call_args.kwargs["data"]
|
|
_url = mock_post.call_args.kwargs["url"]
|
|
print("Arguments passed to API=", args_to_api)
|
|
print("url = ", _url)
|
|
assert _url == "https://api.infinity.ai/rerank"
|
|
|
|
request_data = json.loads(args_to_api)
|
|
assert request_data["query"] == expected_payload["query"]
|
|
assert request_data["documents"] == expected_payload["documents"]
|
|
assert request_data["top_n"] == expected_payload["top_n"]
|
|
assert request_data["model"] == expected_payload["model"]
|
|
|
|
assert response.id is not None
|
|
assert response.results is not None
|
|
assert response.meta["tokens"]["input_tokens"] == 100
|
|
assert response.meta["tokens"]["output_tokens"] == 50 # total_tokens - prompt_tokens
|
|
|
|
assert response.id == "cmpl-mockid"
|
|
assert response.results == [{"index": 0, "relevance_score": 0.95}]
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_infinity_rerank_with_return_documents():
|
|
mock_response = AsyncMock()
|
|
|
|
mock_response = AsyncMock()
|
|
|
|
def return_val():
|
|
return {
|
|
"id": "cmpl-mockid",
|
|
"results": [{"index": 0, "relevance_score": 0.95, "document": "hello"}],
|
|
"usage": {"prompt_tokens": 100, "total_tokens": 150},
|
|
}
|
|
|
|
mock_response.json = return_val
|
|
mock_response.headers = {"key": "value"}
|
|
mock_response.status_code = 200
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
response = await litellm.arerank(
|
|
model="infinity/rerank-model",
|
|
query="hello",
|
|
documents=["hello", "world"],
|
|
top_n=3,
|
|
return_documents=True,
|
|
api_base="https://api.infinity.ai",
|
|
)
|
|
assert response.results[0]["document"] == {"text": "hello"}
|
|
assert response.id == "cmpl-mockid"
|
|
assert response.results == [{"index": 0, "relevance_score": 0.95, "document": {"text": "hello"}}]
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_infinity_rerank_with_env(monkeypatch):
|
|
# Set up mock response
|
|
mock_response = AsyncMock()
|
|
|
|
def return_val():
|
|
return {
|
|
"id": "cmpl-mockid",
|
|
"results": [{"index": 0, "relevance_score": 0.95}],
|
|
"usage": {"prompt_tokens": 100, "total_tokens": 150},
|
|
}
|
|
|
|
mock_response.json = return_val
|
|
mock_response.headers = {"key": "value"}
|
|
mock_response.status_code = 200
|
|
|
|
# Set environment variable
|
|
monkeypatch.setenv("INFINITY_API_BASE", "https://env.infinity.ai")
|
|
|
|
expected_payload = {
|
|
"model": "rerank-model",
|
|
"query": "hello",
|
|
"top_n": 3,
|
|
"documents": ["hello", "world"],
|
|
}
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
response = await litellm.arerank(
|
|
model="infinity/rerank-model",
|
|
query="hello",
|
|
documents=["hello", "world"],
|
|
top_n=3,
|
|
)
|
|
|
|
print("async re rank response: ", response)
|
|
|
|
# Assert
|
|
mock_post.assert_called_once()
|
|
print("call args", mock_post.call_args)
|
|
args_to_api = mock_post.call_args.kwargs["data"]
|
|
_url = mock_post.call_args.kwargs["url"]
|
|
print("Arguments passed to API=", args_to_api)
|
|
print("url = ", _url)
|
|
assert _url == "https://env.infinity.ai/rerank"
|
|
|
|
request_data = json.loads(args_to_api)
|
|
assert request_data["query"] == expected_payload["query"]
|
|
assert request_data["documents"] == expected_payload["documents"]
|
|
assert request_data["top_n"] == expected_payload["top_n"]
|
|
assert request_data["model"] == expected_payload["model"]
|
|
|
|
assert response.id is not None
|
|
assert response.results is not None
|
|
assert response.meta["tokens"]["input_tokens"] == 100
|
|
assert response.meta["tokens"]["output_tokens"] == 50 # total_tokens - prompt_tokens
|
|
|
|
assert response.id == "cmpl-mockid"
|
|
assert response.results == [{"index": 0, "relevance_score": 0.95}]
|
|
|
|
|
|
#### Embedding Tests
|
|
@pytest.mark.asyncio()
|
|
async def test_infinity_embedding():
|
|
mock_response = AsyncMock()
|
|
|
|
def return_val():
|
|
return {
|
|
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}],
|
|
"usage": {"prompt_tokens": 100, "total_tokens": 150},
|
|
"model": "custom-model/embedding-v1",
|
|
"object": "list",
|
|
}
|
|
|
|
mock_response.json = return_val
|
|
mock_response.headers = {"key": "value"}
|
|
mock_response.status_code = 200
|
|
|
|
expected_payload = {
|
|
"model": "custom-model/embedding-v1",
|
|
"input": ["hello world"],
|
|
"encoding_format": "float",
|
|
"output_dimension": 512,
|
|
}
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
response = await litellm.aembedding(
|
|
model="infinity/custom-model/embedding-v1",
|
|
input=["hello world"],
|
|
dimensions=512,
|
|
encoding_format="float",
|
|
api_base="https://api.infinity.ai/embeddings",
|
|
)
|
|
|
|
# Assert
|
|
mock_post.assert_called_once()
|
|
print("call args", mock_post.call_args)
|
|
request_data = mock_post.call_args.kwargs["json"]
|
|
_url = mock_post.call_args.kwargs["url"]
|
|
assert _url == "https://api.infinity.ai/embeddings"
|
|
|
|
assert request_data["input"] == expected_payload["input"]
|
|
assert request_data["model"] == expected_payload["model"]
|
|
assert request_data["output_dimension"] == expected_payload["output_dimension"]
|
|
assert request_data["encoding_format"] == expected_payload["encoding_format"]
|
|
|
|
assert response.data is not None
|
|
assert response.usage.prompt_tokens == 100
|
|
assert response.usage.total_tokens == 150
|
|
assert response.model == "custom-model/embedding-v1"
|
|
assert response.object == "list"
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_infinity_embedding_with_env(monkeypatch):
|
|
# Set up mock response
|
|
mock_response = AsyncMock()
|
|
|
|
def return_val():
|
|
return {
|
|
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}],
|
|
"usage": {"prompt_tokens": 100, "total_tokens": 150},
|
|
"model": "custom-model/embedding-v1",
|
|
"object": "list",
|
|
}
|
|
|
|
mock_response.json = return_val
|
|
mock_response.headers = {"key": "value"}
|
|
mock_response.status_code = 200
|
|
|
|
expected_payload = {
|
|
"model": "custom-model/embedding-v1",
|
|
"input": ["hello world"],
|
|
"encoding_format": "float",
|
|
"output_dimension": 512,
|
|
}
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
response = await litellm.aembedding(
|
|
model="infinity/custom-model/embedding-v1",
|
|
input=["hello world"],
|
|
dimensions=512,
|
|
encoding_format="float",
|
|
api_base="https://api.infinity.ai/embeddings",
|
|
)
|
|
|
|
# Assert
|
|
mock_post.assert_called_once()
|
|
print("call args", mock_post.call_args)
|
|
request_data = mock_post.call_args.kwargs["json"]
|
|
_url = mock_post.call_args.kwargs["url"]
|
|
assert _url == "https://api.infinity.ai/embeddings"
|
|
|
|
assert request_data["input"] == expected_payload["input"]
|
|
assert request_data["model"] == expected_payload["model"]
|
|
assert request_data["output_dimension"] == expected_payload["output_dimension"]
|
|
assert request_data["encoding_format"] == expected_payload["encoding_format"]
|
|
|
|
assert response.data is not None
|
|
assert response.usage.prompt_tokens == 100
|
|
assert response.usage.total_tokens == 150
|
|
assert response.model == "custom-model/embedding-v1"
|
|
assert response.object == "list"
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_infinity_embedding_extra_params():
|
|
mock_response = AsyncMock()
|
|
|
|
def return_val():
|
|
return {
|
|
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}],
|
|
"usage": {"prompt_tokens": 100, "total_tokens": 150},
|
|
"model": "custom-model/embedding-v1",
|
|
"object": "list",
|
|
}
|
|
|
|
mock_response.json = return_val
|
|
mock_response.headers = {"key": "value"}
|
|
mock_response.status_code = 200
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
response = await litellm.aembedding(
|
|
model="infinity/custom-model/embedding-v1",
|
|
input=["test input"],
|
|
dimensions=512,
|
|
encoding_format="float",
|
|
modality="text",
|
|
api_base="https://api.infinity.ai/embeddings",
|
|
)
|
|
|
|
mock_post.assert_called_once()
|
|
request_data = mock_post.call_args.kwargs["json"]
|
|
|
|
# Assert the request parameters
|
|
assert request_data["input"] == ["test input"]
|
|
assert request_data["model"] == "custom-model/embedding-v1"
|
|
assert request_data["output_dimension"] == 512
|
|
assert request_data["encoding_format"] == "float"
|
|
assert request_data["modality"] == "text"
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_infinity_embedding_prompt_token_mapping():
|
|
mock_response = AsyncMock()
|
|
|
|
def return_val():
|
|
return {
|
|
"data": [{"embedding": [0.1, 0.2, 0.3], "index": 0}],
|
|
"usage": {"total_tokens": 1, "prompt_tokens": 1},
|
|
"model": "custom-model/embedding-v1",
|
|
"object": "list",
|
|
}
|
|
|
|
mock_response.json = return_val
|
|
mock_response.headers = {"key": "value"}
|
|
mock_response.status_code = 200
|
|
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
response = await litellm.aembedding(
|
|
model="infinity/custom-model/embedding-v1",
|
|
input=["a"],
|
|
dimensions=512,
|
|
encoding_format="float",
|
|
api_base="https://api.infinity.ai/embeddings",
|
|
)
|
|
|
|
mock_post.assert_called_once()
|
|
# Assert the response
|
|
assert response.usage.prompt_tokens == 1
|
|
assert response.usage.total_tokens == 1
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _vcr_outcome_gate(request, vcr):
|
|
install_live_call_probe(request, vcr)
|
|
yield
|
|
record_vcr_outcome(request, vcr)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def event_loop():
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
loop = asyncio.new_event_loop()
|
|
yield loop
|
|
loop.close()
|
|
|
|
|
|
@pytest.fixture(scope="function", autouse=True)
|
|
def setup_and_teardown(event_loop):
|
|
import litellm
|
|
|
|
original_state = {}
|
|
for attr in (
|
|
"callbacks",
|
|
"success_callback",
|
|
"failure_callback",
|
|
"_async_success_callback",
|
|
"_async_failure_callback",
|
|
):
|
|
if hasattr(litellm, attr):
|
|
val = getattr(litellm, attr)
|
|
original_state[attr] = val.copy() if val else []
|
|
for attr in _SCALAR_DEFAULTS:
|
|
if hasattr(litellm, attr):
|
|
original_state[attr] = getattr(litellm, attr)
|
|
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
|
|
|
asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
|
|
importlib.reload(litellm)
|
|
asyncio.set_event_loop(event_loop)
|
|
yield
|
|
for attr, original_value in original_state.items():
|
|
if hasattr(litellm, attr):
|
|
setattr(litellm, attr, original_value)
|
|
pending = asyncio.all_tasks(event_loop)
|
|
for task in pending:
|
|
task.cancel()
|
|
if pending:
|
|
event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
|
|
|
|
|
|
_SCALAR_DEFAULTS = {
|
|
"num_retries": getattr(litellm, "num_retries", None),
|
|
"set_verbose": getattr(litellm, "set_verbose", False),
|
|
"cache": getattr(litellm, "cache", None),
|
|
"allowed_fails": getattr(litellm, "allowed_fails", 3),
|
|
"disable_aiohttp_transport": getattr(litellm, "disable_aiohttp_transport", False),
|
|
"force_ipv4": getattr(litellm, "force_ipv4", False),
|
|
"drop_params": getattr(litellm, "drop_params", None),
|
|
"modify_params": getattr(litellm, "modify_params", False),
|
|
"api_base": getattr(litellm, "api_base", None),
|
|
"api_key": getattr(litellm, "api_key", None),
|
|
"cohere_key": getattr(litellm, "cohere_key", None),
|
|
}
|