litellm/tests/unit/test_router/test_router_endpoints.py
yuneng-jiang 4846cfabc7
fix(vector-stores): keep proxy context out of provider vector store metadata (#47021)
* 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
2026-10-17 12:51:47 -07:00

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()