mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
* fix(vector-stores): keep proxy context out of provider vector store metadata The proxy merged its request context (user_api_key_* fields, endpoint, user_agent, requester_ip_address, litellm_api_version) and the router's deployment info into the vector store request's metadata. For vector stores that field is also OpenAI's own metadata, which is capped at 16 keys, so a store created with a team key was full on creation and every later metadata update returned 400. Vector store routes now keep that context in litellm_metadata, like batches, files and responses, and the router does the same for avector_store_create. Provider-side metadata holds only what the caller sent. Adds db_scripts/strip_vector_store_internal_metadata.py, an opt-in script that removes the proxy-written keys from stores created before this change. It is a dry run unless --apply is passed. * fix(vector-stores): page the metadata cleanup script without recursion and tolerate non-string values Listing every store recursed once per page, which overflows the interpreter stack on large projects. Pages are now walked iteratively. Metadata values are no longer required to be strings, and a malformed response exits 2 with a fixed message instead of a traceback that could echo metadata values. * fix(vector-stores): read proxy agent id, tags and spend_logs_metadata from litellm_metadata Vector store routes now keep proxy context in litellm_metadata, but the v3 limiter only read agent_id from metadata, and logging let provider metadata tags or spend_logs_metadata replace the proxy's values. A string tags value also crashed get_request_tags. The limiter now prefers the proxy-resolved agent_id in litellm_metadata. Logging combines list tags from both buckets (metadata first, deduplicated) and merges dict spend_logs_metadata with litellm_metadata winning, ignoring values of the wrong type. * fix(vector-stores): narrow the logging fallback to mistyped caller values The tag union and spend_logs_metadata merge changed logged tags and spend_logs_metadata on every route that has both metadata buckets, deduped tags quadratically, and stopped background interaction polling from keeping string tags as sent. Caller values now win exactly as before, and the proxy's litellm_metadata tags or spend_logs_metadata only replace a caller value of the wrong type. The limiter reads agent ids through a pydantic model instead of a cast and ignores non-string ids. * fix(vector-stores): log the proxy's full tag list on litellm_metadata routes On routes that keep proxy context in litellm_metadata, the proxy builds litellm_metadata.tags from the caller's tags plus key, team and header tags, but get_request_tags preferred the caller's metadata.tags. A caller sending metadata.tags on a vector store call dropped the key's tags from request_tags, which tag spend tracking reads. get_request_tags now prefers a non-empty litellm_metadata tag list, and falls back to metadata.tags as before. * fix(vector-stores): strip LiteLLM-written keys at the provider boundary instead of moving metadata buckets Moving vector store proxy context into litellm_metadata broke shared readers (agent limits, request tags, merged logging metadata), and fixing them widened the change to every litellm_metadata route. This reverts the bucket move, the router change and the reader fixes, so proxy context stays in metadata exactly as on main. Vector store create and update now drop the keys LiteLLM itself writes (user_api_key*, requester and deployment keys, agent ids) from the metadata sent to the provider, in the shared optional-param helper both paths use. The cleanup script reuses the same key list. * fix(vector-stores): pass non-string metadata keys through the provider metadata filter * fix(vector-stores): strip LiteLLM-written metadata keys only on proxy and router calls * fix(vector-stores): send the provider exactly the metadata the caller sent through the proxy * fix(vector-stores): snapshot caller metadata from the raw request body and keep it out of logged request bodies * fix(vector-stores): copy the parsed-body metadata when the raw request body cannot be reparsed
1261 lines
38 KiB
Python
1261 lines
38 KiB
Python
import asyncio
|
|
import json
|
|
import os
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
|
|
|
import httpx
|
|
import litellm
|
|
import pytest
|
|
from litellm import Router
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
|
|
|
|
|
@pytest.fixture
|
|
def model_list():
|
|
return [
|
|
{
|
|
"model_name": "gpt-5-mini",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-5.5",
|
|
"litellm_params": {
|
|
"model": "gpt-5.5",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt-image-1",
|
|
"litellm_params": {
|
|
"model": "gpt-image-1",
|
|
"api_key": os.getenv("OPENAI_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "cohere-rerank",
|
|
"litellm_params": {
|
|
"model": "cohere/rerank-english-v3.0",
|
|
"api_key": os.getenv("COHERE_API_KEY"),
|
|
},
|
|
},
|
|
{
|
|
"model_name": "claude-sonnet-4-5-20250929",
|
|
"litellm_params": {
|
|
"model": "gpt-5-mini",
|
|
"mock_response": "hi this is macintosh.",
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aspeech_fallbacks_on_deployment_failure():
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "tts-main",
|
|
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
|
|
},
|
|
{
|
|
"model_name": "tts-backup",
|
|
"litellm_params": {"model": "openai/tts-1-hd", "api_key": "fake-key"},
|
|
},
|
|
],
|
|
fallbacks=[{"tts-main": ["tts-backup"]}],
|
|
num_retries=0,
|
|
)
|
|
|
|
called_models = []
|
|
|
|
async def mock_aspeech(*args, **kwargs):
|
|
called_models.append(kwargs["model"])
|
|
if kwargs["model"] == "openai/tts-1":
|
|
raise litellm.InternalServerError(
|
|
message="deployment down",
|
|
llm_provider="openai",
|
|
model="tts-1",
|
|
)
|
|
return MagicMock()
|
|
|
|
with patch("litellm.aspeech", side_effect=mock_aspeech):
|
|
response = await router.aspeech(
|
|
model="tts-main",
|
|
input="the quick brown fox jumped over the lazy dogs",
|
|
voice="alloy",
|
|
)
|
|
|
|
assert response is not None
|
|
assert called_models == ["openai/tts-1", "openai/tts-1-hd"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aspeech_success_returns_response():
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "tts",
|
|
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
|
|
},
|
|
]
|
|
)
|
|
|
|
mock_response = MagicMock()
|
|
with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech:
|
|
response = await router.aspeech(
|
|
model="tts",
|
|
input="the quick brown fox jumped over the lazy dogs",
|
|
voice="alloy",
|
|
)
|
|
|
|
assert response is mock_response
|
|
mock_aspeech.assert_called_once()
|
|
assert mock_aspeech.call_args.kwargs["model"] == "openai/tts-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aspeech_sets_deployment_metadata():
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "tts",
|
|
"litellm_params": {"model": "openai/tts-1", "api_key": "fake-key"},
|
|
},
|
|
]
|
|
)
|
|
|
|
mock_response = MagicMock()
|
|
with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech:
|
|
response = await router._aspeech(
|
|
model="tts",
|
|
input="the quick brown fox jumped over the lazy dogs",
|
|
voice="alloy",
|
|
)
|
|
|
|
assert response is mock_response
|
|
metadata = mock_aspeech.call_args.kwargs["metadata"]
|
|
assert metadata["deployment"] == "openai/tts-1"
|
|
assert metadata["deployment_model_name"] == "tts"
|
|
assert metadata["model_info"]["id"] is not None
|
|
|
|
|
|
@pytest.mark.asyncio()
|
|
async def test_moderation_endpoint_with_api_base():
|
|
"""
|
|
Test that the moderation endpoint respects api_base configuration
|
|
"""
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
custom_api_base = "https://us.api.openai.com/v1"
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "openai/omni-moderation-latest",
|
|
"litellm_params": {
|
|
"model": "openai/omni-moderation-latest",
|
|
"api_base": custom_api_base,
|
|
"api_key": "test-key",
|
|
},
|
|
},
|
|
]
|
|
)
|
|
|
|
# Mock the OpenAI client to verify api_base is passed
|
|
with patch(
|
|
"litellm.main.openai_chat_completions.get_openai_client"
|
|
) as mock_get_client:
|
|
mock_client = AsyncMock()
|
|
mock_response = MagicMock()
|
|
mock_response.model_dump.return_value = {
|
|
"id": "modr-123",
|
|
"model": "omni-moderation-latest",
|
|
"results": [
|
|
{
|
|
"flagged": False,
|
|
"categories": {},
|
|
"category_scores": {},
|
|
"category_applied_input_types": {},
|
|
}
|
|
],
|
|
}
|
|
mock_client.moderations.create = AsyncMock(return_value=mock_response)
|
|
mock_get_client.return_value = mock_client
|
|
|
|
response = await router.amoderation(
|
|
model="openai/omni-moderation-latest", input="hello this is a test"
|
|
)
|
|
|
|
# Verify that get_openai_client was called with the custom api_base
|
|
mock_get_client.assert_called()
|
|
call_kwargs = mock_get_client.call_args.kwargs
|
|
assert (
|
|
call_kwargs.get("api_base") == custom_api_base
|
|
), f"Expected api_base to be {custom_api_base}, but got {call_kwargs.get('api_base')}"
|
|
|
|
print(f"✓ Moderation endpoint correctly uses api_base: {custom_api_base}")
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_aaaaatext_completion_endpoint(model_list, sync_mode):
|
|
router = Router(model_list=model_list)
|
|
|
|
if sync_mode:
|
|
response = router.text_completion(
|
|
model="gpt-5-mini",
|
|
prompt="Hello, how are you?",
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
else:
|
|
response = await router.atext_completion(
|
|
model="gpt-5-mini",
|
|
prompt="Hello, how are you?",
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
|
|
response_2 = await router._atext_completion(
|
|
model="gpt-5-mini",
|
|
prompt="Hello, how are you?",
|
|
mock_response="I'm fine, thank you!",
|
|
)
|
|
assert response_2.choices[0].text == "I'm fine, thank you!"
|
|
|
|
assert response.choices[0].text == "I'm fine, thank you!"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_with_empty_choices(model_list):
|
|
"""
|
|
https://github.com/BerriAI/litellm/issues/8306
|
|
"""
|
|
router = Router(model_list=model_list)
|
|
mock_response = litellm.ModelResponse(
|
|
choices=[],
|
|
usage=litellm.Usage(
|
|
prompt_tokens=10,
|
|
completion_tokens=10,
|
|
total_tokens=20,
|
|
),
|
|
model="gpt-5-mini",
|
|
object="chat.completion",
|
|
created=1723081200,
|
|
).model_dump()
|
|
response = await router.acompletion(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "Hello, how are you?"}],
|
|
mock_response=mock_response,
|
|
)
|
|
assert response is not None
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
def test_generic_api_call_with_fallbacks_basic(sync_mode):
|
|
"""
|
|
Test both the sync and async versions of generic_api_call_with_fallbacks with a basic successful call
|
|
"""
|
|
# Create a mock function that will be passed to generic_api_call_with_fallbacks
|
|
if sync_mode:
|
|
from unittest.mock import Mock
|
|
|
|
mock_function = Mock()
|
|
mock_function.__name__ = "test_function"
|
|
else:
|
|
mock_function = AsyncMock()
|
|
mock_function.__name__ = "test_function"
|
|
|
|
# Create a mock response
|
|
mock_response = {
|
|
"id": "resp_123456",
|
|
"role": "assistant",
|
|
"content": "This is a test response",
|
|
"model": "test-model",
|
|
"usage": {"input_tokens": 10, "output_tokens": 20},
|
|
}
|
|
mock_function.return_value = mock_response
|
|
|
|
# Create a router with a test model
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model-alias",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
# Call the appropriate generic_api_call_with_fallbacks method
|
|
if sync_mode:
|
|
response = router._generic_api_call_with_fallbacks(
|
|
model="test-model-alias",
|
|
original_function=mock_function,
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
max_tokens=100,
|
|
)
|
|
else:
|
|
response = asyncio.run(
|
|
router._ageneric_api_call_with_fallbacks(
|
|
model="test-model-alias",
|
|
original_function=mock_function,
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
max_tokens=100,
|
|
)
|
|
)
|
|
|
|
# Verify the mock function was called
|
|
mock_function.assert_called_once()
|
|
|
|
# Verify the response
|
|
assert response == mock_response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aadapter_completion():
|
|
"""
|
|
Test the aadapter_completion method which uses async_function_with_fallbacks
|
|
"""
|
|
|
|
mock_response = {
|
|
"id": "adapter_resp_123",
|
|
"object": "adapter.completion",
|
|
"created": 1677858242,
|
|
"model": "test-model-with-adapter",
|
|
"choices": [
|
|
{
|
|
"text": "This is a test adapter response",
|
|
"index": 0,
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
|
}
|
|
|
|
with patch.object(Router, "_aadapter_completion", new_callable=AsyncMock) as mock_method:
|
|
mock_method.return_value = mock_response
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-adapter-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router.async_function_with_fallbacks = AsyncMock(return_value=mock_response)
|
|
|
|
response = await router.aadapter_completion(
|
|
adapter_id="test-adapter-id",
|
|
model="test-adapter-model",
|
|
prompt="This is a test prompt",
|
|
max_tokens=100,
|
|
)
|
|
|
|
assert response == mock_response
|
|
|
|
router.async_function_with_fallbacks.assert_called_once()
|
|
call_kwargs = router.async_function_with_fallbacks.call_args.kwargs
|
|
assert call_kwargs["adapter_id"] == "test-adapter-id"
|
|
assert call_kwargs["model"] == "test-adapter-model"
|
|
assert call_kwargs["prompt"] == "This is a test prompt"
|
|
assert call_kwargs["max_tokens"] == 100
|
|
assert call_kwargs["original_function"] == router._aadapter_completion
|
|
assert "metadata" in call_kwargs
|
|
assert call_kwargs["metadata"]["model_group"] == "test-adapter-model"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test__aadapter_completion():
|
|
"""
|
|
Test the _aadapter_completion method directly
|
|
"""
|
|
|
|
mock_response = {
|
|
"id": "adapter_resp_123",
|
|
"object": "adapter.completion",
|
|
"created": 1677858242,
|
|
"model": "test-model-with-adapter",
|
|
"choices": [
|
|
{
|
|
"text": "This is a test adapter response",
|
|
"index": 0,
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
|
}
|
|
|
|
with patch("litellm.aadapter_completion", new_callable=AsyncMock) as mock_adapter_completion:
|
|
mock_adapter_completion.return_value = mock_response
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-adapter-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router.async_get_available_deployment = AsyncMock(
|
|
return_value={
|
|
"model_name": "test-adapter-model",
|
|
"litellm_params": {
|
|
"model": "test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
"model_info": {
|
|
"id": "test-unique-id",
|
|
},
|
|
}
|
|
)
|
|
|
|
router.async_routing_strategy_pre_call_checks = AsyncMock()
|
|
|
|
response = await router._aadapter_completion(
|
|
adapter_id="test-adapter-id",
|
|
model="test-adapter-model",
|
|
prompt="This is a test prompt",
|
|
max_tokens=100,
|
|
)
|
|
|
|
assert response == mock_response
|
|
|
|
mock_adapter_completion.assert_called_once()
|
|
call_kwargs = mock_adapter_completion.call_args.kwargs
|
|
assert call_kwargs["adapter_id"] == "test-adapter-id"
|
|
assert call_kwargs["model"] == "test-model"
|
|
assert call_kwargs["prompt"] == "This is a test prompt"
|
|
assert call_kwargs["max_tokens"] == 100
|
|
assert call_kwargs["api_key"] == "fake-api-key"
|
|
assert call_kwargs["caching"] == router.cache_responses
|
|
|
|
assert router.success_calls["test-model"] == 1
|
|
assert router.total_calls["test-model"] == 1
|
|
|
|
router.async_routing_strategy_pre_call_checks.assert_called_once()
|
|
|
|
|
|
def test_initialize_router_endpoints():
|
|
"""
|
|
Test that initialize_router_endpoints correctly sets up all router endpoints
|
|
"""
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router.initialize_router_endpoints()
|
|
|
|
assert hasattr(router, "amoderation")
|
|
assert hasattr(router, "aanthropic_messages")
|
|
assert hasattr(router, "aresponses")
|
|
assert hasattr(router, "responses")
|
|
assert hasattr(router, "aget_responses")
|
|
assert hasattr(router, "adelete_responses")
|
|
|
|
assert callable(router.amoderation)
|
|
assert callable(router.aanthropic_messages)
|
|
assert callable(router.aresponses)
|
|
assert callable(router.responses)
|
|
assert callable(router.aget_responses)
|
|
assert callable(router.adelete_responses)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_responses_api_endpoints():
|
|
"""
|
|
A simpler test for _init_responses_api_endpoints that focuses on the basic functionality
|
|
"""
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
|
|
# Create a router with a basic model
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
# Just mock the _ageneric_api_call_with_fallbacks method
|
|
router._ageneric_api_call_with_fallbacks = AsyncMock()
|
|
|
|
# Add a mock implementation of _get_model_id_from_response_id to the Router instance
|
|
ResponsesAPIRequestUtils.get_model_id_from_response_id = MagicMock(
|
|
return_value=None
|
|
)
|
|
|
|
# Call without a response_id (no model extraction should happen)
|
|
await router._init_responses_api_endpoints(
|
|
original_function=AsyncMock(), thread_id="thread_xyz"
|
|
)
|
|
|
|
# Verify _ageneric_api_call_with_fallbacks was called but model wasn't changed
|
|
first_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs
|
|
assert "model" not in first_call_kwargs
|
|
assert first_call_kwargs["thread_id"] == "thread_xyz"
|
|
|
|
# Reset the mock
|
|
router._ageneric_api_call_with_fallbacks.reset_mock()
|
|
|
|
# Change the return value for the second call
|
|
ResponsesAPIRequestUtils.get_model_id_from_response_id.return_value = (
|
|
"claude-3-sonnet"
|
|
)
|
|
|
|
# Call with a response_id
|
|
await router._init_responses_api_endpoints(
|
|
original_function=AsyncMock(), response_id="resp_claude_123"
|
|
)
|
|
|
|
# Verify model was updated in the kwargs
|
|
second_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs
|
|
assert second_call_kwargs["model"] == "claude-3-sonnet"
|
|
assert second_call_kwargs["response_id"] == "resp_claude_123"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_vector_store_api_endpoints():
|
|
"""
|
|
Test that _init_vector_store_api_endpoints correctly passes custom_llm_provider to kwargs
|
|
"""
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
mock_original_function = AsyncMock(return_value={"status": "success"})
|
|
|
|
result = await router._init_vector_store_api_endpoints(
|
|
original_function=mock_original_function, vector_store_id="test-store"
|
|
)
|
|
|
|
mock_original_function.assert_called_once_with(vector_store_id="test-store")
|
|
assert result == {"status": "success"}
|
|
|
|
mock_original_function.reset_mock()
|
|
|
|
await router._init_vector_store_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
vector_store_id="test-store",
|
|
)
|
|
|
|
mock_original_function.assert_called_once_with(vector_store_id="test-store", custom_llm_provider="openai")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_avector_store_create_sends_only_the_callers_metadata_to_the_provider():
|
|
def echo_vector_store(request: httpx.Request) -> httpx.Response:
|
|
body: Final = json.loads(request.content)
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"id": "vs_echo",
|
|
"object": "vector_store",
|
|
"created_at": 0,
|
|
"name": body["name"],
|
|
"usage_bytes": 0,
|
|
"file_counts": {"in_progress": 0, "completed": 0, "failed": 0, "cancelled": 0, "total": 0},
|
|
"status": "completed",
|
|
"metadata": body.get("metadata"),
|
|
},
|
|
)
|
|
|
|
router: Final = Router(
|
|
model_list=[{"model_name": "gpt-5", "litellm_params": {"model": "openai/gpt-5", "api_key": "sk-test"}}]
|
|
)
|
|
|
|
response: Final = await router.avector_store_create(
|
|
model="gpt-5",
|
|
name="x",
|
|
metadata={"team": "llmproxy"},
|
|
client=AsyncHTTPHandler(transport=httpx.MockTransport(echo_vector_store)),
|
|
)
|
|
|
|
assert response["id"] == "vs_echo"
|
|
assert response["metadata"] == {"team": "llmproxy"}
|
|
|
|
|
|
def test_apply_default_settings():
|
|
"""
|
|
Test the apply_default_settings method.
|
|
|
|
This test verifies that apply_default_settings correctly initializes
|
|
default pre-call checks and doesn't modify existing router state.
|
|
"""
|
|
|
|
router = Router()
|
|
initial_optional_callbacks = router.optional_callbacks
|
|
|
|
result = router.apply_default_settings()
|
|
|
|
assert result is None
|
|
|
|
assert router.optional_callbacks == initial_optional_callbacks
|
|
|
|
router_with_callbacks = Router()
|
|
mock_callback = MagicMock()
|
|
router_with_callbacks.optional_callbacks = [mock_callback]
|
|
|
|
result = router_with_callbacks.apply_default_settings()
|
|
|
|
assert result is None
|
|
|
|
assert mock_callback in router_with_callbacks.optional_callbacks
|
|
|
|
with patch.object(Router, "apply_default_settings") as mock_apply:
|
|
Router()
|
|
mock_apply.assert_called_once()
|
|
|
|
router_test = Router()
|
|
with patch.object(router_test, "add_optional_pre_call_checks") as mock_add_checks:
|
|
router_test.apply_default_settings()
|
|
|
|
mock_add_checks.assert_called_once_with([])
|
|
|
|
|
|
def test_initialize_core_endpoints():
|
|
"""
|
|
Test that _initialize_core_endpoints correctly sets up all core router endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_core_endpoints()
|
|
|
|
core_endpoints = [
|
|
"amoderation",
|
|
"aanthropic_messages",
|
|
"agenerate_content",
|
|
"aadapter_generate_content",
|
|
"aresponses",
|
|
"afile_delete",
|
|
"afile_content",
|
|
"responses",
|
|
"aget_responses",
|
|
"acancel_responses",
|
|
"adelete_responses",
|
|
"alist_input_items",
|
|
"_arealtime",
|
|
"acreate_fine_tuning_job",
|
|
"acancel_fine_tuning_job",
|
|
"alist_fine_tuning_jobs",
|
|
"aretrieve_fine_tuning_job",
|
|
"afile_list",
|
|
"aimage_edit",
|
|
"allm_passthrough_route",
|
|
]
|
|
|
|
for endpoint in core_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_specialized_endpoints():
|
|
"""
|
|
Test that _initialize_specialized_endpoints correctly sets up specialized endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_specialized_endpoints()
|
|
|
|
specialized_endpoints = [
|
|
"avector_store_search",
|
|
"avector_store_create",
|
|
"vector_store_search",
|
|
"vector_store_create",
|
|
"agenerate_content",
|
|
"generate_content",
|
|
"agenerate_content_stream",
|
|
"generate_content_stream",
|
|
"aocr",
|
|
"ocr",
|
|
"asearch",
|
|
"search",
|
|
"avideo_generation",
|
|
"video_generation",
|
|
"avideo_list",
|
|
"video_list",
|
|
"avideo_status",
|
|
"video_status",
|
|
"avideo_content",
|
|
"video_content",
|
|
"avideo_remix",
|
|
"video_remix",
|
|
"acreate_container",
|
|
"create_container",
|
|
"alist_containers",
|
|
"list_containers",
|
|
"aretrieve_container",
|
|
"retrieve_container",
|
|
"adelete_container",
|
|
"delete_container",
|
|
"acreate_skill",
|
|
"alist_skills",
|
|
"aget_skill",
|
|
"adelete_skill",
|
|
]
|
|
|
|
for endpoint in specialized_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_vector_store_endpoints():
|
|
"""
|
|
Test that _initialize_vector_store_endpoints correctly sets up vector store endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_vector_store_endpoints()
|
|
|
|
vector_store_endpoints = [
|
|
"avector_store_search",
|
|
"avector_store_create",
|
|
"vector_store_search",
|
|
"vector_store_create",
|
|
]
|
|
|
|
for endpoint in vector_store_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_vector_store_file_endpoints():
|
|
"""
|
|
Test that _initialize_vector_store_file_endpoints correctly sets up vector store file endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_vector_store_file_endpoints()
|
|
|
|
vector_store_file_endpoints = [
|
|
"avector_store_file_create",
|
|
"vector_store_file_create",
|
|
"avector_store_file_list",
|
|
"vector_store_file_list",
|
|
"avector_store_file_retrieve",
|
|
"vector_store_file_retrieve",
|
|
"avector_store_file_content",
|
|
"vector_store_file_content",
|
|
"avector_store_file_update",
|
|
"vector_store_file_update",
|
|
"avector_store_file_delete",
|
|
"vector_store_file_delete",
|
|
]
|
|
|
|
for endpoint in vector_store_file_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_google_genai_endpoints():
|
|
"""
|
|
Test that _initialize_google_genai_endpoints correctly sets up Google GenAI endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_google_genai_endpoints()
|
|
|
|
google_genai_endpoints = [
|
|
"agenerate_content",
|
|
"generate_content",
|
|
"agenerate_content_stream",
|
|
"generate_content_stream",
|
|
]
|
|
|
|
for endpoint in google_genai_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_ocr_search_endpoints():
|
|
"""
|
|
Test that _initialize_ocr_search_endpoints correctly sets up OCR and search endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_ocr_search_endpoints()
|
|
|
|
ocr_search_endpoints = [
|
|
"aocr",
|
|
"ocr",
|
|
"asearch",
|
|
"search",
|
|
]
|
|
|
|
for endpoint in ocr_search_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_video_endpoints():
|
|
"""
|
|
Test that _initialize_video_endpoints correctly sets up video endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_video_endpoints()
|
|
|
|
video_endpoints = [
|
|
"avideo_generation",
|
|
"video_generation",
|
|
"avideo_list",
|
|
"video_list",
|
|
"avideo_status",
|
|
"video_status",
|
|
"avideo_content",
|
|
"video_content",
|
|
"avideo_remix",
|
|
"video_remix",
|
|
]
|
|
|
|
for endpoint in video_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_container_endpoints():
|
|
"""
|
|
Test that _initialize_container_endpoints correctly sets up container endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "openai/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_container_endpoints()
|
|
|
|
container_endpoints = [
|
|
"acreate_container",
|
|
"create_container",
|
|
"alist_containers",
|
|
"list_containers",
|
|
"aretrieve_container",
|
|
"retrieve_container",
|
|
"adelete_container",
|
|
"delete_container",
|
|
]
|
|
|
|
for endpoint in container_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
def test_initialize_skills_endpoints():
|
|
"""
|
|
Test that _initialize_skills_endpoints correctly sets up skills endpoints.
|
|
"""
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "test-model",
|
|
"litellm_params": {
|
|
"model": "anthropic/test-model",
|
|
"api_key": "fake-api-key",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
router._initialize_skills_endpoints()
|
|
|
|
skills_endpoints = [
|
|
"acreate_skill",
|
|
"alist_skills",
|
|
"aget_skill",
|
|
"adelete_skill",
|
|
]
|
|
|
|
for endpoint in skills_endpoints:
|
|
assert hasattr(router, endpoint)
|
|
assert callable(getattr(router, endpoint))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints():
|
|
"""
|
|
Test that _init_containers_api_endpoints calls the original function
|
|
directly when there is no managed container ID (no embedded model_id).
|
|
"""
|
|
router = Router(model_list=[])
|
|
|
|
mock_response = {"id": "cntr_test", "name": "Test Container"}
|
|
mock_original_function = AsyncMock(return_value=mock_response)
|
|
|
|
result = await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
name="Test Container",
|
|
)
|
|
|
|
mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container")
|
|
assert result == mock_response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_managed_id_routes_via_generic_fallbacks():
|
|
"""
|
|
Managed ``cntr_`` IDs embed ``model_id``; router should decode and use
|
|
``_ageneric_api_call_with_fallbacks`` so deployment credentials apply.
|
|
"""
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "azure-router-model",
|
|
"litellm_params": {
|
|
"model": "azure/gpt-5.5",
|
|
"api_key": "fake-key",
|
|
"api_base": "https://westus.api.cognitive.microsoft.com",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
router._ageneric_api_call_with_fallbacks = AsyncMock()
|
|
|
|
managed_id = ResponsesAPIRequestUtils.build_container_id(
|
|
custom_llm_provider="azure",
|
|
model_id="azure-router-model",
|
|
container_id="cfile_upstream_abc",
|
|
)
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=AsyncMock(),
|
|
custom_llm_provider="openai",
|
|
container_id=managed_id,
|
|
file_id="cfile_xyz",
|
|
)
|
|
|
|
router._ageneric_api_call_with_fallbacks.assert_called_once()
|
|
call_kw = router._ageneric_api_call_with_fallbacks.call_args.kwargs
|
|
assert call_kw["model"] == "azure-router-model"
|
|
assert call_kw["container_id"] == "cfile_upstream_abc"
|
|
assert call_kw["file_id"] == "cfile_xyz"
|
|
assert call_kw["custom_llm_provider"] == "azure"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_managed_id_without_model_id_unwraps():
|
|
"""
|
|
Managed ``cntr_`` IDs may be encoded with an empty ``model_id`` (e.g. when a
|
|
streaming response had no router metadata). The router must still unwrap the
|
|
managed ID before calling the upstream provider — otherwise the raw
|
|
``cntr_...`` token leaks downstream and the provider rejects it.
|
|
"""
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
|
|
router = Router(model_list=[])
|
|
mock_original_function = AsyncMock(return_value={"ok": True})
|
|
|
|
managed_id = ResponsesAPIRequestUtils.build_container_id(
|
|
custom_llm_provider="openai",
|
|
model_id=None,
|
|
container_id="cfile_upstream_abc",
|
|
)
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
container_id=managed_id,
|
|
file_id="cfile_xyz",
|
|
)
|
|
|
|
mock_original_function.assert_called_once()
|
|
call_kw = mock_original_function.call_args.kwargs
|
|
assert call_kw["container_id"] == "cfile_upstream_abc"
|
|
assert call_kw["file_id"] == "cfile_xyz"
|
|
assert call_kw["custom_llm_provider"] == "openai"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_managed_id_without_model_id_applies_decoded_provider():
|
|
"""
|
|
A managed ``cntr_`` ID can encode a non-OpenAI provider (e.g. ``azure``) with
|
|
an empty ``model_id`` (streaming events without router ``model_info.id``).
|
|
The router must still apply the decoded provider so the request routes to
|
|
the correct upstream — not stay on the default ``openai``.
|
|
"""
|
|
from litellm.responses.utils import ResponsesAPIRequestUtils
|
|
|
|
router = Router(model_list=[])
|
|
mock_original_function = AsyncMock(return_value={"ok": True})
|
|
|
|
managed_id = ResponsesAPIRequestUtils.build_container_id(
|
|
custom_llm_provider="azure",
|
|
model_id=None,
|
|
container_id="cfile_upstream_abc",
|
|
)
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
container_id=managed_id,
|
|
file_id="cfile_xyz",
|
|
)
|
|
|
|
mock_original_function.assert_called_once()
|
|
call_kw = mock_original_function.call_args.kwargs
|
|
assert call_kw["container_id"] == "cfile_upstream_abc"
|
|
assert call_kw["file_id"] == "cfile_xyz"
|
|
assert call_kw["custom_llm_provider"] == "azure"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_create_with_model_uses_deployment_credentials(monkeypatch):
|
|
"""
|
|
``POST /v1/containers`` carries no container ID, so a ``model`` in the request
|
|
body is the only way to pick a deployment. The upstream call must receive that
|
|
deployment's ``api_key``/``api_base`` instead of falling back to the global
|
|
``OPENAI_API_KEY`` (which may be unset on the proxy).
|
|
"""
|
|
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-5.4",
|
|
"litellm_params": {
|
|
"model": "openai/gpt-5.4",
|
|
"api_key": "sk-model-list-key",
|
|
"api_base": "https://custom.openai.example/v1",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
mock_original_function = AsyncMock(return_value={"id": "cntr_test", "name": "Test Container"})
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
name="Test Container",
|
|
model="gpt-5.4",
|
|
)
|
|
|
|
mock_original_function.assert_called_once()
|
|
call_kw = mock_original_function.call_args.kwargs
|
|
assert call_kw["api_key"] == "sk-model-list-key"
|
|
assert call_kw["api_base"] == "https://custom.openai.example/v1"
|
|
assert call_kw["model"] == "openai/gpt-5.4"
|
|
assert call_kw["name"] == "Test Container"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_create_without_model_calls_directly():
|
|
"""
|
|
Without ``model`` (or with ``model=None`` as the proxy forwards it), create/list
|
|
must keep calling the handler directly with global provider credentials.
|
|
"""
|
|
router = Router(model_list=[])
|
|
router._ageneric_api_call_with_fallbacks = AsyncMock()
|
|
mock_original_function = AsyncMock(return_value={"id": "cntr_test"})
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
name="Test Container",
|
|
model=None,
|
|
)
|
|
|
|
router._ageneric_api_call_with_fallbacks.assert_not_called()
|
|
mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container", model=None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_init_containers_api_endpoints_create_with_unknown_model_passes_through(monkeypatch):
|
|
"""
|
|
A ``model`` that names no configured deployment must not turn into a 400. The call
|
|
falls through to the handler with the caller's model and no injected deployment
|
|
credentials, matching the behaviour before model-based routing existed.
|
|
"""
|
|
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-5.4",
|
|
"litellm_params": {"model": "openai/gpt-5.4", "api_key": "sk-model-list-key"},
|
|
}
|
|
]
|
|
)
|
|
mock_original_function = AsyncMock(return_value={"id": "cntr_test"})
|
|
|
|
await router._init_containers_api_endpoints(
|
|
original_function=mock_original_function,
|
|
custom_llm_provider="openai",
|
|
name="Test Container",
|
|
model="does-not-exist",
|
|
)
|
|
|
|
mock_original_function.assert_called_once()
|
|
call_kw = mock_original_function.call_args.kwargs
|
|
assert call_kw["model"] == "does-not-exist"
|
|
assert call_kw["name"] == "Test Container"
|
|
assert "api_key" not in call_kw
|
|
assert "api_base" not in call_kw
|
|
|
|
|
|
def test_router_model_group_encrypted_content_affinity_callback_registration():
|
|
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
|
|
DeploymentAffinityCheck,
|
|
)
|
|
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
|
|
EncryptedContentAffinityCheck,
|
|
)
|
|
|
|
model_group = "openai.gpt-5.1-codex"
|
|
model_group_affinity_config = {
|
|
model_group: ["encrypted_content_affinity"],
|
|
}
|
|
router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": model_group,
|
|
"litellm_params": {
|
|
"model": "openai/gpt-5.1-codex",
|
|
"api_key": "mock-api-key",
|
|
},
|
|
}
|
|
],
|
|
model_group_affinity_config=model_group_affinity_config,
|
|
num_retries=0,
|
|
)
|
|
|
|
try:
|
|
callbacks = router.optional_callbacks or []
|
|
encrypted_content_callbacks = [
|
|
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
|
|
]
|
|
deployment_callback = next(
|
|
cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)
|
|
)
|
|
assert len(encrypted_content_callbacks) == 1
|
|
assert encrypted_content_callbacks[0].enable_global_affinity is False
|
|
assert (
|
|
encrypted_content_callbacks[0].model_group_affinity_config
|
|
== model_group_affinity_config
|
|
)
|
|
assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(
|
|
deployment_callback
|
|
)
|
|
|
|
router._add_encrypted_content_affinity_check(enable_global_affinity=True)
|
|
|
|
callbacks = router.optional_callbacks or []
|
|
encrypted_content_callbacks = [
|
|
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
|
|
]
|
|
assert len(encrypted_content_callbacks) == 1
|
|
assert encrypted_content_callbacks[0].enable_global_affinity is True
|
|
assert encrypted_content_callbacks[0].router is router
|
|
finally:
|
|
router.discard()
|