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

Merge branch 'BerriAI:main' into main

Browse Source
This commit is contained in:
abhijitjavelin
2025-09-28 01:08:16 +05:30
committed by GitHub
parent 7bba372ac0 973d4b71e7
commit adea1bfbc4
54 changed files with 1303 additions and 668 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.circleci/config.yml
+39 -11
View File
@@ -51,6 +51,33 @@ jobs:
command: |
python -m pytest tests/windows_tests/test_litellm_on_windows.py -v
mypy_linting:
docker:
- image: cimg/python:3.12
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
working_directory: ~/project
resource_class: medium
steps:
- checkout
- setup_google_dns
- run:
name: Install Dependencies
command: |
python -m pip install --upgrade pip
python -m pip install -r requirements.txt
pip uninstall fastuuid -y
pip install "mypy==1.18.2"
- run:
name: MyPy Type Checking
command: |
cd litellm
# Use the same approach as GitHub Actions, explicitly exclude fastuuid to avoid segfaults
python -m mypy .
cd ..
no_output_timeout: 10m
local_testing:
docker:
- image: cimg/python:3.12
@@ -140,13 +167,6 @@ jobs:
python -m pip install black
python -m black .
cd ..
- run:
name: Linting Testing
command: |
cd litellm
# Use the same simple approach that works in GitHub Actions
python -m mypy . --ignore-missing-imports --show-traceback
cd ..
# Run pytest and generate JUnit XML report
- run:
@@ -154,7 +174,7 @@ jobs:
command: |
pwd
ls
python -m pytest -vv tests/local_testing --cov=litellm --cov-report=xml -x --junitxml=test-results/junit.xml --durations=5 -k "not test_python_38.py and not test_basic_python_version.py and not router and not assistants and not langfuse and not caching and not cache" -n 4
python -m pytest -vv tests/local_testing --cov=litellm --cov-report=xml --junitxml=test-results/junit.xml --durations=5 -k "not test_python_38.py and not test_basic_python_version.py and not router and not assistants and not langfuse and not caching and not cache" -n 4
no_output_timeout: 120m
- run:
name: Rename the coverage files
@@ -464,7 +484,7 @@ jobs:
command: |
pwd
ls
python -m pytest tests/local_testing --cov=litellm --cov-report=xml -vv -k "router" -x -v --junitxml=test-results/junit.xml --durations=5
python -m pytest tests/local_testing --cov=litellm --cov-report=xml -vv -k "router" -v --junitxml=test-results/junit.xml --durations=5
no_output_timeout: 120m
- run:
name: Rename the coverage files
@@ -810,7 +830,7 @@ jobs:
command: |
pwd
ls
python -m pytest -vv tests/llm_translation --cov=litellm --cov-report=xml -x -v --junitxml=test-results/junit.xml --durations=5 -n 4
python -m pytest -vv tests/llm_translation --cov=litellm --cov-report=xml -v --junitxml=test-results/junit.xml --durations=5 -n 4
no_output_timeout: 120m
- run:
name: Rename the coverage files
@@ -1042,7 +1062,7 @@ jobs:
command: |
pwd
ls
python -m pytest -vv tests/test_litellm --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit-litellm.xml --durations=10 -n 8
python -m pytest -vv tests/test_litellm --cov=litellm --cov-report=xml -s -v --junitxml=test-results/junit-litellm.xml --durations=10 -n 8
no_output_timeout: 120m
- run:
name: Rename the coverage files
@@ -1180,6 +1200,7 @@ jobs:
pip install "pytest-cov==5.0.0"
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install pytest-mock
# Run pytest and generate JUnit XML report
- run:
name: Run tests
@@ -3071,6 +3092,12 @@ workflows:
only:
- main
- /litellm_.*/
- mypy_linting:
filters:
branches:
only:
- main
- /litellm_.*/
- local_testing:
filters:
branches:
@@ -3317,6 +3344,7 @@ workflows:
- main
- publish_to_pypi:
requires:
- mypy_linting
- local_testing
- build_and_test
- e2e_openai_endpoints
.github/workflows/test-linting.yml
+18 -1
View File
@@ -11,6 +11,9 @@ jobs:
steps:
- uses: actions/checkout@v4
with:
fetch-depth: 0
clean: true
- name: Set up Python
uses: actions/setup-python@v4
@@ -20,6 +23,11 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
- name: Clean Python cache
run: |
find . -type d -name "__pycache__" -exec rm -rf {} + || true
find . -name "*.pyc" -delete || true
- name: Install dependencies
run: |
poetry install --with dev
@@ -31,6 +39,15 @@ jobs:
poetry run black .
cd ..
- name: Debug - Check file state
run: |
echo "Current branch:"
git branch --show-current
echo "Last 3 commits:"
git log --oneline -3
echo "File content around line 43:"
head -50 litellm/litellm_core_utils/custom_logger_registry.py | tail -10
- name: Run Ruff linting
run: |
cd litellm
@@ -44,7 +61,7 @@ jobs:
- name: Run MyPy type checking
run: |
cd litellm
poetry run mypy . --ignore-missing-imports
poetry run mypy . --ignore-missing-imports --disable-error-code=var-annotated
cd ..
- name: Check for circular imports
.github/workflows/test-litellm.yml
+1 -1
View File
@@ -40,4 +40,4 @@ jobs:
cd ..
- name: Run tests
run: |
poetry run pytest tests/test_litellm -x -vv -n 4
poetry run pytest tests/test_litellm --tb=short -vv --maxfail=10 -n 4
docs/my-website/docs/pass_through/azure_passthrough.md
+89
View File
@@ -0,0 +1,89 @@
# Azure Passthrough
Pass-through endpoints for `/azure`
## Overview
| Feature | Supported | Notes |
|-------|-------|-------|
| Cost Tracking | ❌ | Not supported |
| Logging | ✅ | Works across all integrations |
| Streaming | ✅ | Fully supported |
### When to use this?
- For most use cases, you should use the [native LiteLLM Azure OpenAI Integration](../providers/azure/azure) (`/chat/completions`, `/embeddings`, `/completions`, `/images`, etc.)
- Use this passthrough to call newer or less common Azure OpenAI endpoints that LiteLLM doesn't fully support yet, such as `/assistants`, `/threads`, `/vector_stores`
Simply replace your Azure endpoint (e.g. `https://<your-resource-name>.openai.azure.com`) with `LITELLM_PROXY_BASE_URL/azure`
## Usage Examples
### Assistants API
#### Create Azure OpenAI Client
Make sure you do the following:
- Point `azure_endpoint` to your `LITELLM_PROXY_BASE_URL/azure`
- Use your `LITELLM_API_KEY` as the `api_key`
```python
import openai
client = openai.AzureOpenAI(
azure_endpoint="http://0.0.0.0:4000/azure", # <your-proxy-url>/azure
api_key="sk-anything", # <your-proxy-api-key>
api_version="2024-05-01-preview" # required Azure API version
)
```
#### Create an Assistant
```python
assistant = client.beta.assistants.create(
name="Math Tutor",
instructions="You are a math tutor. Help solve equations.",
model="gpt-4o",
)
```
#### Create a Thread
```python
thread = client.beta.threads.create()
```
#### Add a Message to the Thread
```python
message = client.beta.threads.messages.create(
thread_id=thread.id,
role="user",
content="Solve 3x + 11 = 14",
)
```
#### Run the Assistant
```python
run = client.beta.threads.runs.create(
thread_id=thread.id,
assistant_id=assistant.id,
)
# Check run status
run_status = client.beta.threads.runs.retrieve(
thread_id=thread.id,
run_id=run.id
)
```
#### Retrieve Messages
```python
messages = client.beta.threads.messages.list(
thread_id=thread.id
)
```
#### Delete the Assistant
```python
client.beta.assistants.delete(assistant.id)
```
docs/my-website/sidebars.js
+2
View File
@@ -343,6 +343,7 @@ const sidebars = {
"pass_through/anthropic_completion",
"pass_through/assembly_ai",
"pass_through/bedrock",
"pass_through/azure_passthrough",
"pass_through/cohere",
"pass_through/google_ai_studio",
"pass_through/langfuse",
@@ -584,6 +585,7 @@ const sidebars = {
"budget_manager",
"caching/all_caches",
"completion/token_usage",
"sdk_custom_pricing",
"embedding/async_embedding",
"embedding/moderation",
"migration",
litellm/experimental_mcp_client/client.py
+12 -5
View File
@@ -1,10 +1,11 @@
"""
LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
"""
import asyncio
import base64
from datetime import timedelta
from typing import List, Optional
from typing import Dict, List, Optional
from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client
@@ -46,6 +47,7 @@ class MCPClient:
auth_value: Optional[str] = None,
timeout: float = 60.0,
stdio_config: Optional[MCPStdioConfig] = None,
extra_headers: Optional[Dict[str, str]] = None,
):
self.server_url: str = server_url
self.transport_type: MCPTransport = transport_type
@@ -59,7 +61,7 @@ class MCPClient:
self._session_ctx = None
self._task: Optional[asyncio.Task] = None
self.stdio_config: Optional[MCPStdioConfig] = stdio_config
self.extra_headers: Optional[Dict[str, str]] = extra_headers
# handle the basic auth value if provided
if auth_value:
self.update_auth_value(auth_value)
@@ -115,6 +117,9 @@ class MCPClient:
await self._session.initialize()
else: # http
headers = self._get_auth_headers()
verbose_logger.debug(
"litellm headers for streamablehttp_client: ", headers
)
self._transport_ctx = streamablehttp_client(
url=self.server_url,
timeout=timedelta(seconds=self.timeout),
@@ -186,9 +191,7 @@ class MCPClient:
def _get_auth_headers(self) -> dict:
"""Generate authentication headers based on auth type."""
headers = {
"MCP-Protocol-Version": "2025-06-18"
}
headers = {"MCP-Protocol-Version": "2025-06-18"}
if self._mcp_auth_value:
if self.auth_type == MCPAuth.bearer_token:
@@ -200,6 +203,10 @@ class MCPClient:
elif self.auth_type == MCPAuth.authorization:
headers["Authorization"] = self._mcp_auth_value
# update the headers with the extra headers
if self.extra_headers:
headers.update(self.extra_headers)
return headers
async def list_tools(self) -> List[MCPTool]:
litellm/litellm_core_utils/custom_logger_registry.py
+1 -1
View File
@@ -15,6 +15,7 @@ from litellm.integrations.agentops import AgentOps
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.integrations.argilla import ArgillaLogger
from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger
from litellm.integrations.bitbucket import BitBucketPromptManager
from litellm.integrations.braintrust_logging import BraintrustLogger
from litellm.integrations.datadog.datadog import DataDogLogger
from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
@@ -39,7 +40,6 @@ try:
from litellm_enterprise.integrations.prometheus import PrometheusLogger
except Exception:
PrometheusLogger = None
from litellm.integrations.bitbucket import BitBucketPromptManager
from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
from litellm.integrations.dotprompt import DotpromptManager
from litellm.integrations.s3_v2 import S3Logger
litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py
+110 -61
View File
@@ -47,7 +47,7 @@ class StandardBuiltInToolCostTracking:
- Code Interpreter (Azure)
"""
standard_built_in_tools_params = standard_built_in_tools_params or {}
# Handle web search
if StandardBuiltInToolCostTracking.response_object_includes_web_search_call(
response_object=response_object, usage=usage
@@ -58,7 +58,7 @@ class StandardBuiltInToolCostTracking:
usage=usage,
standard_built_in_tools_params=standard_built_in_tools_params,
)
# Handle file search
if StandardBuiltInToolCostTracking.response_object_includes_file_search_call(
response_object=response_object
@@ -68,7 +68,7 @@ class StandardBuiltInToolCostTracking:
custom_llm_provider=custom_llm_provider,
standard_built_in_tools_params=standard_built_in_tools_params,
)
# Handle Azure assistant features
return StandardBuiltInToolCostTracking._handle_azure_assistant_costs(
model=model,
@@ -85,14 +85,14 @@ class StandardBuiltInToolCostTracking:
) -> float:
"""Handle web search cost calculation."""
from litellm.llms import get_cost_for_web_search_request
model_info = StandardBuiltInToolCostTracking._safe_get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
if custom_llm_provider is None and model_info is not None:
custom_llm_provider = model_info["litellm_provider"]
if (
model_info is not None
and usage is not None
@@ -105,9 +105,11 @@ class StandardBuiltInToolCostTracking:
)
if result is not None:
return result
return StandardBuiltInToolCostTracking.get_cost_for_web_search(
web_search_options=standard_built_in_tools_params.get("web_search_options", None),
web_search_options=standard_built_in_tools_params.get(
"web_search_options", None
),
model_info=model_info,
)
@@ -121,12 +123,17 @@ class StandardBuiltInToolCostTracking:
model_info = StandardBuiltInToolCostTracking._safe_get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
file_search_usage = standard_built_in_tools_params.get("file_search", {})
file_search_raw: Any = standard_built_in_tools_params.get("file_search", {})
file_search_usage: Optional[FileSearchTool] = (
FileSearchTool(**file_search_raw) if file_search_raw else None
)
# Convert model_info to dict and extract usage parameters
model_info_dict = dict(model_info) if model_info is not None else None
storage_gb, days = StandardBuiltInToolCostTracking._extract_file_search_params(file_search_usage)
storage_gb, days = StandardBuiltInToolCostTracking._extract_file_search_params(
file_search_usage
)
return StandardBuiltInToolCostTracking.get_cost_for_file_search(
file_search=file_search_usage,
provider=custom_llm_provider,
@@ -144,11 +151,11 @@ class StandardBuiltInToolCostTracking:
"""Handle Azure assistant features cost calculation."""
if custom_llm_provider != "azure":
return 0.0
model_info = StandardBuiltInToolCostTracking._safe_get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
total_cost = 0.0
total_cost += StandardBuiltInToolCostTracking._get_vector_store_cost(
model_info, custom_llm_provider, standard_built_in_tools_params
@@ -159,31 +166,33 @@ class StandardBuiltInToolCostTracking:
total_cost += StandardBuiltInToolCostTracking._get_code_interpreter_cost(
model_info, custom_llm_provider, standard_built_in_tools_params
)
return total_cost
@staticmethod
def _extract_file_search_params(file_search_usage: Any) -> Tuple[Optional[float], Optional[float]]:
def _extract_file_search_params(
file_search_usage: Any,
) -> Tuple[Optional[float], Optional[float]]:
"""Extract and convert file search parameters safely."""
storage_gb = None
days = None
if isinstance(file_search_usage, dict):
storage_gb_val = file_search_usage.get("storage_gb")
days_val = file_search_usage.get("days")
if storage_gb_val is not None:
try:
storage_gb = float(storage_gb_val) # type: ignore
except (TypeError, ValueError):
storage_gb = None
if days_val is not None:
try:
days = float(days_val) # type: ignore
except (TypeError, ValueError):
days = None
return storage_gb, days
@staticmethod
@@ -193,13 +202,17 @@ class StandardBuiltInToolCostTracking:
standard_built_in_tools_params: StandardBuiltInToolsParams,
) -> float:
"""Calculate vector store cost."""
vector_store_usage = standard_built_in_tools_params.get("vector_store_usage", None)
vector_store_usage = standard_built_in_tools_params.get(
"vector_store_usage", None
)
if not vector_store_usage:
return 0.0
model_info_dict = dict(model_info) if model_info is not None else None
vector_store_dict = vector_store_usage if isinstance(vector_store_usage, dict) else {}
vector_store_dict = (
vector_store_usage if isinstance(vector_store_usage, dict) else {}
)
return StandardBuiltInToolCostTracking.get_cost_for_vector_store(
vector_store_usage=vector_store_dict,
provider=custom_llm_provider,
@@ -213,13 +226,17 @@ class StandardBuiltInToolCostTracking:
standard_built_in_tools_params: StandardBuiltInToolsParams,
) -> float:
"""Calculate computer use cost."""
computer_use_usage = standard_built_in_tools_params.get("computer_use_usage", {})
computer_use_usage = standard_built_in_tools_params.get(
"computer_use_usage", {}
)
if not computer_use_usage:
return 0.0
model_info_dict = dict(model_info) if model_info is not None else None
input_tokens, output_tokens = StandardBuiltInToolCostTracking._extract_token_counts(computer_use_usage)
input_tokens, output_tokens = (
StandardBuiltInToolCostTracking._extract_token_counts(computer_use_usage)
)
return StandardBuiltInToolCostTracking.get_cost_for_computer_use(
input_tokens=input_tokens,
output_tokens=output_tokens,
@@ -234,13 +251,17 @@ class StandardBuiltInToolCostTracking:
standard_built_in_tools_params: StandardBuiltInToolsParams,
) -> float:
"""Calculate code interpreter cost."""
code_interpreter_sessions = standard_built_in_tools_params.get("code_interpreter_sessions", None)
code_interpreter_sessions = standard_built_in_tools_params.get(
"code_interpreter_sessions", None
)
if not code_interpreter_sessions:
return 0.0
model_info_dict = dict(model_info) if model_info is not None else None
sessions = StandardBuiltInToolCostTracking._safe_convert_to_int(code_interpreter_sessions)
sessions = StandardBuiltInToolCostTracking._safe_convert_to_int(
code_interpreter_sessions
)
return StandardBuiltInToolCostTracking.get_cost_for_code_interpreter(
sessions=sessions,
provider=custom_llm_provider,
@@ -248,18 +269,24 @@ class StandardBuiltInToolCostTracking:
)
@staticmethod
def _extract_token_counts(computer_use_usage: Any) -> Tuple[Optional[int], Optional[int]]:
def _extract_token_counts(
computer_use_usage: Any,
) -> Tuple[Optional[int], Optional[int]]:
"""Extract and convert token counts safely."""
input_tokens = None
output_tokens = None
if isinstance(computer_use_usage, dict):
input_tokens_val = computer_use_usage.get("input_tokens")
output_tokens_val = computer_use_usage.get("output_tokens")
input_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(input_tokens_val)
output_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(output_tokens_val)
input_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(
input_tokens_val
)
output_tokens = StandardBuiltInToolCostTracking._safe_convert_to_int(
output_tokens_val
)
return input_tokens, output_tokens
@staticmethod
@@ -400,8 +427,11 @@ class StandardBuiltInToolCostTracking:
if model_info is None:
return 0.0
search_context_raw: Any = model_info.get("search_context_cost_per_query", {})
search_context_pricing: SearchContextCostPerQuery = (
model_info.get("search_context_cost_per_query", {}) or {}
SearchContextCostPerQuery(**search_context_raw)
if search_context_raw
else SearchContextCostPerQuery()
)
if web_search_options.get("search_context_size", None) == "low":
return search_context_pricing.get("search_context_size_low", 0.0)
@@ -424,9 +454,12 @@ class StandardBuiltInToolCostTracking:
"""
if model_info is None:
return 0.0
search_context_raw: Any = model_info.get("search_context_cost_per_query", {}) or {}
search_context_pricing: SearchContextCostPerQuery = (
model_info.get("search_context_cost_per_query", {}) or {}
) or {}
SearchContextCostPerQuery(**search_context_raw)
if search_context_raw
else SearchContextCostPerQuery()
)
return search_context_pricing.get("search_context_size_medium", 0.0)
@staticmethod
@@ -445,22 +478,27 @@ class StandardBuiltInToolCostTracking:
"""
if file_search is None:
return 0.0
# Check if model-specific pricing is available
if model_info and "file_search_cost_per_gb_per_day" in model_info and provider == "azure":
if (
model_info
and "file_search_cost_per_gb_per_day" in model_info
and provider == "azure"
):
if storage_gb and days:
return storage_gb * days * model_info["file_search_cost_per_gb_per_day"]
elif model_info and "file_search_cost_per_1k_calls" in model_info:
return model_info["file_search_cost_per_1k_calls"]
# Azure has storage-based pricing for file search
if provider == "azure":
from litellm.constants import AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY
if storage_gb and days:
return storage_gb * days * AZURE_FILE_SEARCH_COST_PER_GB_PER_DAY
# Default to 0 if no storage info provided
return 0.0
# Default to OpenAI pricing (per-call based)
return OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
@@ -472,24 +510,25 @@ class StandardBuiltInToolCostTracking:
) -> float:
"""
Calculate cost for vector store usage.
Azure charges based on storage size and duration.
"""
if vector_store_usage is None:
return 0.0
storage_gb = vector_store_usage.get("storage_gb", 0.0)
days = vector_store_usage.get("days", 0.0)
# Check if model-specific pricing is available
if model_info and "vector_store_cost_per_gb_per_day" in model_info:
return storage_gb * days * model_info["vector_store_cost_per_gb_per_day"]
# Azure has different pricing structure for vector store
if provider == "azure":
from litellm.constants import AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY
return storage_gb * days * AZURE_VECTOR_STORE_COST_PER_GB_PER_DAY
# OpenAI doesn't charge separately for vector store (included in embeddings)
return 0.0
@@ -502,14 +541,18 @@ class StandardBuiltInToolCostTracking:
) -> float:
"""
Calculate cost for computer use feature.
Azure: $0.003 USD per 1K input tokens, $0.012 USD per 1K output tokens
"""
if provider == "azure" and (input_tokens or output_tokens):
# Check if model-specific pricing is available
if model_info:
input_cost = model_info.get("computer_use_input_cost_per_1k_tokens", 0.0)
output_cost = model_info.get("computer_use_output_cost_per_1k_tokens", 0.0)
input_cost = model_info.get(
"computer_use_input_cost_per_1k_tokens", 0.0
)
output_cost = model_info.get(
"computer_use_output_cost_per_1k_tokens", 0.0
)
if input_cost or output_cost:
total_cost = 0.0
if input_tokens:
@@ -517,19 +560,24 @@ class StandardBuiltInToolCostTracking:
if output_tokens:
total_cost += (output_tokens / 1000.0) * output_cost
return total_cost
# Azure default pricing
from litellm.constants import (
AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS,
AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS,
)
total_cost = 0.0
if input_tokens:
total_cost += (input_tokens / 1000.0) * AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS
total_cost += (
input_tokens / 1000.0
) * AZURE_COMPUTER_USE_INPUT_COST_PER_1K_TOKENS
if output_tokens:
total_cost += (output_tokens / 1000.0) * AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS
total_cost += (
output_tokens / 1000.0
) * AZURE_COMPUTER_USE_OUTPUT_COST_PER_1K_TOKENS
return total_cost
# OpenAI doesn't charge separately for computer use yet
return 0.0
@@ -541,21 +589,22 @@ class StandardBuiltInToolCostTracking:
) -> float:
"""
Calculate cost for code interpreter feature.
Azure: $0.03 USD per session
"""
if sessions is None or sessions == 0:
return 0.0
# Check if model-specific pricing is available
if model_info and "code_interpreter_cost_per_session" in model_info:
return sessions * model_info["code_interpreter_cost_per_session"]
# Azure pricing for code interpreter
if provider == "azure":
from litellm.constants import AZURE_CODE_INTERPRETER_COST_PER_SESSION
return sessions * AZURE_CODE_INTERPRETER_COST_PER_SESSION
# OpenAI doesn't charge separately for code interpreter yet
return 0.0
litellm/llms/azure/batches/handler.py
+4 -4
View File
@@ -62,7 +62,7 @@ class AzureBatchesAPI(BaseAzureLLM):
)
if azure_client is None:
raise ValueError(
"Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment."
"OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
)
if _is_async is True:
@@ -108,7 +108,7 @@ class AzureBatchesAPI(BaseAzureLLM):
)
if azure_client is None:
raise ValueError(
"Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment."
"OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
)
if _is_async is True:
@@ -156,7 +156,7 @@ class AzureBatchesAPI(BaseAzureLLM):
)
if azure_client is None:
raise ValueError(
"Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment."
"OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
)
response = azure_client.batches.cancel(**cancel_batch_data)
return response
@@ -195,7 +195,7 @@ class AzureBatchesAPI(BaseAzureLLM):
)
if azure_client is None:
raise ValueError(
"Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment."
"OpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
)
if _is_async is True:
litellm/llms/azure/common_utils.py
+1 -12
View File
@@ -576,19 +576,8 @@ class BaseAzureLLM(BaseOpenAILLM):
verbose_logger.debug(
f"Initializing Azure OpenAI Client for {model_name}, Api Base: {str(api_base)}, Api Key:{_api_key}"
)
# Extract API key from multiple sources with proper precedence
resolved_api_key = (
api_key
or litellm_params.get("api_key")
or litellm.api_key
or litellm.azure_key
or get_secret_str("AZURE_OPENAI_API_KEY")
or get_secret_str("AZURE_API_KEY")
)
azure_client_params = {
"api_key": resolved_api_key,
"api_key": api_key,
"azure_endpoint": api_base,
"api_version": api_version,
"azure_ad_token": azure_ad_token,
litellm/llms/azure/files/handler.py
+5 -5
View File
@@ -58,7 +58,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
)
if openai_client is None:
raise ValueError(
"Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment."
"AzureOpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
)
if _is_async is True:
@@ -106,7 +106,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
)
if openai_client is None:
raise ValueError(
"Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment."
"AzureOpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
)
if _is_async is True:
@@ -156,7 +156,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
)
if openai_client is None:
raise ValueError(
"Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment."
"AzureOpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
)
if _is_async is True:
@@ -208,7 +208,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
)
if openai_client is None:
raise ValueError(
"Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment."
"AzureOpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
)
if _is_async is True:
@@ -262,7 +262,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM):
)
if openai_client is None:
raise ValueError(
"Azure OpenAI client is not initialized. Make sure api_key is passed or AZURE_API_KEY/AZURE_OPENAI_API_KEY is set in the environment."
"AzureOpenAI client is not initialized. Make sure api_key is passed or OPENAI_API_KEY is set in the environment."
)
if _is_async is True:
litellm/mypy.ini
+2 -1
View File
@@ -5,7 +5,8 @@ mypy_path = litellm/stubs
namespace_packages = True
disable_error_code =
valid-type,
annotation-unchecked
annotation-unchecked,
import-untyped
[mypy-google.*]
ignore_missing_imports = True
litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py
+13 -3
View File
@@ -1,4 +1,4 @@
from typing import List, Optional, Dict
from typing import Dict, List, Optional
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
@@ -8,17 +8,27 @@ from litellm.proxy._types import UserAPIKeyAuth
class MCPAuthenticatedUser(AuthenticatedUser):
"""
Wrapper class to make LiteLLM's authentication and configuration compatible with MCP's AuthenticatedUser.
This class handles:
1. User API key authentication information
2. MCP authentication header (deprecated)
3. MCP server configuration (can include access groups)
4. Server-specific authentication headers
5. OAuth2 headers
"""
def __init__(self, user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, mcp_server_auth_headers: Optional[Dict[str, str]] = None, mcp_protocol_version: Optional[str] = None):
def __init__(
self,
user_api_key_auth: UserAPIKeyAuth,
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
):
self.user_api_key_auth = user_api_key_auth
self.mcp_auth_header = mcp_auth_header
self.mcp_servers = mcp_servers
self.mcp_server_auth_headers = mcp_server_auth_headers or {}
self.mcp_protocol_version = mcp_protocol_version
self.oauth2_headers = oauth2_headers
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+36 -13
View File
@@ -1,4 +1,4 @@
from typing import List, Optional, Tuple, Dict, Set
from typing import Dict, List, Optional, Set, Tuple
from starlette.datastructures import Headers
from starlette.requests import Request
@@ -36,7 +36,11 @@ class MCPRequestHandler:
async def process_mcp_request(
scope: Scope,
) -> Tuple[
UserAPIKeyAuth, Optional[str], Optional[List[str]], Optional[Dict[str, str]]
UserAPIKeyAuth,
Optional[str],
Optional[List[str]],
Optional[Dict[str, str]],
Optional[Dict[str, str]],
]:
"""
Process and validate MCP request headers from the ASGI scope.
@@ -44,6 +48,7 @@ class MCPRequestHandler:
1. Extracting and validating authentication headers
2. Processing MCP server configuration
3. Handling MCP-specific headers
4. Handling oauth2 headers
Args:
scope: ASGI scope containing request information
@@ -70,6 +75,9 @@ class MCPRequestHandler:
MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
)
# Get the oauth2 headers
oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
# Parse MCP servers from header
mcp_servers_header = headers.get(
MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME
@@ -96,14 +104,18 @@ class MCPRequestHandler:
return b"{}"
request.body = mock_body # type: ignore
validated_user_api_key_auth = await user_api_key_auth(
api_key=litellm_api_key, request=request
)
if ".well-known" in str(request.url): # public routes
validated_user_api_key_auth = UserAPIKeyAuth()
else:
validated_user_api_key_auth = await user_api_key_auth(
api_key=litellm_api_key, request=request
)
return (
validated_user_api_key_auth,
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
)
@staticmethod
@@ -174,6 +186,17 @@ class MCPRequestHandler:
return server_auth_headers
@staticmethod
def _get_oauth2_headers_from_headers(headers: Headers) -> Dict[str, str]:
"""
Get the oauth2 headers from the request headers.
"""
oauth2_headers = {}
for header_name, header_value in headers.items():
if header_name.lower().startswith("authorization"):
oauth2_headers["Authorization"] = header_value
return oauth2_headers
@staticmethod
def _get_mcp_client_side_auth_header_name() -> str:
"""
@@ -359,10 +382,10 @@ class MCPRequestHandler:
return []
try:
team_obj: Optional[
LiteLLM_TeamTable
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
team_obj: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
)
)
if team_obj is None:
verbose_logger.debug("team_obj is None")
@@ -535,10 +558,10 @@ class MCPRequestHandler:
verbose_logger.debug("prisma_client is None")
return []
team_obj: Optional[
LiteLLM_TeamTable
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
team_obj: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
)
)
if team_obj is None:
verbose_logger.debug("team_obj is None")
litellm/proxy/_experimental/mcp_server/cost_calculator.py
+2 -1
View File
@@ -38,7 +38,8 @@ class MCPCostCalculator:
# Unpack the mcp_tool_call_metadata
#########################################################
mcp_tool_call_metadata: StandardLoggingMCPToolCall = cast(StandardLoggingMCPToolCall, litellm_logging_obj.model_call_details.get("mcp_tool_call_metadata", {})) or {}
mcp_server_cost_info: MCPServerCostInfo = mcp_tool_call_metadata.get("mcp_server_cost_info", {}) or {}
mcp_server_cost_info_raw = mcp_tool_call_metadata.get("mcp_server_cost_info", {}) or {}
mcp_server_cost_info: MCPServerCostInfo = cast(MCPServerCostInfo, mcp_server_cost_info_raw)
#########################################################
# User defined cost per query
#########################################################
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
+252
View File
@@ -0,0 +1,252 @@
import json
from typing import Optional
from urllib.parse import urlencode, urlparse, urlunparse
from fastapi import APIRouter, Form, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
router = APIRouter(
tags=["mcp"],
)
def encode_state_with_base_url(base_url: str, original_state: str) -> str:
"""
Encode the base_url and original state using encryption.
Args:
base_url: The base URL to encode
original_state: The original state parameter
Returns:
An encrypted string that encodes both values
"""
state_data = {"base_url": base_url, "original_state": original_state}
state_json = json.dumps(state_data, sort_keys=True)
encrypted_state = encrypt_value_helper(state_json)
return encrypted_state
def decode_state_hash(encrypted_state: str) -> tuple[str, str]:
"""
Decode an encrypted state to retrieve the base_url and original state.
Args:
encrypted_state: The encrypted string to decode
Returns:
A tuple of (base_url, original_state)
Raises:
Exception: If decryption fails or data is malformed
"""
decrypted_json = decrypt_value_helper(encrypted_state, "oauth_state")
if decrypted_json is None:
raise ValueError("Failed to decrypt state parameter")
state_data = json.loads(decrypted_json)
return state_data["base_url"], state_data["original_state"]
@router.get("/{mcp_server_name}/authorize")
@router.get("/authorize")
async def authorize(
request: Request,
client_id: str,
redirect_uri: str,
state: str = "",
mcp_server_name: Optional[str] = None,
):
# Redirect to real GitHub OAuth
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(client_id)
if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found")
if mcp_server.auth_type != "oauth2":
raise HTTPException(status_code=400, detail="MCP server is not OAuth2")
if mcp_server.client_id is None:
raise HTTPException(status_code=400, detail="MCP server client id is not set")
if mcp_server.authorization_url is None:
raise HTTPException(
status_code=400, detail="MCP server authorization url is not set"
)
if mcp_server.scopes is None:
raise HTTPException(status_code=400, detail="MCP server scopes is not set")
# Parse it to remove any existing query
parsed = urlparse(redirect_uri)
base_url = urlunparse(parsed._replace(query=""))
request_base_url = str(request.base_url).rstrip("/")
# Encode the base_url and original state in a unique hash
encoded_state = encode_state_with_base_url(base_url, state)
params = {
"client_id": mcp_server.client_id,
"redirect_uri": f"{request_base_url}/callback",
"scope": " ".join(mcp_server.scopes),
"state": encoded_state,
}
return RedirectResponse(f"{mcp_server.authorization_url}?{urlencode(params)}")
@router.post("/token")
async def token_endpoint(
request: Request,
grant_type: str = Form(...),
code: str = Form(None),
redirect_uri: str = Form(None),
client_id: str = Form(...),
client_secret: str = Form(...),
):
"""
Accept the authorization code from Claude and exchange it for GitHub token.
Forward the GitHub token back to Claude in standard OAuth format.
1. Call the token endpoint
2. Store the user's PAT in the db - and generate a LiteLLM virtual key
2. Return the token
3. Return a virtual key in this response
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(client_id)
if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found")
if grant_type != "authorization_code":
raise HTTPException(status_code=400, detail="Unsupported grant_type")
if mcp_server.token_url is None:
raise HTTPException(status_code=400, detail="MCP server token url is not set")
proxy_base_url = str(request.base_url).rstrip("/")
# Exchange code for real GitHub token
async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
response = await async_client.post(
mcp_server.token_url,
headers={"Accept": "application/json"},
data={
"client_id": mcp_server.client_id,
"client_secret": mcp_server.client_secret,
"code": code,
"redirect_uri": f"{proxy_base_url}/callback",
},
)
response.raise_for_status()
github_token = response.json()["access_token"]
# Return to Claude in expected OAuth 2 format
### return a virtual key in this response
return JSONResponse(
{"access_token": github_token, "token_type": "Bearer", "expires_in": 3600}
)
@router.get("/callback")
async def callback(code: str, state: str):
try:
# Decode the state hash to get base_url and original state
base_url, original_state = decode_state_hash(state)
# Exchange code for token with GitHub
params = {"code": code, "state": original_state}
# Forward token to Claude ephemeral endpoint
complete_returned_url = f"{base_url}?{urlencode(params)}"
return RedirectResponse(url=complete_returned_url, status_code=302)
except Exception:
# fallback if state hash not found
return HTMLResponse(
"<html><body>Authentication incomplete. You can close this window.</body></html>"
)
# ------------------------------
# Optional .well-known endpoints for MCP + OAuth discovery
# ------------------------------
@router.get("/.well-known/oauth-protected-resource/{mcp_server_name}/mcp")
@router.get("/.well-known/oauth-protected-resource")
async def oauth_protected_resource_mcp(
request: Request, mcp_server_name: Optional[str] = None
):
request_base_url = str(request.base_url).rstrip("/")
return {
"authorization_servers": [
(
f"{request_base_url}/{mcp_server_name}"
if mcp_server_name
else f"{request_base_url}"
)
],
"resource": (
f"{request_base_url}/{mcp_server_name}/mcp"
if mcp_server_name
else f"{request_base_url}/mcp"
), # this is what Claude will call
}
@router.get("/.well-known/oauth-authorization-server/{mcp_server_name}")
@router.get("/.well-known/oauth-authorization-server")
async def oauth_authorization_server_mcp(
request: Request, mcp_server_name: Optional[str] = None
):
request_base_url = str(request.base_url).rstrip("/")
return {
"issuer": request_base_url, # point to your proxy
"authorization_endpoint": f"{request_base_url}/authorize",
"token_endpoint": f"{request_base_url}/token",
"response_types_supported": ["code"],
"grant_types_supported": ["authorization_code"],
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["client_secret_post"],
# Claude expects a registration endpoint, even if we just fake it
"registration_endpoint": f"{request_base_url}/{mcp_server_name}/register",
}
# Alias for standard OpenID discovery
@router.get("/.well-known/openid-configuration")
async def openid_configuration(request: Request):
return await oauth_authorization_server_mcp(request)
@router.get("/.well-known/oauth-authorization-server/{mcp_server_name}/mcp")
@router.get("/.well-known/oauth-authorization-server")
async def oauth_authorization_server_root(
request: Request, mcp_server_name: Optional[str] = None
):
return await oauth_authorization_server_mcp(request, mcp_server_name)
@router.post("/{mcp_server_name}/register")
@router.post("/register")
async def register_client(request: Request, mcp_server_name: Optional[str] = None):
request_base_url = str(request.base_url).rstrip("/")
# return fixed GitHub client credentials
return {
"client_id": mcp_server_name or "dummy_client",
"client_secret": "dummy",
"redirect_uris": [f"{request_base_url}/mcp/callback"],
}
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+111 -61
View File
@@ -39,7 +39,7 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.utils import ProxyLogging
from litellm.types.mcp import MCPStdioConfig
from litellm.types.mcp import MCPAuth, MCPStdioConfig
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
@@ -199,6 +199,12 @@ class MCPServerManager:
command=server_config.get("command", None) or "",
args=server_config.get("args", None) or [],
env=server_config.get("env", None) or {},
# oauth specific fields
client_id=server_config.get("client_id", None),
client_secret=server_config.get("client_secret", None),
scopes=server_config.get("scopes", None),
authorization_url=server_config.get("authorization_url", None),
token_url=server_config.get("token_url", None),
# TODO: utility fn the default values
transport=server_config.get("transport", MCPTransport.http),
auth_type=server_config.get("auth_type", None),
@@ -376,6 +382,7 @@ class MCPServerManager:
self,
server: MCPServer,
mcp_auth_header: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> MCPClient:
"""
Create an MCPClient instance for the given server.
@@ -405,6 +412,7 @@ class MCPServerManager:
auth_value=mcp_auth_header or server.authentication_token,
timeout=60.0,
stdio_config=stdio_config,
extra_headers=extra_headers,
)
else:
# For HTTP/SSE transports
@@ -415,12 +423,14 @@ class MCPServerManager:
auth_type=server.auth_type,
auth_value=mcp_auth_header or server.authentication_token,
timeout=60.0,
extra_headers=extra_headers,
)
async def _get_tools_from_server(
self,
server: MCPServer,
mcp_auth_header: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
Helper method to get tools from a single MCP server with prefixed names.
@@ -441,6 +451,7 @@ class MCPServerManager:
client = self._create_mcp_client(
server=server,
mcp_auth_header=mcp_auth_header,
extra_headers=extra_headers,
)
tools = await self._fetch_tools_with_timeout(client, server.name)
@@ -550,6 +561,77 @@ class MCPServerManager:
)
return prefixed_tools
async def pre_call_tool_check(
self,
name: str,
arguments: Dict[str, Any],
server_name_from_prefix: str,
user_api_key_auth: Optional[UserAPIKeyAuth],
proxy_logging_obj: ProxyLogging,
):
pre_hook_kwargs = {
"name": name,
"arguments": arguments,
"server_name": server_name_from_prefix,
"user_api_key_auth": user_api_key_auth,
"user_api_key_user_id": (
getattr(user_api_key_auth, "user_id", None)
if user_api_key_auth
else None
),
"user_api_key_team_id": (
getattr(user_api_key_auth, "team_id", None)
if user_api_key_auth
else None
),
"user_api_key_end_user_id": (
getattr(user_api_key_auth, "end_user_id", None)
if user_api_key_auth
else None
),
"user_api_key_hash": (
getattr(user_api_key_auth, "api_key_hash", None)
if user_api_key_auth
else None
),
}
# Create MCP request object for processing
mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs(
pre_hook_kwargs
)
# Convert to LLM format for existing guardrail compatibility
synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(
mcp_request_obj, pre_hook_kwargs
)
try:
# Use standard pre_call_hook with call_type="mcp_call"
modified_data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_auth, # type: ignore
data=synthetic_llm_data,
call_type="mcp_call", # type: ignore
)
if modified_data:
# Convert response back to MCP format and apply modifications
modified_kwargs = (
proxy_logging_obj._convert_mcp_hook_response_to_kwargs(
modified_data, pre_hook_kwargs
)
)
if modified_kwargs.get("arguments") != arguments:
arguments = modified_kwargs["arguments"]
except (
BlockedPiiEntityError,
GuardrailRaisedException,
HTTPException,
) as e:
# Re-raise guardrail exceptions to properly fail the MCP call
verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {str(e)}")
raise e
async def call_tool(
self,
name: str,
@@ -558,6 +640,7 @@ class MCPServerManager:
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> CallToolResult:
"""
Call a tool with the given name and arguments (handles prefixed tool names)
@@ -602,65 +685,14 @@ class MCPServerManager:
# Using standard pre_call_hook with call_type="mcp_call"
#########################################################
if proxy_logging_obj:
pre_hook_kwargs = {
"name": name,
"arguments": arguments,
"server_name": server_name_from_prefix,
"user_api_key_auth": user_api_key_auth,
"user_api_key_user_id": getattr(user_api_key_auth, "user_id", None)
if user_api_key_auth
else None,
"user_api_key_team_id": getattr(user_api_key_auth, "team_id", None)
if user_api_key_auth
else None,
"user_api_key_end_user_id": getattr(
user_api_key_auth, "end_user_id", None
)
if user_api_key_auth
else None,
"user_api_key_hash": getattr(user_api_key_auth, "api_key_hash", None)
if user_api_key_auth
else None,
}
# Create MCP request object for processing
mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs(
pre_hook_kwargs
await self.pre_call_tool_check(
name=original_tool_name,
arguments=arguments,
server_name_from_prefix=server_name_from_prefix,
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
)
# Convert to LLM format for existing guardrail compatibility
synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(
mcp_request_obj, pre_hook_kwargs
)
try:
# Use standard pre_call_hook with call_type="mcp_call"
modified_data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_auth, # type: ignore
data=synthetic_llm_data,
call_type="mcp_call", # type: ignore
)
if modified_data:
# Convert response back to MCP format and apply modifications
modified_kwargs = (
proxy_logging_obj._convert_mcp_hook_response_to_kwargs(
modified_data, pre_hook_kwargs
)
)
if modified_kwargs.get("arguments") != arguments:
arguments = modified_kwargs["arguments"]
except (
BlockedPiiEntityError,
GuardrailRaisedException,
HTTPException,
) as e:
# Re-raise guardrail exceptions to properly fail the MCP call
verbose_logger.error(
f"Guardrail blocked MCP tool call pre call: {str(e)}"
)
raise e
# Get server-specific auth header if available
server_auth_header = None
if mcp_server_auth_headers and mcp_server.alias:
@@ -672,9 +704,15 @@ class MCPServerManager:
if server_auth_header is None:
server_auth_header = mcp_auth_header
# oauth2 headers
extra_headers: Optional[Dict[str, str]] = None
if mcp_server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
client = self._create_mcp_client(
server=mcp_server,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
)
async with client:
@@ -834,6 +872,16 @@ class MCPServerManager:
return server
return None
def get_mcp_server_by_name(self, server_name: str) -> Optional[MCPServer]:
"""
Get the MCP Server from the server name
"""
registry = self.get_registry()
for server in registry.values():
if server.server_name == server_name:
return server
return None
def _generate_stable_server_id(
self,
server_name: str,
@@ -1023,9 +1071,11 @@ class MCPServerManager:
auth_type=_server_config.auth_type,
created_at=datetime.datetime.now(),
updated_at=datetime.datetime.now(),
description=_server_config.mcp_info.get("description")
if _server_config.mcp_info
else None,
description=(
_server_config.mcp_info.get("description")
if _server_config.mcp_info
else None
),
mcp_info=_server_config.mcp_info,
mcp_access_groups=_server_config.access_groups or [],
# Stdio-specific fields
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
+3 -3
View File
@@ -177,9 +177,9 @@ if MCP_AVAILABLE:
return {
"tools": list_tools_result,
"error": "partial_failure" if error_message else None,
"message": error_message
if error_message
else "Successfully retrieved tools",
"message": (
error_message if error_message else "Successfully retrieved tools"
),
}
except Exception as e:
litellm/proxy/_experimental/mcp_server/server.py
+65 -25
View File
@@ -22,6 +22,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_VERSION,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
from litellm.types.utils import StandardLoggingMCPToolCall
from litellm.utils import client
@@ -178,6 +179,7 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = get_auth_context()
verbose_logger.debug(
f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}"
@@ -195,6 +197,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
)
verbose_logger.info(
f"MCP list_tools - Successfully returned {len(tools)} tools"
@@ -235,6 +238,7 @@ if MCP_AVAILABLE:
mcp_auth_header,
_,
mcp_server_auth_headers,
oauth2_headers,
) = get_auth_context()
verbose_logger.debug(
@@ -266,6 +270,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
**data, # for logging
)
except BlockedPiiEntityError as e:
@@ -357,6 +362,7 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str],
mcp_servers: Optional[List[str]],
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
Helper method to fetch tools from MCP servers based on server filtering criteria.
@@ -365,7 +371,8 @@ if MCP_AVAILABLE:
user_api_key_auth: User authentication info for access control
mcp_auth_header: Optional auth header for MCP server (deprecated)
mcp_servers: Optional list of server names/aliases to filter by
mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value}
mcp_server_auth_headers: Optional dict of server-specific auth headers
oauth2_headers: Optional dict of oauth2 headers
Returns:
List[MCPTool]: Combined list of tools from filtered servers
@@ -398,6 +405,10 @@ if MCP_AVAILABLE:
elif mcp_server_auth_headers and server.server_name is not None:
server_auth_header = mcp_server_auth_headers.get(server.server_name)
extra_headers: Optional[Dict[str, str]] = None
if server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
# Fall back to deprecated mcp_auth_header if no server-specific header found
if server_auth_header is None:
server_auth_header = mcp_auth_header
@@ -406,6 +417,7 @@ if MCP_AVAILABLE:
tools = await global_mcp_server_manager._get_tools_from_server(
server=server,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
)
all_tools.extend(tools)
verbose_logger.debug(
@@ -427,6 +439,7 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
List all available MCP tools.
@@ -450,6 +463,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
)
verbose_logger.debug(
f"Successfully fetched {len(managed_tools)} tools from managed MCP servers"
@@ -492,6 +506,7 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
**kwargs: Any,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
"""
@@ -519,16 +534,16 @@ if MCP_AVAILABLE:
"litellm_logging_obj", None
)
if litellm_logging_obj:
litellm_logging_obj.model_call_details[
"mcp_tool_call_metadata"
] = standard_logging_mcp_tool_call
litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = (
standard_logging_mcp_tool_call
)
litellm_logging_obj.model = f"MCP: {name}"
# Try managed server tool first (pass the full prefixed name)
# Primary and recommended way to use MCP servers
#########################################################
mcp_server: Optional[
MCPServer
] = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
mcp_server: Optional[MCPServer] = (
global_mcp_server_manager._get_mcp_server_from_tool_name(name)
)
if mcp_server:
standard_logging_mcp_tool_call["mcp_server_cost_info"] = (
mcp_server.mcp_info or {}
@@ -539,6 +554,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
litellm_logging_obj=litellm_logging_obj,
)
@@ -591,6 +607,7 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
litellm_logging_obj: Optional[Any] = None,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
"""Handle tool execution for managed server tools"""
@@ -603,6 +620,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
proxy_logging_obj=proxy_logging_obj,
)
verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
@@ -638,26 +656,32 @@ if MCP_AVAILABLE:
mcp_path_match = re.match(r"^/mcp/([^?#]+)(?:\?.*)?(?:#.*)?$", path)
if mcp_path_match:
servers_and_path = mcp_path_match.group(1)
if servers_and_path:
# Check if it contains commas (comma-separated servers)
if ',' in servers_and_path:
if "," in servers_and_path:
# For comma-separated, look for a path at the end
# Common patterns: /tools, /chat/completions, etc.
path_match = re.search(r'/([^/,]+(?:/[^/,]+)*)$', servers_and_path)
path_match = re.search(r"/([^/,]+(?:/[^/,]+)*)$", servers_and_path)
if path_match:
# Path found at the end, remove it from servers
path_part = '/' + path_match.group(1)
servers_part = servers_and_path[:-len(path_part)]
mcp_servers_from_path = [s.strip() for s in servers_part.split(',') if s.strip()]
path_part = "/" + path_match.group(1)
servers_part = servers_and_path[: -len(path_part)]
mcp_servers_from_path = [
s.strip() for s in servers_part.split(",") if s.strip()
]
else:
# No path, just comma-separated servers
mcp_servers_from_path = [s.strip() for s in servers_and_path.split(',') if s.strip()]
mcp_servers_from_path = [
s.strip() for s in servers_and_path.split(",") if s.strip()
]
else:
# Single server case - use regex approach for server/path separation
# This handles cases like "custom_solutions/user_123/chat/completions"
# where we want to extract "custom_solutions/user_123" as the server name
single_server_match = re.match(r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path)
single_server_match = re.match(
r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path
)
if single_server_match:
server_name = single_server_match.group(1)
mcp_servers_from_path = [server_name]
@@ -677,6 +701,7 @@ if MCP_AVAILABLE:
mcp_auth_header,
_,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
mcp_servers = mcp_servers_from_path
else:
@@ -685,8 +710,15 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
return user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers
return (
user_api_key_auth,
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
)
async def handle_streamable_http_mcp(
scope: Scope, receive: Receive, send: Send
@@ -699,6 +731,7 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await extract_mcp_auth_context(scope, path)
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
@@ -712,6 +745,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
)
# Ensure session managers are initialized
@@ -750,6 +784,7 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await extract_mcp_auth_context(scope, path)
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
@@ -762,6 +797,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
)
if not _SESSION_MANAGERS_INITIALIZED:
@@ -809,6 +845,8 @@ if MCP_AVAILABLE:
# Mount the MCP handlers
app.mount("/", handle_streamable_http_mcp)
app.mount("/mcp", handle_streamable_http_mcp)
app.mount("/{mcp_server_name}/mcp", handle_streamable_http_mcp)
app.mount("/sse", handle_sse_mcp)
app.add_middleware(AuthContextMiddleware)
@@ -821,6 +859,7 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> None:
"""
Set the UserAPIKeyAuth in the auth context variable.
@@ -836,17 +875,17 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
)
auth_context_var.set(auth_user)
def get_auth_context() -> (
Tuple[
Optional[UserAPIKeyAuth],
Optional[str],
Optional[List[str]],
Optional[Dict[str, str]],
]
):
def get_auth_context() -> Tuple[
Optional[UserAPIKeyAuth],
Optional[str],
Optional[List[str]],
Optional[Dict[str, str]],
Optional[Dict[str, str]],
]:
"""
Get the UserAPIKeyAuth from the auth context variable.
@@ -861,8 +900,9 @@ if MCP_AVAILABLE:
auth_user.mcp_auth_header,
auth_user.mcp_servers,
auth_user.mcp_server_auth_headers,
auth_user.oauth2_headers,
)
return None, None, None, None
return None, None, None, None, None
########################################################
############ End of Auth Context Functions #############
litellm/proxy/_new_secret_config.yaml
+13
View File
@@ -15,3 +15,16 @@ model_list:
model: hosted_vllm/whisper-v3
api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5"
api_key: dummy
mcp_servers:
github_mcp:
url: "https://api.githubcopilot.com/mcp"
auth_type: oauth2
authorization_url: https://github.com/login/oauth/authorize
token_url: https://github.com/login/oauth/access_token
client_id: os.environ/GITHUB_OAUTH_CLIENT_ID
client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET
scopes: ["public_repo", "user:email"]
# allowed_tools: ["list_tools"]
# disallowed_tools: ["repo_delete"]
litellm/proxy/auth/oauth2_proxy_hook.py
+2 -2
View File
@@ -15,8 +15,8 @@ async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth:
verbose_proxy_logger.debug("Handling oauth2 proxy request")
# Define the OAuth2 config mappings
oauth2_config_mappings: Dict[str, str] = general_settings.get(
"oauth2_config_mappings", None
)
"oauth2_config_mappings", {}
) or {}
verbose_proxy_logger.debug(f"Oauth2 config mappings: {oauth2_config_mappings}")
if not oauth2_config_mappings:
litellm/proxy/auth/user_api_key_auth.py
+19 -6
View File
@@ -390,7 +390,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
pass_through_endpoints: Optional[List[dict]] = general_settings.get(
"pass_through_endpoints", None
)
passed_in_key: Optional[str] = None
## CHECK IF X-LITELM-API-KEY IS PASSED IN - supercedes Authorization header
api_key, passed_in_key = get_api_key(
custom_litellm_key_header=custom_litellm_key_header,
@@ -502,7 +501,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
end_user_object = result["end_user_object"]
org_id = result["org_id"]
token = result["token"]
team_membership: Optional[LiteLLM_TeamMembership] = result.get("team_membership", None)
team_membership: Optional[LiteLLM_TeamMembership] = result.get(
"team_membership", None
)
global_proxy_spend = await get_global_proxy_spend(
litellm_proxy_admin_name=litellm_proxy_admin_name,
@@ -537,10 +538,22 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
org_id=org_id,
parent_otel_span=parent_otel_span,
end_user_id=end_user_id,
user_tpm_limit=user_object.tpm_limit if user_object is not None else None,
user_rpm_limit=user_object.rpm_limit if user_object is not None else None,
team_member_rpm_limit=team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None,
team_member_tpm_limit=team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None,
user_tpm_limit=(
user_object.tpm_limit if user_object is not None else None
),
user_rpm_limit=(
user_object.rpm_limit if user_object is not None else None
),
team_member_rpm_limit=(
team_membership.safe_get_team_member_rpm_limit()
if team_membership is not None
else None
),
team_member_tpm_limit=(
team_membership.safe_get_team_member_tpm_limit()
if team_membership is not None
else None
),
)
# run through common checks
_ = await common_checks(
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
+7 -7
View File
@@ -363,7 +363,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
prepared_request.headers,
)
response = await self.async_handler.post(
httpx_response = await self.async_handler.post(
url=prepared_request.url,
data=prepared_request.body, # type: ignore
headers=prepared_request.headers, # type: ignore
@@ -373,19 +373,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
#########################################################
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response=response.json(),
guardrail_json_response=httpx_response.json(),
request_data=request_data or {},
guardrail_status=self._get_bedrock_guardrail_response_status(
response=response
response=httpx_response
),
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
)
#########################################################
if response.status_code == 200:
if httpx_response.status_code == 200:
# check if the response was flagged
_json_response = response.json()
_json_response = httpx_response.json()
redacted_response = _redact_pii_matches(_json_response)
verbose_proxy_logger.debug("Bedrock AI response : %s", redacted_response)
bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
@@ -398,8 +398,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
else:
verbose_proxy_logger.error(
"Bedrock AI: error in response. Status code: %s, response: %s",
response.status_code,
response.text,
httpx_response.status_code,
httpx_response.text,
)
return bedrock_guardrail_response
litellm/proxy/hooks/parallel_request_limiter_v3.py
+2 -2
View File
@@ -297,7 +297,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
Group keys by their Redis hash tag to ensure cluster compatibility.
Keys with the same hash tag will be processed together.
"""
groups = {}
groups: Dict[str, List[str]] = {}
for key in keys:
# Extract hash tag from key like "{api_key:sk-123}:requests"
if "{" in key and "}" in key:
@@ -378,7 +378,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
for descriptor in descriptors:
descriptor_key = descriptor["key"]
descriptor_value = descriptor["value"]
rate_limit = descriptor.get("rate_limit", {}) or {}
rate_limit: Optional[RateLimitDescriptorRateLimitObject] = descriptor.get("rate_limit", {}) or {}
requests_limit = rate_limit.get("requests_per_unit")
tokens_limit = rate_limit.get("tokens_per_unit")
max_parallel_requests_limit = rate_limit.get("max_parallel_requests")
litellm/proxy/proxy_server.py
+76 -6
View File
@@ -151,6 +151,9 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
router as mcp_discoverable_endpoints_router,
)
from litellm.proxy._experimental.mcp_server.rest_endpoints import (
router as mcp_rest_endpoints_router,
)
@@ -249,9 +252,7 @@ from litellm.proxy.management_endpoints.customer_endpoints import (
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,
@@ -298,9 +299,7 @@ from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMi
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,
)
@@ -9512,5 +9511,76 @@ app.include_router(ui_discovery_endpoints_router)
########################################################
# MCP Server
########################################################
# Dynamic MCP server routes - handle /{mcp_server_name}/mcp
@app.api_route(
"/{mcp_server_name}/mcp",
methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"],
)
async def dynamic_mcp_route(mcp_server_name: str, request: Request):
"""Handle dynamic MCP server routes like /github_mcp/mcp"""
try:
# Validate that the MCP server exists
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.types.mcp import MCPAuth
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name)
if mcp_server is None:
raise HTTPException(
status_code=404, detail=f"MCP server '{mcp_server_name}' not found"
)
# Create a new scope with the correct path format that the MCP handler expects
# Transform /{mcp_server_name}/mcp to /mcp/{mcp_server_name}
scope = dict(request.scope)
scope["path"] = f"/mcp/{mcp_server_name}"
# Import the MCP handler
from litellm.proxy._experimental.mcp_server.server import (
handle_streamable_http_mcp,
)
# Create a custom send function to capture the response
response_started = False
response_body = b""
response_status = 200
response_headers = []
async def custom_send(message):
nonlocal response_started, response_body, response_status, response_headers
if message["type"] == "http.response.start":
response_started = True
response_status = message["status"]
response_headers = message.get("headers", [])
elif message["type"] == "http.response.body":
response_body += message.get("body", b"")
# Call the existing MCP handler
await handle_streamable_http_mcp(
scope, receive=request.receive, send=custom_send
)
# Return the response
from starlette.responses import Response
headers_dict = {k.decode(): v.decode() for k, v in response_headers}
return Response(
content=response_body,
status_code=response_status,
headers=headers_dict,
media_type=headers_dict.get("content-type", "application/json"),
)
except Exception as e:
verbose_proxy_logger.error(
f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}"
)
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
app.mount(path=BASE_MCP_ROUTE, app=mcp_app)
app.include_router(mcp_rest_endpoints_router)
app.include_router(mcp_discoverable_endpoints_router)
litellm/types/mcp.py
+16 -4
View File
@@ -15,6 +15,7 @@ else:
MCPImageContent = Any
MCPTextContent = Any
class MCPTransport(str, enum.Enum):
sse = "sse"
http = "http"
@@ -26,17 +27,21 @@ class MCPSpecVersion(str, enum.Enum):
mar_2025 = "2025-03-26"
jun_2025 = "2025-06-18"
class MCPAuth(str, enum.Enum):
none = "none"
api_key = "api_key"
bearer_token = "bearer_token"
basic = "basic"
authorization = "authorization"
oauth2 = "oauth2"
# MCP Literals
MCPTransportType = Literal[MCPTransport.sse, MCPTransport.http, MCPTransport.stdio]
MCPSpecVersionType = Literal[MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, MCPSpecVersion.jun_2025]
MCPSpecVersionType = Literal[
MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, MCPSpecVersion.jun_2025
]
MCPAuthType = Optional[
Literal[
MCPAuth.none,
@@ -44,11 +49,11 @@ MCPAuthType = Optional[
MCPAuth.bearer_token,
MCPAuth.basic,
MCPAuth.authorization,
MCPAuth.oauth2,
]
]
class MCPServerCostInfo(TypedDict, total=False):
default_cost_per_query: Optional[float]
"""
@@ -82,6 +87,7 @@ class MCPPreCallRequestObject(BaseModel):
"""
Pydantic object used for MCP pre_call_hook request validation and modification
"""
tool_name: str
arguments: Dict[str, Any]
server_name: Optional[str] = None
@@ -93,6 +99,7 @@ class MCPPreCallResponseObject(BaseModel):
"""
Pydantic object used for MCP pre_call_hook response
"""
should_proceed: bool = True
modified_arguments: Optional[Dict[str, Any]] = None
error_message: Optional[str] = None
@@ -103,6 +110,7 @@ class MCPDuringCallRequestObject(BaseModel):
"""
Pydantic object used for MCP during_call_hook request
"""
tool_name: str
arguments: Dict[str, Any]
server_name: Optional[str] = None
@@ -114,6 +122,7 @@ class MCPDuringCallResponseObject(BaseModel):
"""
Pydantic object used for MCP during_call_hook response
"""
should_continue: bool = True
error_message: Optional[str] = None
hidden_params: HiddenParams = HiddenParams()
@@ -123,5 +132,8 @@ class MCPPostCallResponseObject(BaseModel):
"""
Pydantic object used for MCP post_call_hook response
"""
mcp_tool_call_response: List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]
hidden_params: HiddenParams
mcp_tool_call_response: List[
Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]
]
hidden_params: HiddenParams
litellm/types/mcp_server/mcp_server_manager.py
+6 -1
View File
@@ -6,7 +6,6 @@ from typing_extensions import TypedDict
from litellm.proxy._types import MCPAuthType, MCPTransportType
from litellm.types.mcp import MCPServerCostInfo
# MCPInfo now allows arbitrary additional fields for custom metadata
MCPInfo = Dict[str, Any]
@@ -21,6 +20,12 @@ class MCPServer(BaseModel):
auth_type: Optional[MCPAuthType] = None
authentication_token: Optional[str] = None
mcp_info: Optional[MCPInfo] = None
# OAuth-specific fields
client_id: Optional[str] = None
client_secret: Optional[str] = None
scopes: Optional[List[str]] = None
authorization_url: Optional[str] = None
token_url: Optional[str] = None
# Stdio-specific fields
command: Optional[str] = None
args: Optional[List[str]] = None
proxy_server_config.yaml
+8 -8
View File
@@ -7,14 +7,14 @@ model_list:
id: "1"
- model_name: gpt-3.5-turbo-end-user-test
litellm_params:
model: azure/gpt-4o-new-test
api_base: https://openai-gpt-4-test-v-1.openai.azure.com/
model: azure/gpt-4.1-nano
api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com/
api_version: "2023-05-15"
api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
- model_name: gpt-3.5-turbo
litellm_params:
model: azure/gpt-4o-new-test
api_base: https://openai-gpt-4-test-v-1.openai.azure.com/
model: azure/gpt-4.1-nano
api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com/
api_version: "2023-05-15"
api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
- model_name: gpt-3.5-turbo-large
@@ -26,8 +26,8 @@ model_list:
stream_timeout: 60
- model_name: gpt-4
litellm_params:
model: azure/gpt-4o-new-test
api_base: https://openai-gpt-4-test-v-1.openai.azure.com/
model: azure/gpt-4.1-nano
api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com/
api_version: "2023-05-15"
api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault
rpm: 480
@@ -39,9 +39,9 @@ model_list:
input_cost_per_second: 0.000420
- model_name: text-embedding-ada-002
litellm_params:
model: azure/azure-embedding-model
model: azure/text-embedding-ada-002
api_key: os.environ/AZURE_API_KEY
api_base: https://openai-gpt-4-test-v-1.openai.azure.com/
api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com/
api_version: "2023-05-15"
model_info:
mode: embedding
tests/local_testing/test_health_check.py → tests/litellm_utils_tests/test_health_check.py
+2 -2
View File
@@ -20,7 +20,7 @@ import litellm
async def test_azure_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/chatgpt-v-3",
"model": "azure/gpt-4.1-nano",
"messages": [{"role": "user", "content": "Hey, how's it going?"}],
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
@@ -51,7 +51,7 @@ async def test_text_completion_health_check():
async def test_azure_embedding_health_check():
response = await litellm.ahealth_check(
model_params={
"model": "azure/azure-embedding-model",
"model": "azure/text-embedding-ada-002",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
tests/llm_translation/test_azure_ai.py
+2 -2
View File
@@ -282,8 +282,8 @@ async def test_azure_ai_request_format():
# Set up the test parameters
api_key = os.getenv("AZURE_API_KEY")
api_base = f"{os.getenv('AZURE_API_BASE')}/openai/deployments/gpt-4o-new-test/chat/completions?api-version=2024-08-01-preview"
model = "azure_ai/gpt-4o"
api_base = os.getenv("AZURE_API_BASE")
model = "azure_ai/gpt-4.1-nano"
messages = [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": "Hello! How can I assist you today?"},
tests/llm_translation/test_azure_openai.py
+6 -6
View File
@@ -204,7 +204,7 @@ def test_process_azure_endpoint_url(api_base, model, expected_endpoint):
class TestAzureEmbedding(BaseLLMEmbeddingTest):
def get_base_embedding_call_args(self) -> dict:
return {
"model": "azure/azure-embedding-model",
"model": "azure/text-embedding-ada-002",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
}
@@ -339,7 +339,7 @@ def test_azure_gpt_4o_with_tool_call_and_response_format(api_version):
with patch.object(client.chat.completions.with_raw_response, "create") as mock_post:
response = litellm.completion(
model="azure/gpt-4o-new-test",
model="azure/gpt-4.1-nano",
messages=[
{
"role": "system",
@@ -474,7 +474,7 @@ def test_azure_max_retries_0(
try:
completion(
model="azure/gpt-4o-new-test",
model="azure/gpt-4.1-nano",
messages=[{"role": "user", "content": "Hello world"}],
max_retries=max_retries,
stream=stream,
@@ -502,7 +502,7 @@ async def test_async_azure_max_retries_0(
try:
await acompletion(
model="azure/gpt-4o-new-test",
model="azure/gpt-4.1-nano",
messages=[{"role": "user", "content": "Hello world"}],
max_retries=max_retries,
stream=stream,
@@ -565,7 +565,7 @@ async def test_azure_embedding_max_retries_0(
from litellm import aembedding, embedding
args = {
"model": "azure/azure-embedding-model",
"model": "azure/text-embedding-ada-002",
"input": "Hello world",
"max_retries": max_retries,
}
@@ -598,7 +598,7 @@ def test_azure_safety_result():
litellm._turn_on_debug()
response = completion(
model="azure/gpt-4o-new-test",
model="azure/gpt-4.1-nano",
messages=[{"role": "user", "content": "Hello world"}],
)
print(f"response: {response}")
tests/llm_translation/test_gemini.py
+1 -1
View File
@@ -452,7 +452,7 @@ def test_gemini_finish_reason():
litellm._turn_on_debug()
response = completion(
model="gemini/gemini-1.5-pro",
model="gemini/gemini-2.5-flash-lite",
messages=[{"role": "user", "content": "give me 3 random words"}],
max_tokens=2,
)
tests/local_testing/test_acooldowns_router.py
+3 -7
View File
@@ -143,20 +143,16 @@ async def test_cooldown_same_model_name(sync_mode):
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/chatgpt-v-3",
"model": "gpt-4.1-nano",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"tpm": 90,
},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/chatgpt-v-3",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"model": "gpt-4.1-nano",
"api_key": os.getenv("OPENAI_API_KEY"),
"tpm": 1,
},
},
tests/local_testing/test_caching.py
+16 -16
View File
@@ -325,7 +325,7 @@ def test_caching_with_models_v2():
litellm.set_verbose = True
response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
response3 = completion(model="azure/chatgpt-v-3", messages=messages, caching=True)
response3 = completion(model="gpt-4.1-nano", messages=messages, caching=True)
print(f"response1: {response1}")
print(f"response2: {response2}")
print(f"response3: {response3}")
@@ -527,7 +527,7 @@ def test_embedding_caching_azure():
print(api_key)
print(api_base)
embedding1 = embedding(
model="azure/azure-embedding-model",
model="azure/text-embedding-ada-002",
input=["good morning from litellm", "this is another item"],
api_key=api_key,
api_base=api_base,
@@ -540,7 +540,7 @@ def test_embedding_caching_azure():
time.sleep(1)
start_time = time.time()
embedding2 = embedding(
model="azure/azure-embedding-model",
model="azure/text-embedding-ada-002",
input=["good morning from litellm", "this is another item"],
api_key=api_key,
api_base=api_base,
@@ -595,10 +595,10 @@ async def test_embedding_caching_azure_individual_items():
]
embedding_val_1 = await aembedding(
model="azure/azure-embedding-model", input=embedding_1, caching=True
model="text-embedding-ada-002", input=embedding_1, caching=True
)
embedding_val_2 = await aembedding(
model="azure/azure-embedding-model", input=embedding_2, caching=True
model="text-embedding-ada-002", input=embedding_2, caching=True
)
print(f"embedding_val_2._hidden_params: {embedding_val_2._hidden_params}")
assert embedding_val_2._hidden_params["cache_hit"] == True
@@ -633,11 +633,11 @@ async def test_embedding_caching_azure_individual_items_reordered():
]
embedding_val_1 = await aembedding(
model="azure/azure-embedding-model", input=embedding_1, caching=True
model="text-embedding-ada-002", input=embedding_1, caching=True
)
print("embedding val 1", embedding_val_1)
embedding_val_2 = await aembedding(
model="azure/azure-embedding-model", input=embedding_2, caching=True
model="text-embedding-ada-002", input=embedding_2, caching=True
)
print("embedding val 2", embedding_val_2)
print(f"embedding_val_2._hidden_params: {embedding_val_2._hidden_params}")
@@ -667,7 +667,7 @@ async def test_embedding_caching_base_64():
]
embedding_val_1 = await aembedding(
model="azure/azure-embedding-model",
model="text-embedding-ada-002",
input=inputs,
caching=True,
encoding_format="base64",
@@ -675,7 +675,7 @@ async def test_embedding_caching_base_64():
await asyncio.sleep(5)
print("\n\nCALL2\n\n")
embedding_val_2 = await aembedding(
model="azure/azure-embedding-model",
model="text-embedding-ada-002",
input=inputs,
caching=True,
encoding_format="base64",
@@ -718,7 +718,7 @@ async def test_embedding_caching_redis_ttl():
# Call the embedding method
embedding_val_1 = await litellm.aembedding(
model="azure/azure-embedding-model",
model="text-embedding-ada-002",
input=inputs,
encoding_format="base64",
)
@@ -1226,7 +1226,7 @@ async def test_s3_cache_stream_azure(sync_mode):
if sync_mode:
response1 = litellm.completion(
model="azure/chatgpt-v-3",
model="azure/gpt-4.1-nano",
messages=messages,
max_tokens=40,
temperature=1,
@@ -1239,7 +1239,7 @@ async def test_s3_cache_stream_azure(sync_mode):
print(response_1_content)
else:
response1 = await litellm.acompletion(
model="azure/chatgpt-v-3",
model="azure/gpt-4.1-nano",
messages=messages,
max_tokens=40,
temperature=1,
@@ -1259,7 +1259,7 @@ async def test_s3_cache_stream_azure(sync_mode):
if sync_mode:
response2 = litellm.completion(
model="azure/chatgpt-v-3",
model="azure/gpt-4.1-nano",
messages=messages,
max_tokens=40,
temperature=1,
@@ -1272,7 +1272,7 @@ async def test_s3_cache_stream_azure(sync_mode):
print(response_2_content)
else:
response2 = await litellm.acompletion(
model="azure/chatgpt-v-3",
model="azure/gpt-4.1-nano",
messages=messages,
max_tokens=40,
temperature=1,
@@ -1335,7 +1335,7 @@ async def test_s3_cache_acompletion_azure():
print("s3 Cache: test for caching, streaming + completion")
response1 = await litellm.acompletion(
model="azure/chatgpt-v-3",
model="azure/gpt-4.1-nano",
messages=messages,
max_tokens=40,
temperature=1,
@@ -1345,7 +1345,7 @@ async def test_s3_cache_acompletion_azure():
time.sleep(2)
response2 = await litellm.acompletion(
model="azure/chatgpt-v-3",
model="azure/gpt-4.1-nano",
messages=messages,
max_tokens=40,
temperature=1,
tests/local_testing/test_custom_callback_input.py
-119
View File
@@ -914,125 +914,6 @@ async def test_async_embedding_bedrock():
pytest.fail(f"An exception occurred: {str(e)}")
# asyncio.run(test_async_embedding_bedrock())
# CACHING
## Test Azure - completion, embedding
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_async_completion_azure_caching():
litellm.set_verbose = True
customHandler_caching = CompletionCustomHandler()
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
litellm.callbacks = [customHandler_caching]
unique_time = time.time()
response1 = await litellm.acompletion(
model="azure/chatgpt-v-3",
messages=[
{"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"}
],
caching=True,
)
await asyncio.sleep(1)
print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}")
response2 = await litellm.acompletion(
model="azure/chatgpt-v-3",
messages=[
{"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"}
],
caching=True,
)
await asyncio.sleep(1) # success callbacks are done in parallel
print(
f"customHandler_caching.states post-cache hit: {customHandler_caching.states}"
)
assert len(customHandler_caching.errors) == 0
assert len(customHandler_caching.states) == 4 # pre, post, success, success
@pytest.mark.asyncio
async def test_async_completion_azure_caching_streaming():
import copy
litellm.set_verbose = True
customHandler_caching = CompletionCustomHandler()
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
litellm.callbacks = [customHandler_caching]
unique_time = uuid.uuid4()
response1 = await litellm.acompletion(
model="azure/chatgpt-v-3",
messages=[
{"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"}
],
caching=True,
stream=True,
)
async for chunk in response1:
print(f"chunk in response1: {chunk}")
await asyncio.sleep(1)
initial_customhandler_caching_states = len(customHandler_caching.states)
print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}")
response2 = await litellm.acompletion(
model="azure/chatgpt-v-3",
messages=[
{"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"}
],
caching=True,
stream=True,
)
async for chunk in response2:
print(f"chunk in response2: {chunk}")
await asyncio.sleep(1) # success callbacks are done in parallel
print(
f"customHandler_caching.states post-cache hit: {customHandler_caching.states}"
)
assert len(customHandler_caching.errors) == 0
assert (
len(customHandler_caching.states) > initial_customhandler_caching_states
) # pre, post, streaming .., success, success
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=2)
async def test_async_embedding_azure_caching():
print("Testing custom callback input - Azure Caching")
customHandler_caching = CompletionCustomHandler()
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
litellm.callbacks = [customHandler_caching]
unique_time = time.time()
response1 = await litellm.aembedding(
model="azure/azure-embedding-model",
input=[f"good morning from litellm1 {unique_time}"],
caching=True,
)
await asyncio.sleep(1) # set cache is async for aembedding()
response2 = await litellm.aembedding(
model="azure/azure-embedding-model",
input=[f"good morning from litellm1 {unique_time}"],
caching=True,
)
await asyncio.sleep(1) # success callbacks are done in parallel
print(customHandler_caching.states)
print(customHandler_caching.errors)
assert len(customHandler_caching.errors) == 0
assert len(customHandler_caching.states) == 4 # pre, post, success, success
# Image Generation
tests/local_testing/test_custom_logger.py
+1 -1
View File
@@ -392,7 +392,7 @@ async def test_async_custom_handler_embedding_optional_param():
customHandler_optional_params = MyCustomHandler()
litellm.callbacks = [customHandler_optional_params]
response = await litellm.aembedding(
model="azure/azure-embedding-model", input=["hello world"], user="John"
model="text-embedding-ada-002", input=["hello world"], user="John"
)
await asyncio.sleep(1) # success callback is async
assert customHandler_optional_params.user == "John"
tests/local_testing/test_mock_request.py
+1 -1
View File
@@ -157,7 +157,7 @@ def test_router_mock_request_with_mock_timeout_with_fallbacks():
{
"model_name": "azure-gpt",
"litellm_params": {
"model": "azure/chatgpt-v-3",
"model": "azure/gpt-4.1-nano",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
},
tests/local_testing/test_prometheus_service.py
+5 -7
View File
@@ -92,24 +92,22 @@ async def test_router_with_caching():
"""
try:
def get_azure_params(deployment_name: str):
def get_openai_params():
params = {
"model": f"azure/{deployment_name}",
"api_key": os.environ["AZURE_API_KEY"],
"api_version": os.environ["AZURE_API_VERSION"],
"api_base": os.environ["AZURE_API_BASE"],
"model": "gpt-4.1-nano",
"api_key": os.environ["OPENAI_API_KEY"],
}
return params
model_list = [
{
"model_name": "azure/gpt-4",
"litellm_params": get_azure_params("gpt-4o-new-test"),
"litellm_params": get_openai_params(),
"tpm": 100,
},
{
"model_name": "azure/gpt-4",
"litellm_params": get_azure_params("gpt-4o-new-test"),
"litellm_params": get_openai_params(),
"tpm": 1000,
},
]
tests/local_testing/test_router.py
+10 -18
View File
@@ -334,10 +334,8 @@ async def test_router_retries(sync_mode):
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4o-new-test",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"model": "gpt-4.1-nano",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
]
@@ -473,16 +471,12 @@ def test_reading_key_from_model_list():
try:
print("testing if router raises an exception")
old_api_key = os.environ["AZURE_API_KEY"]
os.environ.pop("AZURE_API_KEY", None)
model_list = [
{
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/chatgpt-v-3",
"api_key": old_api_key,
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"model": "gpt-4.1-nano",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"tpm": 240000,
"rpm": 1800,
@@ -521,10 +515,8 @@ def test_reading_key_from_model_list():
print("\n completed_response", completed_response)
assert len(completed_response) > 0
print("\n Passed Streaming")
os.environ["AZURE_API_KEY"] = old_api_key
router.reset()
except Exception as e:
os.environ["AZURE_API_KEY"] = old_api_key
print(f"FAILED TEST")
pytest.fail(f"Got unexpected exception on router! - {e}")
@@ -544,7 +536,7 @@ def test_call_one_endpoint():
{
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4o-new-test",
"model": "azure/gpt-4.1-nano",
"api_key": old_api_key,
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
@@ -555,7 +547,7 @@ def test_call_one_endpoint():
{
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "azure/azure-embedding-model",
"model": "azure/text-embedding-ada-002",
"api_key": os.environ["AZURE_API_KEY"],
"api_base": os.environ["AZURE_API_BASE"],
},
@@ -574,7 +566,7 @@ def test_call_one_endpoint():
async def call_azure_completion():
response = await router.acompletion(
model="azure/gpt-4o-new-test",
model="azure/gpt-4.1-nano",
messages=[{"role": "user", "content": "hello this request will pass"}],
specific_deployment=True,
)
@@ -582,7 +574,7 @@ def test_call_one_endpoint():
async def call_azure_embedding():
response = await router.aembedding(
model="azure/azure-embedding-model",
model="azure/text-embedding-ada-002",
input=["good morning from litellm"],
specific_deployment=True,
)
@@ -620,7 +612,7 @@ def test_router_azure_acompletion():
{
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4o-new-test",
"model": "azure/gpt-4.1-nano",
"api_key": old_api_key,
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
@@ -1274,7 +1266,7 @@ def test_azure_embedding_on_router():
{
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "azure/azure-embedding-model",
"model": "azure/text-embedding-ada-002",
"api_key": os.environ["AZURE_API_KEY"],
"api_base": os.environ["AZURE_API_BASE"],
},
tests/local_testing/test_router_budget_limiter.py
+1 -1
View File
@@ -74,7 +74,7 @@ async def test_provider_budgets_e2e_test():
{
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/chatgpt-v-3",
"model": "azure/gpt-4.1-nano",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
tests/local_testing/test_router_caching.py
+11 -23
View File
@@ -6,7 +6,7 @@ import sys
import time
import traceback
from unittest.mock import patch
from typing import Union
import pytest
sys.path.insert(
@@ -85,21 +85,11 @@ async def test_acompletion_caching_on_router():
litellm.set_verbose = True
model_list = [
{
"model_name": "gpt-3.5-turbo",
"model_name": "gpt-4.1-nano",
"litellm_params": {
"model": "gpt-3.5-turbo",
"model": "gpt-4.1-nano",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"tpm": 100000,
"rpm": 10000,
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/chatgpt-v-3",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"mock_response": "Hello world",
},
"tpm": 100000,
"rpm": 10000,
@@ -120,13 +110,13 @@ async def test_acompletion_caching_on_router():
routing_strategy="simple-shuffle",
)
response1 = await router.acompletion(
model="gpt-3.5-turbo", messages=messages, temperature=1
model="gpt-4.1-nano", messages=messages, temperature=1
)
print(f"response1: {response1}")
await asyncio.sleep(5) # add cache is async, async sleep for cache to get set
response2 = await router.acompletion(
model="gpt-3.5-turbo", messages=messages, temperature=1
model="gpt-4.1-nano", messages=messages, temperature=1
)
print(f"response2: {response2}")
assert response1.id == response2.id
@@ -213,10 +203,8 @@ async def test_acompletion_caching_with_ttl_on_router():
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/chatgpt-v-3",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"model": "gpt-4.1-nano",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"tpm": 100000,
"rpm": 10000,
@@ -279,7 +267,7 @@ async def test_acompletion_caching_on_router_caching_groups():
{
"model_name": "azure-gpt-3.5-turbo",
"litellm_params": {
"model": "azure/chatgpt-v-3",
"model": "azure/gpt-4.1-nano",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
@@ -343,8 +331,8 @@ async def test_acompletion_caching_on_router_caching_groups():
],
)
def test_create_correct_redis_cache_instance(
startup_nodes: list[dict] | None,
expected_cache_type: type[RedisClusterCache | RedisCache],
startup_nodes: Union[list[dict], None],
expected_cache_type: Union[type[RedisClusterCache], type[RedisCache]],
):
cache_config = dict(
host="mockhost",
tests/local_testing/test_router_cooldown_handlers.py
+1 -1
View File
@@ -44,7 +44,7 @@ async def test_cooldown_badrequest_error():
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/chatgpt-v-3",
"model": "azure/gpt-4.1-nano",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
tests/local_testing/test_router_fallbacks.py
+7 -13
View File
@@ -250,10 +250,8 @@ def test_sync_fallbacks_embeddings():
{ # list of model deployments
"model_name": "good-azure-embedding-model", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/azure-embedding-model",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"model": "text-embedding-ada-002",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"tpm": 240000,
"rpm": 1800,
@@ -302,10 +300,8 @@ async def test_async_fallbacks_embeddings():
{ # list of model deployments
"model_name": "good-azure-embedding-model", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/azure-embedding-model",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"model": "text-embedding-ada-002",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"tpm": 240000,
"rpm": 1800,
@@ -993,10 +989,8 @@ async def test_service_unavailable_fallbacks(sync_mode):
{
"model_name": "gpt-3.5-turbo-0125-preview",
"litellm_params": {
"model": "azure/chatgpt-v-3",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
"model": "gpt-4.1-nano",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
],
@@ -1014,7 +1008,7 @@ async def test_service_unavailable_fallbacks(sync_mode):
messages=[{"role": "user", "content": "Hey, how's it going?"}],
)
assert response.model == "gpt-3.5-turbo-0125"
assert "gpt-4.1-nano" in response.model
@pytest.mark.parametrize("sync_mode", [True, False])
tests/local_testing/test_router_get_deployments.py
+1 -1
View File
@@ -397,7 +397,7 @@ def test_usage_based_routing():
"model": f"azure/{deployment_name}",
"api_key": os.environ["AZURE_API_KEY"],
"api_version": os.environ["AZURE_API_VERSION"],
"api_base": os.environ["AZURE_API_BASE"],
"api_base": "https://fake-api.openai.com/v1",
}
return params
tests/local_testing/test_router_policy_violation.py
-137
View File
@@ -1,137 +0,0 @@
#### What this tests ####
# This tests if the router sends back a policy violation, without retries
import sys, os, time
import traceback, asyncio
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
from litellm import Router
from litellm.integrations.custom_logger import CustomLogger
class MyCustomHandler(CustomLogger):
success: bool = False
failure: bool = False
previous_models: int = 0
def log_pre_api_call(self, model, messages, kwargs):
print(f"Pre-API Call")
print(
f"previous_models: {kwargs['litellm_params']['metadata']['previous_models']}"
)
self.previous_models += len(
kwargs["litellm_params"]["metadata"]["previous_models"]
) # {"previous_models": [{"model": litellm_model_name, "exception_type": AuthenticationError, "exception_string": <complete_traceback>}]}
print(f"self.previous_models: {self.previous_models}")
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
print(
f"Post-API Call - response object: {response_obj}; model: {kwargs['model']}"
)
def log_stream_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Stream")
def async_log_stream_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Stream")
def log_success_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Success")
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Success")
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Failure")
kwargs = {
"model": "azure/gpt-3.5-turbo",
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{
"role": "user",
"content": "vorrei vedere la cosa più bella ad Ercolano. Qual’è?",
},
],
}
@pytest.mark.asyncio
async def test_async_fallbacks():
litellm.set_verbose = False
model_list = [
{ # list of model deployments
"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/chatgpt-v-3",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
},
{
"model_name": "azure/gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/chatgpt-functioncalling",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
},
{
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"tpm": 1000000,
"rpm": 9000,
},
{
"model_name": "gpt-3.5-turbo-16k", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "gpt-3.5-turbo-16k",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"tpm": 1000000,
"rpm": 9000,
},
]
router = Router(
model_list=model_list,
num_retries=3,
fallbacks=[{"azure/gpt-3.5-turbo": ["gpt-3.5-turbo"]}],
# context_window_fallbacks=[
# {"azure/gpt-3.5-turbo-context-fallback": ["gpt-3.5-turbo-16k"]},
# {"gpt-3.5-turbo": ["gpt-3.5-turbo-16k"]},
# ],
set_verbose=False,
)
customHandler = MyCustomHandler()
litellm.callbacks = [customHandler]
try:
response = await router.acompletion(**kwargs)
pytest.fail(
f"An exception occurred: {e}"
) # should've raised azure policy error
except litellm.Timeout as e:
pass
except Exception as e:
await asyncio.sleep(
0.05
) # allow a delay as success_callbacks are on a separate thread
assert customHandler.previous_models == 0 # 0 retries, 0 fallback
router.reset()
finally:
router.reset()
tests/logging_callback_tests/test_custom_callback_router.py
+113 -11
View File
@@ -393,21 +393,21 @@ async def test_async_chat_azure():
litellm.set_verbose = True
model_list = [
{
"model_name": "gpt-3.5-turbo", # openai model name
"model_name": "gpt-4.1-nano", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/gpt-4o-new-test",
"model": "azure/gpt-4.1-nano",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
},
"model_info": {"base_model": "azure/gpt-4-1106-preview"},
"model_info": {"base_model": "azure/gpt-4.1-nano"},
"tpm": 240000,
"rpm": 1800,
},
]
router = Router(model_list=model_list, num_retries=0) # type: ignore
response = await router.acompletion(
model="gpt-3.5-turbo",
model="gpt-4.1-nano",
messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}],
)
print("got response, sleeping 5 seconds....")
@@ -422,7 +422,7 @@ async def test_async_chat_azure():
litellm.callbacks = [customHandler_streaming_azure_router]
router2 = Router(model_list=model_list, num_retries=0) # type: ignore
response = await router2.acompletion(
model="gpt-3.5-turbo",
model="gpt-4.1-nano",
messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}],
stream=True,
)
@@ -471,7 +471,6 @@ async def test_async_chat_azure():
pytest.fail(f"An exception occurred - {str(e)}")
# asyncio.run(test_async_chat_azure())
## EMBEDDING
@pytest.mark.asyncio
async def test_async_embedding_azure():
@@ -483,7 +482,7 @@ async def test_async_embedding_azure():
{
"model_name": "azure-embedding-model", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/azure-embedding-model",
"model": "azure/text-embedding-ada-002",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
@@ -606,9 +605,9 @@ async def test_async_completion_azure_caching():
unique_time = time.time()
model_list = [
{
"model_name": "gpt-3.5-turbo", # openai model name
"model_name": "gpt-4.1-nano", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/chatgpt-v-3",
"model": "azure/gpt-4.1-nano",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
@@ -627,7 +626,7 @@ async def test_async_completion_azure_caching():
]
router = Router(model_list=model_list) # type: ignore
response1 = await router.acompletion(
model="gpt-3.5-turbo",
model="gpt-4.1-nano",
messages=[
{"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"}
],
@@ -636,7 +635,7 @@ async def test_async_completion_azure_caching():
await asyncio.sleep(1)
print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}")
response2 = await router.acompletion(
model="gpt-3.5-turbo",
model="gpt-4.1-nano",
messages=[
{"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"}
],
@@ -650,6 +649,108 @@ async def test_async_completion_azure_caching():
assert len(customHandler_caching.states) == 4 # pre, post, success, success
@pytest.mark.asyncio
async def test_async_completion_azure_caching_streaming():
import copy
import uuid
litellm.set_verbose = True
customHandler_caching = CompletionCustomHandler()
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
litellm.callbacks = [customHandler_caching]
unique_time = uuid.uuid4()
# Use Router instead of direct litellm.acompletion to get router-specific metadata
model_list = [
{
"model_name": "gpt-4.1-nano",
"litellm_params": {
"model": "azure/gpt-4.1-nano",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
},
"tpm": 240000,
"rpm": 1800,
},
]
router = Router(model_list=model_list)
response1 = await router.acompletion(
model="gpt-4.1-nano",
messages=[
{"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"}
],
caching=True,
stream=True,
)
async for chunk in response1:
print(f"chunk in response1: {chunk}")
await asyncio.sleep(1)
initial_customhandler_caching_states = len(customHandler_caching.states)
print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}")
response2 = await router.acompletion(
model="gpt-4.1-nano",
messages=[
{"role": "user", "content": f"Hi 👋 - i'm async azure {unique_time}"}
],
caching=True,
stream=True,
)
async for chunk in response2:
print(f"chunk in response2: {chunk}")
await asyncio.sleep(1) # success callbacks are done in parallel
print(
f"customHandler_caching.states post-cache hit: {customHandler_caching.states}"
)
assert len(customHandler_caching.errors) == 0
assert (
len(customHandler_caching.states) > initial_customhandler_caching_states
) # pre, post, streaming .., success, success
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=2)
async def test_async_embedding_azure_caching():
print("Testing custom callback input - Azure Caching")
customHandler_caching = CompletionCustomHandler()
litellm.cache = Cache(
type="redis",
host=os.environ["REDIS_HOST"],
port=os.environ["REDIS_PORT"],
password=os.environ["REDIS_PASSWORD"],
)
router = Router(model_list=[{
"model_name": "text-embedding-ada-002",
"litellm_params": {
"model": "openai/text-embedding-ada-002",
},
}])
litellm.callbacks = [customHandler_caching]
unique_time = time.time()
response1 = await router.aembedding(
model="text-embedding-ada-002",
input=[f"good morning from litellm1 {unique_time}"],
caching=True,
)
await asyncio.sleep(1) # set cache is async for aembedding()
response2 = await router.aembedding(
model="text-embedding-ada-002",
input=[f"good morning from litellm1 {unique_time}"],
caching=True,
)
await asyncio.sleep(1) # success callbacks are done in parallel
print(customHandler_caching.states)
print(customHandler_caching.errors)
assert len(customHandler_caching.errors) == 0
assert len(customHandler_caching.states) == 4 # pre, post, success, success
@pytest.mark.asyncio
async def test_rate_limit_error_callback():
"""
@@ -717,3 +818,4 @@ async def test_rate_limit_error_callback():
assert "original_model_group" in mock_client.call_args.kwargs
assert mock_client.call_args.kwargs["original_model_group"] == "my-test-gpt"
tests/logging_callback_tests/test_unit_tests_init_callbacks.py
+24 -1
View File
@@ -102,6 +102,26 @@ async def use_callback_in_llm_call(
elif callback == "openmeter":
# it's currently handled in jank way, TODO: fix openmete and then actually run it's test
return
elif callback == "bitbucket":
# Set up mock bitbucket configuration required for initialization
litellm.global_bitbucket_config = {
"workspace": "test-workspace",
"repository": "test-repo",
"access_token": "test-token",
"branch": "main"
}
# Mock BitBucket HTTP calls to prevent actual API requests
import httpx
from unittest.mock import MagicMock
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"values": []}
mock_response.text = ""
patch.object(
litellm.module_level_client, "get", return_value=mock_response
).start()
elif callback == "prometheus":
# pytest teardown - clear existing prometheus collectors
collectors = list(REGISTRY._collector_to_names.keys())
@@ -175,7 +195,10 @@ async def use_callback_in_llm_call(
if callback == "argilla":
patch.stopall()
if callback == "argilla":
if callback == "bitbucket":
# Clean up bitbucket configuration and patches
if hasattr(litellm, 'global_bitbucket_config'):
delattr(litellm, 'global_bitbucket_config')
patch.stopall()
tests/mcp_tests/test_mcp_server.py
+21 -21
View File
@@ -98,9 +98,9 @@ async def test_mcp_server_manager_https_server():
assert tools[0].name == f"{expected_prefix}-gmail_send_email"
# Manually set up the tool mapping for the call_tool test
mcp_server_manager.tool_name_to_mcp_server_name_mapping[
"gmail_send_email"
] = expected_prefix
mcp_server_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = (
expected_prefix
)
mcp_server_manager.tool_name_to_mcp_server_name_mapping[
f"{expected_prefix}-gmail_send_email"
] = expected_prefix
@@ -260,9 +260,9 @@ async def test_mcp_http_transport_call_tool_mock():
)
# Manually set up tool mapping (normally done by list_tools)
test_manager.tool_name_to_mcp_server_name_mapping[
"gmail_send_email"
] = "test_http_server"
test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = (
"test_http_server"
)
# Call the tool
result = await test_manager.call_tool(
@@ -326,9 +326,9 @@ async def test_mcp_http_transport_call_tool_error_mock():
)
# Manually set up tool mapping
test_manager.tool_name_to_mcp_server_name_mapping[
"gmail_send_email"
] = "test_http_server"
test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = (
"test_http_server"
)
# Call the tool with invalid data
result = await test_manager.call_tool(
@@ -792,18 +792,18 @@ async def test_get_tools_from_mcp_servers():
assert len(result) == 1, "Should only return tools from server1"
assert result[0].name == "tool1", "Should return tool from server1"
# Test Case 2: Without specific MCP servers
# Create a different mock manager for the second test case
mock_manager_2 = AsyncMock()
mock_manager_2.get_allowed_mcp_servers = AsyncMock(
return_value=["server1_id", "server2_id"]
)
mock_manager_2.get_mcp_server_by_id = mock_get_server_by_id
mock_manager_2._get_tools_from_server = AsyncMock(
side_effect=lambda server, mcp_auth_header=None: [mock_tool_1]
if server.server_id == "server1_id"
else [mock_tool_2]
)
# Test Case 2: Without specific MCP servers
# Create a different mock manager for the second test case
mock_manager_2 = AsyncMock()
mock_manager_2.get_allowed_mcp_servers = AsyncMock(
return_value=["server1_id", "server2_id"]
)
mock_manager_2.get_mcp_server_by_id = mock_get_server_by_id
mock_manager_2._get_tools_from_server = AsyncMock(
side_effect=lambda server, mcp_auth_header=None, extra_headers=None: (
[mock_tool_1] if server.server_id == "server1_id" else [mock_tool_2]
)
)
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
tests/proxy_unit_tests/test_configs/test_custom_logger.yaml
+4 -8
View File
@@ -1,20 +1,16 @@
model_list:
- model_name: Azure OpenAI GPT-4 Canada
litellm_params:
model: azure/chatgpt-v-3
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_version: "2023-07-01-preview"
model: gpt-4.1-nano
api_key: os.environ/OPENAI_API_KEY
model_info:
mode: chat
input_cost_per_token: 0.0002
id: gm
- model_name: azure-embedding-model
litellm_params:
model: azure/azure-embedding-model
api_base: os.environ/AZURE_API_BASE
api_key: os.environ/AZURE_API_KEY
api_version: "2023-07-01-preview"
model: text-embedding-ada-002
api_key: os.environ/OPENAI_API_KEY
model_info:
mode: embedding
input_cost_per_token: 0.002
tests/proxy_unit_tests/test_proxy_custom_logger.py
+1 -1
View File
@@ -164,7 +164,7 @@ def test_chat_completion(client):
my_custom_logger.async_success == True
) # checks if the status of async_success is True, only the async_log_success_event can set this to true
assert (
my_custom_logger.async_completion_kwargs["model"] == "chatgpt-v-3"
my_custom_logger.async_completion_kwargs["model"] == "gpt-4.1-nano"
) # checks if kwargs passed to async_log_success_event are correct
print(
"\n\n Custom Logger Async Completion args",
tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py
+26 -11
View File
@@ -1,8 +1,19 @@
import sys
import os
import sys
import pytest
from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig
from litellm.types.rerank import OptionalRerankParams, RerankResponse, RerankResponseResult, RerankResponseMeta, RerankBilledUnits, RerankTokens, RerankResponseDocument
from litellm.types.rerank import (
OptionalRerankParams,
RerankBilledUnits,
RerankResponse,
RerankResponseDocument,
RerankResponseMeta,
RerankResponseResult,
RerankTokens,
)
class TestHostedVLLMRerankTransform:
def setup_method(self):
@@ -27,23 +38,27 @@ class TestHostedVLLMRerankTransform:
assert params["return_documents"] is True
def test_map_cohere_rerank_params_raises_on_max_chunks_per_doc(self):
with pytest.raises(ValueError, match="Hosted VLLM does not support max_chunks_per_doc"):
with pytest.raises(
ValueError, match="Hosted VLLM does not support max_chunks_per_doc"
):
self.config.map_cohere_rerank_params(
non_default_params=None,
model=self.model,
drop_params=False,
query="test query",
documents=["doc1"],
max_chunks_per_doc=5
max_chunks_per_doc=5,
)
def test_get_complete_url(self):
base = "https://api.example.com"
url = self.config.get_complete_url(base, self.model)
assert url == "https://api.example.com/v1/rerank"
# Already ends with /v1/rerank
url2 = self.config.get_complete_url("https://api.example.com/v1/rerank", self.model)
assert url2 == "https://api.example.com/v1/rerank"
assert url == "https://api.example.com/rerank"
# Already ends with /rerank
url2 = self.config.get_complete_url(
"https://api.example.com/rerank", self.model
)
assert url2 == "https://api.example.com/rerank"
# Raises if api_base is None
with pytest.raises(ValueError):
self.config.get_complete_url(None, self.model)
@@ -55,7 +70,7 @@ class TestHostedVLLMRerankTransform:
{"index": 0, "relevance_score": 0.9, "document": {"text": "doc1 text"}},
{"index": 1, "relevance_score": 0.7, "document": {"text": "doc2 text"}},
],
"usage": {"total_tokens": 42}
"usage": {"total_tokens": 42},
}
result = self.config._transform_response(response_dict)
assert result.id == "abc123"
@@ -75,7 +90,7 @@ class TestHostedVLLMRerankTransform:
response_dict = {
"id": "abc123",
"results": [{"relevance_score": 0.5}],
"usage": {"total_tokens": 10}
"usage": {"total_tokens": 10},
}
with pytest.raises(ValueError, match="Missing required fields in the result="):
self.config._transform_response(response_dict)
self.config._transform_response(response_dict)
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+6
View File
@@ -364,6 +364,7 @@ class TestMCPRequestHandler:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@@ -543,6 +544,7 @@ class TestMCPRequestHandler:
mcp_auth_header,
mcp_servers_result,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
assert auth_result == mock_auth_result
assert mcp_auth_header == expected_result["mcp_auth"]
@@ -716,6 +718,7 @@ class TestMCPCustomHeaderName:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@@ -886,6 +889,7 @@ class TestMCPAccessGroupsE2E:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@@ -935,6 +939,7 @@ class TestMCPAccessGroupsE2E:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@@ -963,6 +968,7 @@ def test_mcp_path_based_server_segregation(monkeypatch):
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = get_auth_context()
# Capture the MCP servers for testing
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+118 -2
View File
@@ -102,7 +102,9 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails():
working_server if server_id == "working_server" else failing_server
)
async def mock_get_tools_from_server(server, mcp_auth_header=None):
async def mock_get_tools_from_server(
server, mcp_auth_header=None, extra_headers=None
):
if server.name == "working_server":
# Working server returns tools
tool1 = MagicMock()
@@ -184,7 +186,9 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing():
failing_server1 if server_id == "failing_server1" else failing_server2
)
async def mock_get_tools_from_server(server, mcp_auth_header=None):
async def mock_get_tools_from_server(
server, mcp_auth_header=None, extra_headers=None
):
# All servers fail
raise Exception(f"Server {server.name} connection failed")
@@ -448,3 +452,115 @@ async def test_mcp_routing_with_conflicting_alias_and_group_name():
assert (
called_servers[0].server_id == specific_server.server_id
), "Should have contacted the specific server alias, not the group."
@pytest.mark.asyncio
async def test_oauth2_headers_passed_to_mcp_client():
"""Test that OAuth2 headers are properly passed through to the MCP client for OAuth2 servers like github_mcp"""
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.server import (
_get_tools_from_mcp_servers,
set_auth_context,
)
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
except ImportError:
pytest.skip("MCP server not available")
# Clear the registry to avoid conflicts with other tests
global_mcp_server_manager.registry.clear()
# Create an OAuth2 MCP server similar to github_mcp configuration
oauth2_server = MCPServer(
server_id="github_mcp_server_id",
name="github_mcp",
alias="github_mcp",
transport=MCPTransport.http,
url="https://api.githubcopilot.com/mcp",
auth_type=MCPAuth.oauth2,
client_id="test_github_client_id",
client_secret="test_github_client_secret",
scopes=["public_repo", "user:email"],
authorization_url="https://github.com/login/oauth/authorize",
token_url="https://github.com/login/oauth/access_token",
)
global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server
# Mock user auth
user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
# Set up OAuth2 headers that would come from the client
oauth2_headers = {"Authorization": "Bearer github_oauth_token_12345"}
# Set auth context with OAuth2 headers
set_auth_context(user_api_key_auth=user_api_key_auth, oauth2_headers=oauth2_headers)
# This will capture the arguments passed to _create_mcp_client
captured_client_args = {}
def mock_create_mcp_client(server, mcp_auth_header=None, extra_headers=None):
# Capture the arguments for verification
captured_client_args.update(
{
"server": server,
"mcp_auth_header": mcp_auth_header,
"extra_headers": extra_headers,
}
)
# Return a mock client that doesn't actually connect
mock_client = MagicMock()
mock_client.disconnect = AsyncMock()
return mock_client
# Mock _fetch_tools_with_timeout to avoid actual network calls
async def mock_fetch_tools_with_timeout(client, server_name):
return [] # Return empty list of tools
with patch.object(
global_mcp_server_manager,
"_create_mcp_client",
side_effect=mock_create_mcp_client,
) as mock_create_client, patch.object(
global_mcp_server_manager,
"_fetch_tools_with_timeout",
side_effect=mock_fetch_tools_with_timeout,
), patch.object(
global_mcp_server_manager,
"get_allowed_mcp_servers",
AsyncMock(return_value=[oauth2_server.server_id]),
):
# Call _get_tools_from_mcp_servers which should eventually call _create_mcp_client
await _get_tools_from_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=None,
mcp_servers=None, # Will use all allowed servers
oauth2_headers=oauth2_headers,
)
# Verify that _create_mcp_client was called
assert (
mock_create_client.call_count == 1
), "Expected _create_mcp_client to be called once"
# Verify the server passed to _create_mcp_client is the OAuth2 server
assert captured_client_args["server"].server_id == oauth2_server.server_id
assert captured_client_args["server"].auth_type == MCPAuth.oauth2
# Most importantly: verify that OAuth2 headers were passed as extra_headers
assert (
captured_client_args["extra_headers"] is not None
), "Expected extra_headers to be passed for OAuth2 server"
assert (
captured_client_args["extra_headers"] == oauth2_headers
), f"Expected OAuth2 headers to be passed as extra_headers, got {captured_client_args['extra_headers']}"
# Verify the Authorization header specifically
assert "Authorization" in captured_client_args["extra_headers"]
assert (
captured_client_args["extra_headers"]["Authorization"]
== "Bearer github_oauth_token_12345"
)
tests/test_models.py
+7 -13
View File
@@ -83,10 +83,8 @@ async def add_models(
data = {
"model_name": model_name,
"litellm_params": {
"model": "azure/chatgpt-v-3",
"api_key": "os.environ/AZURE_API_KEY",
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"api_version": "2023-05-15",
"model": "openai/gpt-4.1-nano",
"api_key": "os.environ/OPENAI_API_KEY",
},
"model_info": {"id": model_id},
}
@@ -119,10 +117,8 @@ async def update_model(
data = {
"model_name": model_name,
"litellm_params": {
"model": "azure/chatgpt-v-3",
"api_key": "os.environ/AZURE_API_KEY",
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"api_version": "2023-05-15",
"model": "openai/gpt-4.1-nano",
"api_key": "os.environ/OPENAI_API_KEY",
},
"model_info": {"id": model_id},
}
@@ -311,10 +307,8 @@ async def add_model_for_health_checking(session, model_id="123"):
data = {
"model_name": f"azure-model-health-check-{model_id}",
"litellm_params": {
"model": "azure/chatgpt-v-3",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"api_version": "2023-05-15",
"model": "gpt-4.1-nano",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"model_info": {"id": model_id},
}
@@ -436,7 +430,7 @@ async def test_add_model_run_health():
assert _health_info["healthy_count"] == 1
assert (
_healthy_endpooint["model"] == "azure/chatgpt-v-3"
_healthy_endpooint["model"] == "gpt-4.1-nano"
) # this is the model that got added
# assert httpx client is is unchanges