scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36
litellm/tests/router_unit_tests/test_router_endpoints.py
yuneng-jiang 6a0d03914c
test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 
* test: drop the cwd-relative sys.path.insert calls from the test suite

TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.

Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.

Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.

* test: drop the duplicate imports the sys.path sweep exposed to F811

* test(pre-call-utils): restore the os import the new bedrock tests need
2026-08-22 09:25:58 -07:00

1381 lines
42 KiB
Python
Raw Permalink Blame History

import os
import json
import traceback
from typing import Optional
from dotenv import load_dotenv
from fastapi import Request
from datetime import datetime
from unittest.mock import AsyncMock, patch, MagicMock
from litellm import Router, CustomLogger
from litellm.types.utils import StandardLoggingPayload
## Get the current directory of the file being run
pwd = os.path.dirname(os.path.realpath(__file__))
print(pwd)
file_path = os.path.join(pwd, "gettysburg.wav")
audio_file = open(file_path, "rb")
from pathlib import Path
import litellm
import pytest
import asyncio
@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.",
},
},
]
# This file includes the custom callbacks for LiteLLM Proxy
# Once defined, these can be passed in proxy_config.yaml
class MyCustomHandler(CustomLogger):
def __init__(self):
self.openai_client = None
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
try:
# init logging config
print("logging a transcript kwargs: ", kwargs)
print("openai client=", kwargs.get("client"))
self.openai_client = kwargs.get("client")
self.standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
"standard_logging_object"
)
except Exception:
pass
# Set litellm.callbacks = [proxy_handler_instance] on the proxy
@pytest.mark.asyncio
@pytest.mark.flaky(retries=6, delay=10)
async def test_transcription_on_router():
proxy_handler_instance = MyCustomHandler()
litellm.set_verbose = True
litellm.callbacks = [proxy_handler_instance]
print("\n Testing async transcription on router\n")
try:
model_list = [
{
"model_name": "whisper",
"litellm_params": {
"model": "whisper-1",
},
},
{
"model_name": "whisper",
"litellm_params": {
"model": "azure/azure-whisper",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com/",
"api_key": os.getenv("AZURE_EUROPE_API_KEY"),
"api_version": "2024-02-15-preview",
},
},
]
router = Router(model_list=model_list)
router_level_clients = []
for deployment in router.model_list:
_deployment_openai_client = router._get_client(
deployment=deployment,
kwargs={"model": "whisper-1"},
client_type="async",
)
router_level_clients.append(str(_deployment_openai_client))
## test 1: user facing function
response = await router.atranscription(
model="whisper",
file=audio_file,
)
## test 2: underlying function
response = await router._atranscription(
model="whisper",
file=audio_file,
)
print(response)
# PROD Test
# Ensure we ONLY use OpenAI/Azure client initialized on the router level
await asyncio.sleep(5)
print("OpenAI Client used= ", proxy_handler_instance.openai_client)
print("all router level clients= ", router_level_clients)
assert proxy_handler_instance.openai_client in router_level_clients
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize("mode", ["iterator"]) # "file",
@pytest.mark.asyncio
async def test_audio_speech_router(mode):
litellm.set_verbose = True
test_logger = MyCustomHandler()
litellm.callbacks = [test_logger]
from litellm import Router
client = Router(
model_list=[
{
"model_name": "tts",
"litellm_params": {
"model": "openai/tts-1",
},
},
]
)
response = await client.aspeech(
model="tts",
voice="alloy",
input="the quick brown fox jumped over the lazy dogs",
api_base=None,
api_key=None,
organization=None,
project=None,
max_retries=1,
timeout=600,
client=None,
optional_params={},
)
await asyncio.sleep(3)
from litellm.llms.openai.openai import HttpxBinaryResponseContent
assert isinstance(response, HttpxBinaryResponseContent)
assert test_logger.standard_logging_object is not None
print(
"standard_logging_object=",
json.dumps(test_logger.standard_logging_object, indent=4),
)
assert test_logger.standard_logging_object["model_group"] == "tts"
@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_rerank_endpoint(model_list):
from litellm.types.utils import RerankResponse
router = Router(model_list=model_list)
## Test 1: user facing function
response = await router.arerank(
model="cohere-rerank",
query="hello",
documents=["hello", "world"],
top_n=3,
)
## Test 2: underlying function
response = await router._arerank(
model="cohere-rerank",
query="hello",
documents=["hello", "world"],
top_n=3,
)
print("async re rank response: ", response)
assert response.id is not None
assert response.results is not None
RerankResponse.model_validate(response)
@pytest.mark.asyncio()
@pytest.mark.parametrize(
"model", ["omni-moderation-latest", "openai/omni-moderation-latest", None]
)
async def test_moderation_endpoint(model):
litellm.set_verbose = True
router = Router(
model_list=[
{
"model_name": "openai/*",
"litellm_params": {
"model": "openai/*",
},
},
{
"model_name": "*",
"litellm_params": {
"model": "openai/*",
},
},
]
)
if model is None:
response = await router.amoderation(input="hello this is a test")
else:
response = await router.amoderation(model=model, input="hello this is a test")
print("moderation response: ", response)
@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:
## Test 1: user facing function
response = await router.atext_completion(
model="gpt-5-mini",
prompt="Hello, how are you?",
mock_response="I'm fine, thank you!",
)
## Test 2: underlying function
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
"""
# Create a mock for the _aadapter_completion method
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},
}
# Create a router with a patched _aadapter_completion method
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",
},
}
]
)
# Replace the async_function_with_fallbacks with a mock
router.async_function_with_fallbacks = AsyncMock(return_value=mock_response)
# Call the aadapter_completion method
response = await router.aadapter_completion(
adapter_id="test-adapter-id",
model="test-adapter-model",
prompt="This is a test prompt",
max_tokens=100,
)
# Verify the response
assert response == mock_response
# Verify async_function_with_fallbacks was called with the right parameters
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
"""
# Create a mock response for litellm.aadapter_completion
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},
}
# Create a router with a mocked litellm.aadapter_completion
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",
},
}
]
)
# Mock the async_get_available_deployment method
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",
},
}
)
# Mock the async_routing_strategy_pre_call_checks method
router.async_routing_strategy_pre_call_checks = AsyncMock()
# Call the _aadapter_completion method
response = await router._aadapter_completion(
adapter_id="test-adapter-id",
model="test-adapter-model",
prompt="This is a test prompt",
max_tokens=100,
)
# Verify the response
assert response == mock_response
# Verify litellm.aadapter_completion was called with the right parameters
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
# Verify the success call was recorded
assert router.success_calls["test-model"] == 1
assert router.total_calls["test-model"] == 1
# Verify async_routing_strategy_pre_call_checks was called
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
"""
# Create a router with a basic model
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "anthropic/test-model",
"api_key": "fake-api-key",
},
}
]
)
# Explicitly call initialize_router_endpoints
router.initialize_router_endpoints()
# Verify all expected endpoints are initialized
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")
# Verify the endpoints are callable
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
"""
# 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",
},
}
]
)
# Mock the original function
mock_original_function = AsyncMock(return_value={"status": "success"})
# Call without custom_llm_provider
result = await router._init_vector_store_api_endpoints(
original_function=mock_original_function, vector_store_id="test-store"
)
# Verify original function was called with correct kwargs
mock_original_function.assert_called_once_with(vector_store_id="test-store")
assert result == {"status": "success"}
# Reset the mock
mock_original_function.reset_mock()
# Call with custom_llm_provider
await router._init_vector_store_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
vector_store_id="test-store",
)
# Verify custom_llm_provider was added to kwargs
mock_original_function.assert_called_once_with(
vector_store_id="test-store", custom_llm_provider="openai"
)
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.
"""
# Test with fresh router
router = Router()
initial_optional_callbacks = router.optional_callbacks
# Test that the method runs without error
result = router.apply_default_settings()
# Verify method returns None as expected
assert result is None
# Verify that optional_callbacks remains None if it was initially None
# (since default_pre_call_checks is an empty list)
assert router.optional_callbacks == initial_optional_callbacks
# Test with router that already has some optional_callbacks
router_with_callbacks = Router()
mock_callback = MagicMock()
router_with_callbacks.optional_callbacks = [mock_callback]
# Apply default settings
result = router_with_callbacks.apply_default_settings()
# Verify method returns None
assert result is None
# Verify existing callbacks are preserved (since we're adding empty list)
assert mock_callback in router_with_callbacks.optional_callbacks
# Test that the method is called during router initialization
with patch.object(Router, "apply_default_settings") as mock_apply:
Router()
mock_apply.assert_called_once()
# Test with mocked add_optional_pre_call_checks to verify internal call
router_test = Router()
with patch.object(router_test, "add_optional_pre_call_checks") as mock_add_checks:
router_test.apply_default_settings()
# Verify add_optional_pre_call_checks was called with empty list
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"
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()
Reference in a new issue View git blame Copy permalink