tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-13 15:07:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #19185 from BerriAI/litellm_fallback_endpoints_support

Browse Source 
[Feat] Add fallback endpoints support
This commit is contained in:
Sameer Kankute
2026-01-16 17:39:35 +05:30
committed by GitHub
parent ca06bb4e3a e54af42774
commit cca6632dae
6 changed files with 1220 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/proxy/fallback_management.md
+273
View File
@@ -0,0 +1,273 @@
# [New] Fallback Management Endpoints
Dedicated endpoints for managing model fallbacks separately from the general configuration.
## Overview
These endpoints allow you to configure, retrieve, and delete fallback models without modifying the entire proxy configuration. This provides a cleaner and safer way to manage fallbacks compared to using the `/config/update` endpoint.
## Prerequisites
- Database storage must be enabled: Set `STORE_MODEL_IN_DB=True` in your environment
- Models must exist in the router before configuring fallbacks
## Endpoints
### POST /fallback
Create or update fallbacks for a specific model.
**Request Body:**
```json
{
"model": "gpt-3.5-turbo",
"fallback_models": ["gpt-4", "claude-3-haiku"],
"fallback_type": "general"
}
```
**Parameters:**
- `model` (string, required): The primary model name to configure fallbacks for
- `fallback_models` (array of strings, required): List of fallback model names in priority order
- `fallback_type` (string, optional): Type of fallback. Options:
- `"general"` (default): Standard fallbacks for any error
- `"context_window"`: Fallbacks for context window exceeded errors
- `"content_policy"`: Fallbacks for content policy violations
**Response:**
```json
{
"model": "gpt-3.5-turbo",
"fallback_models": ["gpt-4", "claude-3-haiku"],
"fallback_type": "general",
"message": "Fallback configuration created successfully"
}
```
**Example using cURL:**
```bash
curl -X POST "http://localhost:4000/fallback" \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"model": "gpt-3.5-turbo",
"fallback_models": ["gpt-4", "claude-3-haiku"],
"fallback_type": "general"
}'
```
**Example using Python:**
```python
import requests
response = requests.post(
"http://localhost:4000/fallback",
headers={
"Authorization": "Bearer sk-1234",
"Content-Type": "application/json"
},
json={
"model": "gpt-3.5-turbo",
"fallback_models": ["gpt-4", "claude-3-haiku"],
"fallback_type": "general"
}
)
print(response.json())
```
### GET /fallback/{model}
Get fallback configuration for a specific model.
**Parameters:**
- `model` (path parameter, required): The model name to get fallbacks for
- `fallback_type` (query parameter, optional): Type of fallback to retrieve (default: "general")
**Response:**
```json
{
"model": "gpt-3.5-turbo",
"fallback_models": ["gpt-4", "claude-3-haiku"],
"fallback_type": "general"
}
```
**Example using cURL:**
```bash
curl -X GET "http://localhost:4000/fallback/gpt-3.5-turbo?fallback_type=general" \
-H "Authorization: Bearer sk-1234"
```
**Example using Python:**
```python
import requests
response = requests.get(
"http://localhost:4000/fallback/gpt-3.5-turbo",
headers={"Authorization": "Bearer sk-1234"},
params={"fallback_type": "general"}
)
print(response.json())
```
### DELETE /fallback/{model}
Delete fallback configuration for a specific model.
**Parameters:**
- `model` (path parameter, required): The model name to delete fallbacks for
- `fallback_type` (query parameter, optional): Type of fallback to delete (default: "general")
**Response:**
```json
{
"model": "gpt-3.5-turbo",
"fallback_type": "general",
"message": "Fallback configuration deleted successfully"
}
```
**Example using cURL:**
```bash
curl -X DELETE "http://localhost:4000/fallback/gpt-3.5-turbo?fallback_type=general" \
-H "Authorization: Bearer sk-1234"
```
**Example using Python:**
```python
import requests
response = requests.delete(
"http://localhost:4000/fallback/gpt-3.5-turbo",
headers={"Authorization": "Bearer sk-1234"},
params={"fallback_type": "general"}
)
print(response.json())
```
### Test fallback
</TabItem>
<TabItem value="proxy" label="PROXY">
```bash
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"model": "gpt-3.5-turbo",
"messages": [
{
"role": "user",
"content": "ping"
}
],
"mock_testing_fallbacks": true
}
'
```
</TabItem>
</Tabs>
## Validation
The endpoints perform the following validations:
1. **Model Existence**: Verifies that the primary model exists in the router
2. **Fallback Model Existence**: Ensures all fallback models exist in the router
3. **No Self-Fallback**: Prevents a model from being its own fallback
4. **No Duplicates**: Ensures no duplicate models in the fallback list
5. **Database Enabled**: Requires `STORE_MODEL_IN_DB=True` to be set
## Error Responses
### 400 Bad Request
```json
{
"detail": {
"error": "Invalid fallback models: ['non-existent-model']",
"available_models": ["gpt-3.5-turbo", "gpt-4", "claude-3-haiku"]
}
}
```
### 404 Not Found
```json
{
"detail": {
"error": "Model 'gpt-3.5-turbo' not found in router",
"available_models": ["gpt-4", "claude-3-haiku"]
}
}
```
### 500 Internal Server Error
```json
{
"detail": {
"error": "Router not initialized"
}
}
```
## Fallback Types Explained
### General Fallbacks
Used for any type of error that occurs during model invocation. This is the most common type of fallback.
**Use Case:** When a model is unavailable, rate-limited, or returns an error.
```json
{
"model": "gpt-3.5-turbo",
"fallback_models": ["gpt-4", "claude-3-haiku"],
"fallback_type": "general"
}
```
### Context Window Fallbacks
Specifically triggered when a context window exceeded error occurs.
**Use Case:** When the input is too long for the primary model, fallback to a model with a larger context window.
```json
{
"model": "gpt-3.5-turbo",
"fallback_models": ["gpt-4-32k", "claude-3-opus"],
"fallback_type": "context_window"
}
```
### Content Policy Fallbacks
Specifically triggered when content policy violations occur.
**Use Case:** When the primary model rejects content due to safety filters, fallback to a model with different content policies.
```json
{
"model": "gpt-4",
"fallback_models": ["claude-3-haiku"],
"fallback_type": "content_policy"
}
```
## Benefits Over /config/update
1. **Safety**: Only modifies fallback configuration, won't accidentally change other settings
2. **Simplicity**: Focused API with clear validation messages
3. **Granularity**: Manage fallbacks per model and per type
4. **Validation**: Comprehensive checks ensure configuration is valid before applying
5. **Clarity**: Clear error messages with available models listed
## Notes
- Fallbacks are triggered after the configured number of retries fails
- Fallbacks are attempted in the order specified in `fallback_models`
- The maximum number of fallbacks attempted is controlled by the router's `max_fallbacks` setting
- Changes take effect immediately and are persisted to the database
docs/my-website/sidebars.js
+1
View File
@@ -858,6 +858,7 @@ const sidebars = {
"proxy/load_balancing",
"proxy/provider_budget_routing",
"proxy/reliability",
"proxy/fallback_management",
"proxy/tag_routing",
"proxy/timeout",
"wildcard_routing"
litellm/proxy/management_endpoints/fallback_management_endpoints.py
+375
View File
@@ -0,0 +1,375 @@
"""
FALLBACK MANAGEMENT ENDPOINTS
Dedicated endpoints for managing model fallbacks separately from general config.
POST /fallback - Create or update fallbacks for a specific model
GET /fallback/{model} - Get fallbacks for a specific model
DELETE /fallback/{model} - Delete fallbacks for a specific model
"""
# pyright: reportMissingImports=false
import json
from typing import TYPE_CHECKING, Dict, List, Literal
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.model_checks import get_all_fallbacks
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
if TYPE_CHECKING:
from fastapi import APIRouter, Depends, HTTPException, status
else:
try:
from fastapi import APIRouter, Depends, HTTPException, status
except ImportError:
# fastapi is only required for proxy, not for SDK usage
pass
from litellm.types.management_endpoints.router_settings_endpoints import (
FallbackCreateRequest,
FallbackDeleteResponse,
FallbackGetResponse,
FallbackResponse,
)
router = APIRouter()
@router.post(
"/fallback",
tags=["Fallback Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=FallbackResponse,
status_code=status.HTTP_200_OK,
)
async def create_fallback(
data: FallbackCreateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Create or update fallbacks for a specific model.
This endpoint allows you to configure fallback models separately from the general config.
Fallbacks are triggered when a model call fails after retries.
**Example Request:**
```json
{
"model": "gpt-3.5-turbo",
"fallback_models": ["gpt-4", "claude-3-haiku"],
"fallback_type": "general"
}
```
**Fallback Types:**
- `general`: Standard fallbacks for any error (default)
- `context_window`: Fallbacks specifically for context window exceeded errors
- `content_policy`: Fallbacks specifically for content policy violations
"""
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_config,
store_model_in_db,
)
try:
# Validate that we have a router
if llm_router is None:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": "Router not initialized"},
)
# Validate that the model exists in the router
model_names = llm_router.model_names
if data.model not in model_names:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail={
"error": f"Model '{data.model}' not found in router",
"available_models": list(model_names),
},
)
# Validate that all fallback models exist in the router
invalid_fallback_models = [
m for m in data.fallback_models if m not in model_names
]
if invalid_fallback_models:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": f"Invalid fallback models: {invalid_fallback_models}",
"available_models": list(model_names),
},
)
# Check if fallback model is the same as the primary model
if data.model in data.fallback_models:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": f"Model '{data.model}' cannot be its own fallback"
},
)
# Check if we need to store in DB
if store_model_in_db is not True or prisma_client is None:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": "Database storage not enabled. Set 'STORE_MODEL_IN_DB=True' in your environment to use this feature."
},
)
# Load existing config
config = await proxy_config.get_config()
router_settings = config.get("router_settings", {})
# Get the appropriate fallback list based on type
fallback_key = "fallbacks"
if data.fallback_type == "context_window":
fallback_key = "context_window_fallbacks"
elif data.fallback_type == "content_policy":
fallback_key = "content_policy_fallbacks"
# Get existing fallbacks
existing_fallbacks: List[Dict[str, List[str]]] = router_settings.get(
fallback_key, []
)
# Update or add the fallback configuration
fallback_updated = False
for i, fallback_dict in enumerate(existing_fallbacks):
if data.model in fallback_dict:
# Update existing fallback
existing_fallbacks[i] = {data.model: data.fallback_models}
fallback_updated = True
break
if not fallback_updated:
# Add new fallback
existing_fallbacks.append({data.model: data.fallback_models})
# Update router settings
router_settings[fallback_key] = existing_fallbacks
# Save to database - convert router_settings to JSON string
router_settings_json = json.dumps(router_settings)
await prisma_client.db.litellm_config.upsert(
where={"param_name": "router_settings"},
data={
"create": {
"param_name": "router_settings",
"param_value": router_settings_json,
},
"update": {
"param_value": router_settings_json
},
},
)
# Update the in-memory router configuration
setattr(llm_router, fallback_key, existing_fallbacks)
verbose_proxy_logger.info(
f"Fallback configured: {data.model} -> {data.fallback_models} (type: {data.fallback_type})"
)
return FallbackResponse(
model=data.model,
fallback_models=data.fallback_models,
fallback_type=data.fallback_type,
message=f"Fallback configuration {'updated' if fallback_updated else 'created'} successfully",
)
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.error(f"Error creating fallback: {str(e)}", exc_info=True)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": f"Failed to create fallback: {str(e)}"},
)
@router.get(
"/fallback/{model}",
tags=["Fallback Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=FallbackGetResponse,
)
async def get_fallback(
model: str,
fallback_type: Literal["general", "context_window", "content_policy"] = "general",
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get fallback configuration for a specific model.
**Parameters:**
- `model`: The model name to get fallbacks for
- `fallback_type`: Type of fallback to retrieve (query parameter)
**Example:**
```
GET /fallback/gpt-3.5-turbo?fallback_type=general
```
"""
from litellm.proxy.proxy_server import llm_router
try:
if llm_router is None:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": "Router not initialized"},
)
# Get fallbacks using the existing utility function
fallback_models = get_all_fallbacks(
model=model, llm_router=llm_router, fallback_type=fallback_type
)
if not fallback_models:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail={
"error": f"No {fallback_type} fallbacks configured for model '{model}'"
},
)
return FallbackGetResponse(
model=model,
fallback_models=fallback_models,
fallback_type=fallback_type,
)
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.error(f"Error getting fallback: {str(e)}", exc_info=True)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": f"Failed to get fallback: {str(e)}"},
)
@router.delete(
"/fallback/{model}",
tags=["Fallback Management"],
dependencies=[Depends(user_api_key_auth)],
response_model=FallbackDeleteResponse,
)
async def delete_fallback(
model: str,
fallback_type: Literal["general", "context_window", "content_policy"] = "general",
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Delete fallback configuration for a specific model.
**Parameters:**
- `model`: The model name to delete fallbacks for
- `fallback_type`: Type of fallback to delete (query parameter)
**Example:**
```
DELETE /fallback/gpt-3.5-turbo?fallback_type=general
```
"""
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_config,
store_model_in_db,
)
try:
if llm_router is None:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": "Router not initialized"},
)
if store_model_in_db is not True or prisma_client is None:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail={
"error": "Database storage not enabled. Set 'STORE_MODEL_IN_DB=True' in your environment to use this feature."
},
)
# Load existing config
config = await proxy_config.get_config()
router_settings = config.get("router_settings", {})
# Get the appropriate fallback list based on type
fallback_key = "fallbacks"
if fallback_type == "context_window":
fallback_key = "context_window_fallbacks"
elif fallback_type == "content_policy":
fallback_key = "content_policy_fallbacks"
# Get existing fallbacks
existing_fallbacks: List[Dict[str, List[str]]] = router_settings.get(
fallback_key, []
)
# Find and remove the fallback configuration
fallback_found = False
updated_fallbacks = []
for fallback_dict in existing_fallbacks:
if model not in fallback_dict:
updated_fallbacks.append(fallback_dict)
else:
fallback_found = True
if not fallback_found:
raise HTTPException(
status_code=status.HTTP_404_NOT_FOUND,
detail={
"error": f"No {fallback_type} fallbacks configured for model '{model}'"
},
)
# Update router settings
router_settings[fallback_key] = updated_fallbacks
# Save to database - convert router_settings to JSON string
router_settings_json = json.dumps(router_settings)
await prisma_client.db.litellm_config.upsert(
where={"param_name": "router_settings"},
data={
"create": {
"param_name": "router_settings",
"param_value": router_settings_json,
},
"update": {
"param_value": router_settings_json
},
},
)
# Update the in-memory router configuration
setattr(llm_router, fallback_key, updated_fallbacks)
verbose_proxy_logger.info(
f"Fallback deleted: {model} (type: {fallback_type})"
)
return FallbackDeleteResponse(
model=model,
fallback_type=fallback_type,
message="Fallback configuration deleted successfully",
)
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.error(f"Error deleting fallback: {str(e)}", exc_info=True)
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail={"error": f"Failed to delete fallback: {str(e)}"},
)
litellm/proxy/proxy_server.py
+14 -3
View File
@@ -297,10 +297,15 @@ from litellm.proxy.management_endpoints.cost_tracking_settings import (
from litellm.proxy.management_endpoints.customer_endpoints import (
router as customer_router,
)
from litellm.proxy.management_endpoints.fallback_management_endpoints import (
router as fallback_management_router,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import (
router as internal_user_router,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
from litellm.proxy.management_endpoints.internal_user_endpoints import (
user_update,
)
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_verification_tokens,
duration_in_seconds,
@@ -354,7 +359,9 @@ from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router
from litellm.proxy.openai_files_endpoints.files_endpoints import (
router as openai_files_router,
)
from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config
from litellm.proxy.openai_files_endpoints.files_endpoints import (
set_files_config,
)
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
passthrough_endpoint_router,
)
@@ -449,7 +456,9 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
LiteLLM_UpperboundKeyGenerateParams,
)
from litellm.types.realtime import RealtimeQueryParams
from litellm.types.router import DeploymentTypedDict
from litellm.types.router import (
DeploymentTypedDict,
)
from litellm.types.router import ModelInfo as RouterModelInfo
from litellm.types.router import (
RouterGeneralSettings,
@@ -3261,6 +3270,7 @@ class ProxyConfig:
return None
import json
import yaml
# 1. Try key-level router_settings
@@ -10518,6 +10528,7 @@ app.include_router(model_access_group_management_router)
app.include_router(tag_management_router)
app.include_router(cost_tracking_settings_router)
app.include_router(router_settings_router)
app.include_router(fallback_management_router)
app.include_router(cache_settings_router)
app.include_router(user_agent_analytics_router)
app.include_router(enterprise_router)
litellm/types/management_endpoints/router_settings_endpoints.py
+63 -2
View File
@@ -2,9 +2,70 @@
Types and field definitions for router settings management endpoints
"""
from typing import Any, Dict, List, Optional
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel
from pydantic import BaseModel, Field, field_validator
# Fallback Management Types
class FallbackCreateRequest(BaseModel):
"""Request model for creating/updating fallbacks"""
model: str = Field(
description="The model name to configure fallbacks for (e.g., 'gpt-3.5-turbo')"
)
fallback_models: List[str] = Field(
description="List of fallback model names in order of priority",
min_length=1,
)
fallback_type: Literal["general", "context_window", "content_policy"] = Field(
default="general",
description="Type of fallback: 'general' (default), 'context_window', or 'content_policy'",
)
@field_validator("fallback_models")
@classmethod
def validate_fallback_models(cls, v: List[str]) -> List[str]:
if not v:
raise ValueError("fallback_models must contain at least one model")
if len(v) != len(set(v)):
raise ValueError("fallback_models must not contain duplicates")
return v
@field_validator("model")
@classmethod
def validate_model(cls, v: str) -> str:
if not v or not v.strip():
raise ValueError("model must be a non-empty string")
return v.strip()
class FallbackResponse(BaseModel):
"""Response model for fallback operations"""
model: str = Field(description="The model name")
fallback_models: List[str] = Field(description="List of fallback model names")
fallback_type: str = Field(description="Type of fallback")
message: str = Field(description="Success message")
class FallbackGetResponse(BaseModel):
"""Response model for getting fallbacks"""
model: str = Field(description="The model name")
fallback_models: List[str] = Field(description="List of fallback model names")
fallback_type: str = Field(description="Type of fallback")
class FallbackDeleteResponse(BaseModel):
"""Response model for deleting fallbacks"""
model: str = Field(description="The model name")
fallback_type: str = Field(description="Type of fallback")
message: str = Field(description="Success message")
# Router Settings Types
class RouterSettingsField(BaseModel):
tests/test_litellm/proxy/test_fallback_management_endpoints.py
+494
View File
@@ -0,0 +1,494 @@
"""
Tests for fallback management endpoints
Tests:
1. Create fallback configuration
2. Get fallback configuration
3. Delete fallback configuration
4. Validation tests (invalid models, duplicate fallbacks, etc.)
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from litellm.proxy.management_endpoints.fallback_management_endpoints import (
FallbackCreateRequest,
create_fallback,
delete_fallback,
get_fallback,
)
class TestFallbackCreateRequest:
"""Test the FallbackCreateRequest validation"""
def test_valid_request(self):
"""Test valid fallback request"""
request = FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["gpt-4", "claude-3-haiku"],
fallback_type="general",
)
assert request.model == "gpt-3.5-turbo"
assert request.fallback_models == ["gpt-4", "claude-3-haiku"]
assert request.fallback_type == "general"
def test_default_fallback_type(self):
"""Test default fallback type is 'general'"""
request = FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["gpt-4"],
)
assert request.fallback_type == "general"
def test_empty_fallback_models(self):
"""Test that empty fallback_models raises validation error"""
with pytest.raises(ValueError, match="at least 1 item"):
FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=[],
)
def test_duplicate_fallback_models(self):
"""Test that duplicate fallback models raise validation error"""
with pytest.raises(ValueError, match="fallback_models must not contain duplicates"):
FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["gpt-4", "gpt-4"],
)
def test_empty_model_name(self):
"""Test that empty model name raises validation error"""
with pytest.raises(ValueError, match="model must be a non-empty string"):
FallbackCreateRequest(
model="",
fallback_models=["gpt-4"],
)
def test_whitespace_model_name(self):
"""Test that whitespace-only model name raises validation error"""
with pytest.raises(ValueError, match="model must be a non-empty string"):
FallbackCreateRequest(
model=" ",
fallback_models=["gpt-4"],
)
def test_model_name_trimmed(self):
"""Test that model name is trimmed"""
request = FallbackCreateRequest(
model=" gpt-3.5-turbo ",
fallback_models=["gpt-4"],
)
assert request.model == "gpt-3.5-turbo"
def test_context_window_fallback_type(self):
"""Test context_window fallback type"""
request = FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["gpt-4-32k"],
fallback_type="context_window",
)
assert request.fallback_type == "context_window"
def test_content_policy_fallback_type(self):
"""Test content_policy fallback type"""
request = FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["gpt-4"],
fallback_type="content_policy",
)
assert request.fallback_type == "content_policy"
@pytest.mark.asyncio
class TestCreateFallback:
"""Test the create_fallback endpoint"""
@pytest.fixture
def mock_router(self):
"""Create a mock router"""
router = MagicMock()
router.model_names = {"gpt-3.5-turbo", "gpt-4", "claude-3-haiku"}
router.fallbacks = []
router.context_window_fallbacks = []
router.content_policy_fallbacks = []
return router
@pytest.fixture
def mock_prisma_client(self):
"""Create a mock prisma client"""
client = MagicMock()
client.db.litellm_config.upsert = AsyncMock()
client.jsonify_object = lambda x: x
return client
@pytest.fixture
def mock_proxy_config(self):
"""Create a mock proxy config"""
config = MagicMock()
config.get_config = AsyncMock(return_value={"router_settings": {}})
return config
@pytest.fixture
def mock_user_api_key_dict(self):
"""Create a mock user API key dict"""
return MagicMock()
async def test_create_fallback_success(
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
):
"""Test successful fallback creation"""
request = FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["gpt-4", "claude-3-haiku"],
fallback_type="general",
)
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router,
), patch(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
), patch(
"litellm.proxy.proxy_server.proxy_config",
mock_proxy_config,
), patch(
"litellm.proxy.proxy_server.store_model_in_db",
True,
):
response = await create_fallback(request, mock_user_api_key_dict)
assert response.model == "gpt-3.5-turbo"
assert response.fallback_models == ["gpt-4", "claude-3-haiku"]
assert response.fallback_type == "general"
assert "created" in response.message.lower() or "updated" in response.message.lower()
# Verify database was updated
mock_prisma_client.db.litellm_config.upsert.assert_called_once()
async def test_create_fallback_router_not_initialized(
self, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
):
"""Test error when router is not initialized"""
request = FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["gpt-4"],
)
with patch(
"litellm.proxy.proxy_server.llm_router",
None,
), pytest.raises(HTTPException) as exc_info:
await create_fallback(request, mock_user_api_key_dict)
assert exc_info.value.status_code == 500
assert "Router not initialized" in str(exc_info.value.detail)
async def test_create_fallback_model_not_found(
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
):
"""Test error when model is not found in router"""
request = FallbackCreateRequest(
model="invalid-model",
fallback_models=["gpt-4"],
)
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router,
), patch(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
), patch(
"litellm.proxy.proxy_server.store_model_in_db",
True,
), pytest.raises(HTTPException) as exc_info:
await create_fallback(request, mock_user_api_key_dict)
assert exc_info.value.status_code == 404
assert "not found in router" in str(exc_info.value.detail)
async def test_create_fallback_invalid_fallback_model(
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
):
"""Test error when fallback model is not found in router"""
request = FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["invalid-fallback-model"],
)
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router,
), patch(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
), patch(
"litellm.proxy.proxy_server.store_model_in_db",
True,
), pytest.raises(HTTPException) as exc_info:
await create_fallback(request, mock_user_api_key_dict)
assert exc_info.value.status_code == 400
assert "Invalid fallback models" in str(exc_info.value.detail)
async def test_create_fallback_model_is_own_fallback(
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
):
"""Test error when model is its own fallback"""
request = FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["gpt-3.5-turbo", "gpt-4"],
)
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router,
), patch(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
), patch(
"litellm.proxy.proxy_server.store_model_in_db",
True,
), pytest.raises(HTTPException) as exc_info:
await create_fallback(request, mock_user_api_key_dict)
assert exc_info.value.status_code == 400
assert "cannot be its own fallback" in str(exc_info.value.detail)
async def test_create_fallback_db_not_enabled(
self, mock_router, mock_user_api_key_dict
):
"""Test error when database storage is not enabled"""
request = FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["gpt-4"],
)
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router,
), patch(
"litellm.proxy.proxy_server.store_model_in_db",
False,
), pytest.raises(HTTPException) as exc_info:
await create_fallback(request, mock_user_api_key_dict)
assert exc_info.value.status_code == 400
assert "Database storage not enabled" in str(exc_info.value.detail)
async def test_create_fallback_context_window_type(
self, mock_router, mock_prisma_client, mock_proxy_config, mock_user_api_key_dict
):
"""Test creating context_window fallback"""
request = FallbackCreateRequest(
model="gpt-3.5-turbo",
fallback_models=["gpt-4"],
fallback_type="context_window",
)
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router,
), patch(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
), patch(
"litellm.proxy.proxy_server.proxy_config",
mock_proxy_config,
), patch(
"litellm.proxy.proxy_server.store_model_in_db",
True,
):
response = await create_fallback(request, mock_user_api_key_dict)
assert response.fallback_type == "context_window"
# Verify the correct attribute was updated
assert hasattr(mock_router, "context_window_fallbacks")
@pytest.mark.asyncio
class TestGetFallback:
"""Test the get_fallback endpoint"""
@pytest.fixture
def mock_router_with_fallbacks(self):
"""Create a mock router with fallbacks configured"""
router = MagicMock()
router.fallbacks = [{"gpt-3.5-turbo": ["gpt-4", "claude-3-haiku"]}]
router.context_window_fallbacks = []
router.content_policy_fallbacks = []
return router
@pytest.fixture
def mock_user_api_key_dict(self):
"""Create a mock user API key dict"""
return MagicMock()
async def test_get_fallback_success(
self, mock_router_with_fallbacks, mock_user_api_key_dict
):
"""Test successful fallback retrieval"""
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router_with_fallbacks,
):
response = await get_fallback(
"gpt-3.5-turbo", "general", mock_user_api_key_dict
)
assert response.model == "gpt-3.5-turbo"
assert response.fallback_models == ["gpt-4", "claude-3-haiku"]
assert response.fallback_type == "general"
async def test_get_fallback_not_found(
self, mock_router_with_fallbacks, mock_user_api_key_dict
):
"""Test error when fallback is not found"""
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router_with_fallbacks,
), pytest.raises(HTTPException) as exc_info:
await get_fallback("gpt-4", "general", mock_user_api_key_dict)
assert exc_info.value.status_code == 404
assert "No general fallbacks configured" in str(exc_info.value.detail)
async def test_get_fallback_router_not_initialized(self, mock_user_api_key_dict):
"""Test error when router is not initialized"""
with patch(
"litellm.proxy.proxy_server.llm_router",
None,
), pytest.raises(HTTPException) as exc_info:
await get_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict)
assert exc_info.value.status_code == 500
assert "Router not initialized" in str(exc_info.value.detail)
@pytest.mark.asyncio
class TestDeleteFallback:
"""Test the delete_fallback endpoint"""
@pytest.fixture
def mock_router_with_fallbacks(self):
"""Create a mock router with fallbacks configured"""
router = MagicMock()
router.fallbacks = [{"gpt-3.5-turbo": ["gpt-4", "claude-3-haiku"]}]
router.context_window_fallbacks = []
router.content_policy_fallbacks = []
return router
@pytest.fixture
def mock_prisma_client(self):
"""Create a mock prisma client"""
client = MagicMock()
client.db.litellm_config.upsert = AsyncMock()
client.jsonify_object = lambda x: x
return client
@pytest.fixture
def mock_proxy_config(self):
"""Create a mock proxy config"""
config = MagicMock()
config.get_config = AsyncMock(
return_value={
"router_settings": {
"fallbacks": [{"gpt-3.5-turbo": ["gpt-4", "claude-3-haiku"]}]
}
}
)
return config
@pytest.fixture
def mock_user_api_key_dict(self):
"""Create a mock user API key dict"""
return MagicMock()
async def test_delete_fallback_success(
self,
mock_router_with_fallbacks,
mock_prisma_client,
mock_proxy_config,
mock_user_api_key_dict,
):
"""Test successful fallback deletion"""
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router_with_fallbacks,
), patch(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
), patch(
"litellm.proxy.proxy_server.proxy_config",
mock_proxy_config,
), patch(
"litellm.proxy.proxy_server.store_model_in_db",
True,
):
response = await delete_fallback(
"gpt-3.5-turbo", "general", mock_user_api_key_dict
)
assert response.model == "gpt-3.5-turbo"
assert response.fallback_type == "general"
assert "deleted" in response.message.lower()
# Verify database was updated
mock_prisma_client.db.litellm_config.upsert.assert_called_once()
async def test_delete_fallback_not_found(
self,
mock_router_with_fallbacks,
mock_prisma_client,
mock_proxy_config,
mock_user_api_key_dict,
):
"""Test error when fallback to delete is not found"""
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router_with_fallbacks,
), patch(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
), patch(
"litellm.proxy.proxy_server.proxy_config",
mock_proxy_config,
), patch(
"litellm.proxy.proxy_server.store_model_in_db",
True,
), pytest.raises(HTTPException) as exc_info:
await delete_fallback("gpt-4", "general", mock_user_api_key_dict)
assert exc_info.value.status_code == 404
assert "No general fallbacks configured" in str(exc_info.value.detail)
async def test_delete_fallback_router_not_initialized(self, mock_user_api_key_dict):
"""Test error when router is not initialized"""
with patch(
"litellm.proxy.proxy_server.llm_router",
None,
), pytest.raises(HTTPException) as exc_info:
await delete_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict)
assert exc_info.value.status_code == 500
assert "Router not initialized" in str(exc_info.value.detail)
async def test_delete_fallback_db_not_enabled(
self, mock_router_with_fallbacks, mock_user_api_key_dict
):
"""Test error when database storage is not enabled"""
with patch(
"litellm.proxy.proxy_server.llm_router",
mock_router_with_fallbacks,
), patch(
"litellm.proxy.proxy_server.store_model_in_db",
False,
), pytest.raises(HTTPException) as exc_info:
await delete_fallback("gpt-3.5-turbo", "general", mock_user_api_key_dict)
assert exc_info.value.status_code == 400
assert "Database storage not enabled" in str(exc_info.value.detail)