mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
Merge branch 'main' into litellm_staging_branch_10_11_2025_p1
This commit is contained in:
Krish Dholakia
GitHub
commit
04dc9d091c
17 changed files with 371 additions and 38 deletions
|
|
@ -180,11 +180,11 @@ def completion(
|
|||
|
||||
- `function`: *object* - Required.
|
||||
|
||||
- `tool_choice`: *string or object (optional)* - Controls which (if any) function is called by the model. none means the model will not call a function and instead generates a message. auto means the model can pick between generating a message or calling a function. Specifying a particular function via `{"type: "function", "function": {"name": "my_function"}}` forces the model to call that function.
|
||||
- `tool_choice`: *string or object (optional)* - Controls which (if any) function is called by the model. none means the model will not call a function and instead generates a message. auto means the model can pick between generating a message or calling a function. Specifying a particular function via `{"type": "function", "function": {"name": "my_function"}}` forces the model to call that function.
|
||||
|
||||
- `none` is the default when no functions are present. `auto` is the default if functions are present.
|
||||
|
||||
- `parallel_tool_calls`: *boolean (optional)* - Whether to enable parallel function calling during tool use.. OpenAI default is true.
|
||||
- `parallel_tool_calls`: *boolean (optional)* - Whether to enable parallel function calling during tool use. OpenAI default is true.
|
||||
|
||||
- `frequency_penalty`: *number or null (optional)* - It is used to penalize new tokens based on their frequency in the text so far.
|
||||
|
||||
|
|
|
|||
|
|
@ -1204,6 +1204,8 @@ mcp_servers:
|
|||
scopes: ["public_repo", "user:email"]
|
||||
```
|
||||
|
||||
[**See Claude Code Tutorial**](./tutorials/claude_responses_api#connecting-mcp-servers)
|
||||
|
||||
## Using your MCP with client side credentials
|
||||
|
||||
Use this if you want to pass a client side authentication token to LiteLLM to then pass to your MCP to auth to your MCP.
|
||||
|
|
|
|||
|
|
@ -209,4 +209,81 @@ claude --model claude-bedrock
|
|||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
<Image img={require('../../img/release_notes/claude_code_demo.png')} style={{ width: '500px', height: 'auto' }} />
|
||||
<Image img={require('../../img/release_notes/claude_code_demo.png')} style={{ width: '500px', height: 'auto' }} />
|
||||
|
||||
|
||||
## Connecting MCP Servers
|
||||
|
||||
You can also connect MCP servers to Claude Code via LiteLLM Proxy.
|
||||
|
||||
:::note
|
||||
|
||||
Limitations:
|
||||
|
||||
- Currently, only HTTP MCP servers are supported
|
||||
- Does not work in Cursor IDE yet.
|
||||
|
||||
:::
|
||||
|
||||
1. Add the MCP server to your `config.yaml`
|
||||
|
||||
In this example, we'll add the Github MCP server to our `config.yaml`
|
||||
|
||||
```yaml title="config.yaml" showLineNumbers
|
||||
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"]
|
||||
```
|
||||
|
||||
2. Start LiteLLM Proxy
|
||||
|
||||
```bash
|
||||
litellm --config /path/to/config.yaml
|
||||
|
||||
# RUNNING on http://0.0.0.0:4000
|
||||
```
|
||||
|
||||
3. Use the MCP server in Claude Code
|
||||
|
||||
```bash
|
||||
claude mcp add --transport http litellm_proxy http://0.0.0.0:4000 --header "Authorization: Bearer sk-LITELLM_VIRTUAL_KEY"
|
||||
```
|
||||
|
||||
4. Authenticate via Claude Code
|
||||
|
||||
a. Start Claude Code
|
||||
|
||||
```bash
|
||||
claude
|
||||
```
|
||||
|
||||
b. Authenticate via Claude Code
|
||||
|
||||
```bash
|
||||
/mcp
|
||||
```
|
||||
|
||||
c. Select the MCP server
|
||||
|
||||
```bash
|
||||
> litellm_proxy
|
||||
```
|
||||
|
||||
d. Start Oauth flow via Claude Code
|
||||
|
||||
```bash
|
||||
> 1. Authenticate
|
||||
2. Reconnect
|
||||
3. Disable
|
||||
```
|
||||
|
||||
e. Once completed, you should see this success message:
|
||||
|
||||
<Image img={require('../../img/oauth_2_success.png')} style={{ width: '500px', height: 'auto' }} />
|
||||
|
||||
|
|
|
|||
BIN
docs/my-website/img/mcp_updates.jpg
Normal file
BIN
docs/my-website/img/mcp_updates.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 913 KiB |
BIN
docs/my-website/img/oauth_2_success.png
Normal file
BIN
docs/my-website/img/oauth_2_success.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 67 KiB |
BIN
docs/my-website/img/release_notes/1_78_0_perf.png
Normal file
BIN
docs/my-website/img/release_notes/1_78_0_perf.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 134 KiB |
BIN
docs/my-website/img/release_notes/tool_control.png
Normal file
BIN
docs/my-website/img/release_notes/tool_control.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 798 KiB |
|
|
@ -57,19 +57,6 @@ pip install litellm==1.77.5
|
|||
|
||||
---
|
||||
|
||||
### Scheduled Key Rotations
|
||||
|
||||
<Image img={require('../../img/release_notes/schedule_key_rotations.png')} style={{ width: '800px', height: 'auto' }} />
|
||||
|
||||
<br/>
|
||||
|
||||
This release brings support for scheduling virtual key rotations on LiteLLM AI Gateway.
|
||||
|
||||
This is great for Proxy Admins looking to enforce Enterprise Grade security for use cases going through LiteLLM AI Gateway.
|
||||
|
||||
From this release you can enforce Virtual Keys to rotate on a schedule of your choice e.g every 15 days/30 days/60 days etc.
|
||||
|
||||
---
|
||||
### Performance Improvements - 54% RPS Improvement
|
||||
|
||||
<Image img={require('../../img/release_notes/perf_77_5.png')} style={{ width: '800px', height: 'auto' }} />
|
||||
|
|
|
|||
|
|
@ -103,6 +103,31 @@ View the complete configuration: [gist.github.com/AlexsanderHamir/config.yaml](h
|
|||
|
||||
View the complete load testing script: [gist.github.com/AlexsanderHamir/no_cache_hits.py](https://gist.github.com/AlexsanderHamir/42c33d7a4dc7a57f56a78b560dee3a42)
|
||||
|
||||
### MCP OAuth 2.0 Support
|
||||
|
||||
<Image img={require('../../img/mcp_updates.jpg')} style={{ width: '800px', height: 'auto' }} />
|
||||
|
||||
<br/>
|
||||
|
||||
This release adds support for OAuth 2.0 Client Credentials for MCP servers. This is great for **Internal Dev Tools** use-cases, as it enables your users to call MCP servers, with their own credentials. E.g. Allowing your developers to call the Github MCP, with their own credentials.
|
||||
|
||||
[Set it up today on Claude Code](../../docs/tutorials/claude_responses_api#connecting-mcp-servers)
|
||||
|
||||
### Scheduled Key Rotations
|
||||
|
||||
<Image img={require('../../img/release_notes/schedule_key_rotations.png')} style={{ width: '800px', height: 'auto' }} />
|
||||
|
||||
<br/>
|
||||
|
||||
This release brings support for scheduling virtual key rotations on LiteLLM AI Gateway.
|
||||
|
||||
From this release you can enforce Virtual Keys to rotate on a schedule of your choice e.g every 15 days/30 days/60 days etc.
|
||||
|
||||
This is great for Proxy Admins who need to enforce security policies for production workloads.
|
||||
|
||||
[Get Started](../../docs/proxy/virtual_keys#scheduled-key-rotations)
|
||||
|
||||
|
||||
## New Models / Updated Models
|
||||
|
||||
#### New Model Support
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
---
|
||||
title: "[Preview] v1.78.0-stable - MCP Gateway: Control Tool Access by Team/Key"
|
||||
title: "[Preview] v1.78.0-stable - MCP Gateway: Control Tool Access by Team, Key"
|
||||
slug: "v1-78-0"
|
||||
date: 2025-10-11T10:00:00
|
||||
authors:
|
||||
|
|
@ -37,6 +37,10 @@ import TabItem from '@theme/TabItem';
|
|||
<TabItem value="docker" label="Docker">
|
||||
|
||||
``` showLineNumbers title="docker run litellm"
|
||||
docker run \
|
||||
-e STORE_MODEL_IN_DB=True \
|
||||
-p 4000:4000 \
|
||||
ghcr.io/berriai/litellm:v1.78.0.rc.1
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
|
@ -44,7 +48,7 @@ import TabItem from '@theme/TabItem';
|
|||
<TabItem value="pip" label="Pip">
|
||||
|
||||
``` showLineNumbers title="pip install litellm"
|
||||
pip install litellm==1.78.0
|
||||
pip install litellm==1.78.0.rc.1
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
|
@ -54,12 +58,80 @@ pip install litellm==1.78.0
|
|||
|
||||
## Key Highlights
|
||||
|
||||
- **MCP Gateway Enhancements** - Fine-grained tool control at team/key level, OpenAPI to MCP server conversion, and per-tool parameter allowlists
|
||||
- **MCP Gateway - Control Tool Access by Team, Key** - Control MCP tool access by team/key.
|
||||
- **Performance Improvements** - 70% Lower p99 Latency
|
||||
- **GPT-5 Pro & GPT-Image-1-Mini** - Day 0 support for OpenAI's GPT-5 Pro (400K context) and gpt-image-1-mini image generation
|
||||
- **UI Performance Boost** - Replaces bloated key list calls with lean key aliases endpoint, Turbopack for faster development, and major UI refactors
|
||||
- **EnkryptAI Guardrails** - New guardrail integration for content moderation
|
||||
- **Tag-Based Budgets** - Support for setting budgets based on request tags
|
||||
- **Azure AD & SSO** - Enhanced Azure AD default credentials selection and EntraID app roles support
|
||||
|
||||
---
|
||||
|
||||
### MCP Gateway - Control Tool Access by Team, Key
|
||||
|
||||
<Image
|
||||
img={require('../../img/release_notes/tool_control.png')}
|
||||
style={{width: '100%', display: 'block', margin: '2rem auto'}}
|
||||
/>
|
||||
|
||||
<br/>
|
||||
|
||||
Proxy admins can now control MCP tool access by team or key. This makes it easy to grant different teams selective access to tools from the same MCP server.
|
||||
|
||||
For example, you can now give your Engineering team access to `list_repositories`, `create_issue`, and `search_code` tools, while Sales only gets `search_code` and `close_issue` tools.
|
||||
|
||||
This makes it easier for Proxy Admins to govern MCP Tool Access.
|
||||
|
||||
[Get Started](../../docs/mcp_control#set-allowed-tools-for-a-key-team-or-organization)
|
||||
|
||||
---
|
||||
|
||||
## Performance - 70% Lower p99 Latency
|
||||
|
||||
<Image img={require('../../img/release_notes/1_78_0_perf.png')} style={{ width: '800px', height: 'auto' }} />
|
||||
|
||||
<br/>
|
||||
|
||||
This release cuts p99 latency by 70% on LiteLLM AI Gateway, making it even better for low-latency use cases.
|
||||
|
||||
These gains come from two key enhancements:
|
||||
|
||||
**Reliable Sessions**
|
||||
|
||||
Added support for shared sessions with aiohttp. The shared_session parameter is now consistently used across all calls, enabling connection pooling.
|
||||
|
||||
**Faster Routing**
|
||||
|
||||
A new `model_name_to_deployment_indices` hash map replaces O(n) list scans in `_get_all_deployments()` with O(1) hash lookups, boosting routing performance and scalability.
|
||||
|
||||
As a result, performance improved across all latency percentiles:
|
||||
|
||||
- **Median latency:** 110 ms → **100 ms** (−9.1%)
|
||||
- **p95 latency:** 440 ms → **150 ms** (−65.9%)
|
||||
- **p99 latency:** 810 ms → **240 ms** (−70.4%)
|
||||
- **Average latency:** 310 ms → **111.73 ms** (−64.0%)
|
||||
|
||||
### **Test Setup**
|
||||
|
||||
**Locust**
|
||||
|
||||
- **Concurrent users:** 1,000
|
||||
- **Ramp-up:** 500
|
||||
|
||||
**System Specs**
|
||||
|
||||
- **Database was used**
|
||||
- **CPU:** 4 vCPUs
|
||||
- **Memory:** 8 GB RAM
|
||||
- **LiteLLM Workers:** 4
|
||||
- **Instances**: 4
|
||||
|
||||
**Configuration (config.yaml)**
|
||||
|
||||
View the complete configuration: [gist.github.com/AlexsanderHamir/config.yaml](https://gist.github.com/AlexsanderHamir/53f7d554a5d2afcf2c4edb5b6be68ff4)
|
||||
|
||||
**Load Script (no_cache_hits.py)**
|
||||
|
||||
View the complete load testing script: [gist.github.com/AlexsanderHamir/no_cache_hits.py](https://gist.github.com/AlexsanderHamir/42c33d7a4dc7a57f56a78b560dee3a42)
|
||||
|
||||
---
|
||||
|
||||
|
|
|
|||
|
|
@ -94,10 +94,10 @@ const sidebars = {
|
|||
|
||||
{
|
||||
type: "category",
|
||||
label: "LiteLLM Proxy Server",
|
||||
label: "LiteLLM AI Gateway",
|
||||
link: {
|
||||
type: "generated-index",
|
||||
title: "LiteLLM Proxy Server (LLM Gateway)",
|
||||
title: "LiteLLM AI Gateway (LLM Proxy)",
|
||||
description: `OpenAI Proxy Server (LLM Gateway) to call 100+ LLMs in a unified interface & track spend, set budgets per virtual key/user`,
|
||||
slug: "/simple_proxy",
|
||||
},
|
||||
|
|
|
|||
|
|
@ -36,6 +36,8 @@ async def apply_guardrail(
|
|||
if active_guardrail is None:
|
||||
raise Exception(f"Guardrail {request.guardrail_name} not found")
|
||||
|
||||
return await active_guardrail.apply_guardrail(
|
||||
response_text = await active_guardrail.apply_guardrail(
|
||||
text=request.text, language=request.language, entities=request.entities
|
||||
)
|
||||
|
||||
return ApplyGuardrailResponse(response_text=response_text)
|
||||
|
|
|
|||
|
|
@ -4841,7 +4841,7 @@
|
|||
"input_cost_per_token_above_200k_tokens": 6e-06,
|
||||
"litellm_provider": "anthropic",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 1000000,
|
||||
"max_output_tokens": 64000,
|
||||
"max_tokens": 1000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-05,
|
||||
|
|
@ -20200,12 +20200,12 @@
|
|||
},
|
||||
"us.anthropic.claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
"cache_read_input_token_cost": 33e-07,
|
||||
"input_cost_per_token": 33e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 66e-06,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
"input_cost_per_token": 3.3e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 6.6e-06,
|
||||
"output_cost_per_token_above_200k_tokens": 2.475e-05,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 8.25e-06,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 66e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 6.6e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 64000,
|
||||
|
|
|
|||
|
|
@ -1344,7 +1344,7 @@ class MCPServerManager:
|
|||
get_prisma_client_or_throw,
|
||||
)
|
||||
|
||||
verbose_logger.info("Loading MCP servers from database into registry...")
|
||||
verbose_logger.debug("Loading MCP servers from database into registry...")
|
||||
|
||||
# perform authz check to filter the mcp servers user has access to
|
||||
prisma_client = get_prisma_client_or_throw(
|
||||
|
|
@ -1360,7 +1360,9 @@ class MCPServerManager:
|
|||
)
|
||||
self.add_update_server(server)
|
||||
|
||||
verbose_logger.info(f"Registry now contains {len(self.get_registry())} servers")
|
||||
verbose_logger.debug(
|
||||
f"Registry now contains {len(self.get_registry())} servers"
|
||||
)
|
||||
|
||||
def get_mcp_server_by_id(self, server_id: str) -> Optional[MCPServer]:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1,4 +1,14 @@
|
|||
model_list:
|
||||
- model_name: gpt-5-codex
|
||||
- model_name: gpt-5-mini
|
||||
litellm_params:
|
||||
model: gpt-5-codex
|
||||
model: gpt-5-mini
|
||||
|
||||
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"]
|
||||
|
|
|
|||
|
|
@ -4841,7 +4841,7 @@
|
|||
"input_cost_per_token_above_200k_tokens": 6e-06,
|
||||
"litellm_provider": "anthropic",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 1000000,
|
||||
"max_output_tokens": 64000,
|
||||
"max_tokens": 1000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-05,
|
||||
|
|
@ -20200,12 +20200,12 @@
|
|||
},
|
||||
"us.anthropic.claude-sonnet-4-5-20250929-v1:0": {
|
||||
"cache_creation_input_token_cost": 4.125e-06,
|
||||
"cache_read_input_token_cost": 33e-07,
|
||||
"input_cost_per_token": 33e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 66e-06,
|
||||
"cache_read_input_token_cost": 3.3e-07,
|
||||
"input_cost_per_token": 3.3e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 6.6e-06,
|
||||
"output_cost_per_token_above_200k_tokens": 2.475e-05,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 8.25e-06,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 66e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 6.6e-07,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 64000,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,156 @@
|
|||
"""
|
||||
Test the /guardrails/apply_guardrail endpoint
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
from fastapi import HTTPException
|
||||
from litellm.types.guardrails import ApplyGuardrailRequest, ApplyGuardrailResponse
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_endpoint_returns_correct_response():
|
||||
"""Test that apply_guardrail endpoint returns ApplyGuardrailResponse object"""
|
||||
from enterprise.litellm_enterprise.proxy.guardrails.endpoints import apply_guardrail
|
||||
|
||||
# Mock the guardrail registry
|
||||
with patch("enterprise.litellm_enterprise.proxy.guardrails.endpoints.GUARDRAIL_REGISTRY") as mock_registry:
|
||||
# Create a mock guardrail
|
||||
mock_guardrail = Mock(spec=CustomGuardrail)
|
||||
mock_guardrail.apply_guardrail = AsyncMock(return_value="Redacted text: [REDACTED] and [REDACTED]")
|
||||
|
||||
# Configure the registry to return our mock guardrail
|
||||
mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
|
||||
|
||||
# Create the request
|
||||
request = ApplyGuardrailRequest(
|
||||
guardrail_name="test-guardrail",
|
||||
text="Test text with PII",
|
||||
language="en",
|
||||
entities=["EMAIL_ADDRESS", "PERSON"]
|
||||
)
|
||||
|
||||
# Create a mock user API key
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
||||
|
||||
# Call the endpoint
|
||||
response = await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
# Verify the response is of the correct type
|
||||
assert isinstance(response, ApplyGuardrailResponse)
|
||||
assert response.response_text == "Redacted text: [REDACTED] and [REDACTED]"
|
||||
|
||||
# Verify the guardrail was called with correct parameters
|
||||
mock_guardrail.apply_guardrail.assert_called_once_with(
|
||||
text="Test text with PII",
|
||||
language="en",
|
||||
entities=["EMAIL_ADDRESS", "PERSON"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_endpoint_guardrail_not_found():
|
||||
"""Test that apply_guardrail endpoint raises exception when guardrail not found"""
|
||||
from enterprise.litellm_enterprise.proxy.guardrails.endpoints import apply_guardrail
|
||||
|
||||
# Mock the guardrail registry to return None
|
||||
with patch("enterprise.litellm_enterprise.proxy.guardrails.endpoints.GUARDRAIL_REGISTRY") as mock_registry:
|
||||
mock_registry.get_initialized_guardrail_callback.return_value = None
|
||||
|
||||
# Create the request
|
||||
request = ApplyGuardrailRequest(
|
||||
guardrail_name="non-existent-guardrail",
|
||||
text="Test text",
|
||||
language="en"
|
||||
)
|
||||
|
||||
# Create a mock user API key
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
||||
|
||||
# Verify exception is raised
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
assert "Guardrail non-existent-guardrail not found" in str(exc_info.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_endpoint_with_presidio_guardrail():
|
||||
"""Test apply_guardrail endpoint with a Presidio-like guardrail"""
|
||||
from enterprise.litellm_enterprise.proxy.guardrails.endpoints import apply_guardrail
|
||||
|
||||
# Mock the guardrail registry
|
||||
with patch("enterprise.litellm_enterprise.proxy.guardrails.endpoints.GUARDRAIL_REGISTRY") as mock_registry:
|
||||
# Create a mock guardrail that simulates Presidio behavior
|
||||
mock_guardrail = Mock(spec=CustomGuardrail)
|
||||
# Simulate masking PII entities
|
||||
mock_guardrail.apply_guardrail = AsyncMock(
|
||||
return_value="My name is [PERSON] and my email is [EMAIL_ADDRESS]"
|
||||
)
|
||||
|
||||
# Configure the registry to return our mock guardrail
|
||||
mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
|
||||
|
||||
# Create the request
|
||||
request = ApplyGuardrailRequest(
|
||||
guardrail_name="pii-detection-guard",
|
||||
text="My name is John Doe and my email is john@example.com",
|
||||
language="en",
|
||||
entities=["EMAIL_ADDRESS", "PERSON"]
|
||||
)
|
||||
|
||||
# Create a mock user API key
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
||||
|
||||
# Call the endpoint
|
||||
response = await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
# Verify the response is of the correct type
|
||||
assert isinstance(response, ApplyGuardrailResponse)
|
||||
assert response.response_text == "My name is [PERSON] and my email is [EMAIL_ADDRESS]"
|
||||
assert "john@example.com" not in response.response_text
|
||||
assert "John Doe" not in response.response_text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_endpoint_without_optional_params():
|
||||
"""Test apply_guardrail endpoint without optional language and entities parameters"""
|
||||
from enterprise.litellm_enterprise.proxy.guardrails.endpoints import apply_guardrail
|
||||
|
||||
# Mock the guardrail registry
|
||||
with patch("enterprise.litellm_enterprise.proxy.guardrails.endpoints.GUARDRAIL_REGISTRY") as mock_registry:
|
||||
# Create a mock guardrail
|
||||
mock_guardrail = Mock(spec=CustomGuardrail)
|
||||
mock_guardrail.apply_guardrail = AsyncMock(return_value="Processed text")
|
||||
|
||||
# Configure the registry to return our mock guardrail
|
||||
mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail
|
||||
|
||||
# Create the request without optional parameters
|
||||
request = ApplyGuardrailRequest(
|
||||
guardrail_name="test-guardrail",
|
||||
text="Test text"
|
||||
)
|
||||
|
||||
# Create a mock user API key
|
||||
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
|
||||
|
||||
# Call the endpoint
|
||||
response = await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
# Verify the response is of the correct type
|
||||
assert isinstance(response, ApplyGuardrailResponse)
|
||||
assert response.response_text == "Processed text"
|
||||
|
||||
# Verify the guardrail was called with None for optional parameters
|
||||
mock_guardrail.apply_guardrail.assert_called_once_with(
|
||||
text="Test text",
|
||||
language=None,
|
||||
entities=None
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue