scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

Merge remote-tracking branch 'origin' into litellm_deleted_keys_team

Browse source
This commit is contained in:
yuneng-jiang 2026-01-16 14:16:05 -08:00
commit 5db9e78296
58 changed files with 1150 additions and 3609 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
.circleci/config.yml
View file

@ -144,8 +144,8 @@ jobs:
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install pyarrow
pip install "boto3==1.40.61"
pip install "aioboto3==15.5.0"
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install lunary==0.2.5
pip install "azure-identity==1.16.1"
@ -260,8 +260,8 @@ jobs:
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install pyarrow
pip install "boto3==1.40.61"
pip install "aioboto3==15.5.0"
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install lunary==0.2.5
pip install "azure-identity==1.16.1"
@ -367,8 +367,8 @@ jobs:
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install pyarrow
pip install "boto3==1.40.61"
pip install "aioboto3==15.5.0"
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install lunary==0.2.5
pip install "azure-identity==1.16.1"
@ -637,8 +637,8 @@ jobs:
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install pyarrow
pip install "boto3==1.40.61"
pip install "aioboto3==15.5.0"
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install "langfuse>=2.0.0"
pip install "logfire==0.29.0"
@ -759,8 +759,8 @@ jobs:
pip install "google-cloud-aiplatform==1.43.0"
pip install "google-genai==1.22.0"
pip install pyarrow
pip install "boto3==1.40.61"
pip install "aioboto3==15.5.0"
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install lunary==0.2.5
pip install "azure-identity==1.16.1"
@ -865,8 +865,8 @@ jobs:
pip install "google-cloud-aiplatform==1.43.0"
pip install "google-genai==1.22.0"
pip install pyarrow
pip install "boto3==1.40.61"
pip install "aioboto3==15.5.0"
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install lunary==0.2.5
pip install "azure-identity==1.16.1"
@ -972,8 +972,8 @@ jobs:
pip install "google-cloud-aiplatform==1.43.0"
pip install "google-genai==1.22.0"
pip install pyarrow
pip install "boto3==1.40.61"
pip install "aioboto3==15.5.0"
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install lunary==0.2.5
pip install "azure-identity==1.16.1"
@ -1198,7 +1198,7 @@ jobs:
pip install "pytest-asyncio==0.21.1"
pip install "respx==0.22.0"
pip install "pydantic==2.10.2"
pip install "boto3==1.40.61"
pip install "boto3==1.36.0"
# Run pytest and generate JUnit XML report
- run:
name: Run tests
@ -1879,7 +1879,7 @@ jobs:
pip install aiohttp
pip install openai
pip install click
pip install "boto3==1.40.61"
pip install "boto3==1.36.0"
pip install jinja2
pip install "tokenizers==0.20.0"
pip install "uvloop==0.21.0"
@ -2176,8 +2176,8 @@ jobs:
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install pyarrow
pip install "boto3==1.40.61"
pip install "aioboto3==15.5.0"
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install "langfuse>=2.0.0"
pip install "logfire==0.29.0"
@ -2316,8 +2316,8 @@ jobs:
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install pyarrow
pip install "boto3==1.40.61"
pip install "aioboto3==15.5.0"
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install "langchain_mcp_adapters==0.0.5"
pip install "langfuse>=2.0.0"
@ -2462,8 +2462,8 @@ jobs:
pip install "google-generativeai==0.3.2"
pip install "google-cloud-aiplatform==1.43.0"
pip install pyarrow
pip install "boto3==1.40.61"
pip install "aioboto3==15.5.0"
pip install "boto3==1.36.0"
pip install "aioboto3==13.4.0"
pip install langchain
pip install "langfuse>=2.0.0"
pip install "logfire==0.29.0"
@ -3118,7 +3118,7 @@ jobs:
pip install "pytest==7.3.1"
pip install "pytest-mock==3.12.0"
pip install "pytest-asyncio==0.21.1"
pip install "boto3==1.40.61"
pip install "boto3==1.36.0"
pip install "mypy==1.18.2"
pip install pyarrow
pip install numpydoc

8
.github/ISSUE_TEMPLATE/bug_report.yml vendored
View file

@ -9,6 +9,14 @@ body:
Thanks for taking the time to fill out this bug report!
**💡 Tip:** See our [Troubleshooting Guide](https://docs.litellm.ai/docs/troubleshoot) for what information to include.
- type: checkboxes
id: duplicate-check
attributes:
label: Check for existing issues
description: Please search to see if an issue already exists for the bug you encountered.
options:
- label: I have searched the existing issues and checked that my issue is not a duplicate.
required: true
- type: textarea
id: what-happened
attributes:

8
.github/ISSUE_TEMPLATE/feature_request.yml vendored
View file

@ -7,6 +7,14 @@ body:
attributes:
value: |
Thanks for making LiteLLM better!
- type: checkboxes
id: duplicate-check
attributes:
label: Check for existing issues
description: Please search to see if an issue already exists for the feature you are requesting.
options:
- label: I have searched the existing issues and checked that my issue is not a duplicate.
required: true
- type: textarea
id: the-feature
attributes:

29
.github/workflows/check_duplicate_issues.yml vendored Normal file
View file

@ -0,0 +1,29 @@
name: Check Duplicate Issues
on:
issues:
types: [opened, edited]
jobs:
check-duplicate:
runs-on: ubuntu-latest
permissions:
issues: write
contents: read
steps:
- name: Check for potential duplicates
uses: wow-actions/potential-duplicates@v1
with:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
label: potential-duplicate
threshold: 0.6
reaction: eyes
comment: |
**⚠️ Potential duplicate detected**
This issue appears similar to existing issue(s):
{{#issues}}
- [#{{number}}]({{html_url}}) - {{title}} ({{accuracy}}% similar)
{{/issues}}
Please review the linked issue(s) to see if they address your concern. If this is not a duplicate, please provide additional context to help us understand the difference.

34
.github/workflows/label-component.yml vendored
View file

@ -80,3 +80,37 @@ jobs:
break;
}
}
// Check for 'claude code' keyword (can be applied alongside component labels)
if (/claude code/i.test(body)) {
const claudeLabel = {
name: 'claude code',
color: '7c3aed',
description: 'Issues related to Claude Code usage'
};
try {
await github.rest.issues.getLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: claudeLabel.name
});
} catch (error) {
if (error.status === 404) {
await github.rest.issues.createLabel({
owner: context.repo.owner,
repo: context.repo.repo,
name: claudeLabel.name,
color: claudeLabel.color,
description: claudeLabel.description
});
}
}
await github.rest.issues.addLabels({
owner: context.repo.owner,
repo: context.repo.repo,
issue_number: context.issue.number,
labels: [claudeLabel.name]
});
}

468
docs/my-website/docs/completion/message_sanitization.md
View file

@ -1,468 +0,0 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Message Sanitization for Tool Calling for anthropic models
**Automatically fix common message formatting issues when using tool calling with `modify_params=True`**
LiteLLM can automatically sanitize messages to handle common issues that occur during tool calling workflows, especially when using OpenAI-compatible clients with providers that have strict message format requirements (like Anthropic Claude).
## Overview
When `litellm.modify_params = True` is enabled, LiteLLM automatically sanitizes messages to fix three common issues:
1. **Orphaned Tool Calls** - Assistant messages with tool_calls but missing tool results
2. **Orphaned Tool Results** - Tool messages that reference non-existent tool_call_ids
3. **Empty Message Content** - Messages with empty or whitespace-only text content
This ensures your tool calling workflows work seamlessly across different LLM providers without manual message validation.
## Why Message Sanitization?
Different LLM providers have varying requirements for message formats, especially during tool calling:
- **Anthropic Claude** requires every tool_call to have a corresponding tool result
- Some providers reject messages with empty content
- OpenAI-compatible clients may not always maintain perfect message consistency
Without sanitization, these issues cause API errors that interrupt your workflows. With `modify_params=True`, LiteLLM handles these edge cases automatically.
## Quick Start
<Tabs>
<TabItem value="sdk" label="SDK">
```python
import litellm
# Enable automatic message sanitization
litellm.modify_params = True
# This will work even if messages have formatting issues
response = litellm.completion(
model="anthropic/claude-3-5-sonnet-20241022",
messages=[
{"role": "user", "content": "What's the weather in Boston?"},
{
"role": "assistant",
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {"name": "get_weather", "arguments": '{"city": "Boston"}'}
}
]
# Missing tool result - LiteLLM will add a dummy result automatically
},
{"role": "user", "content": "Thanks!"}
],
tools=[{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather for a city",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"]
}
}
}]
)
```
</TabItem>
<TabItem value="proxy" label="PROXY">
```yaml
litellm_settings:
modify_params: true # Enable automatic message sanitization
model_list:
- model_name: claude-3-5-sonnet
litellm_params:
model: anthropic/claude-3-5-sonnet-20241022
```
</TabItem>
</Tabs>
## Sanitization Cases
### Case A: Orphaned Tool Calls (Missing Tool Results)
**Problem:** An assistant message contains `tool_calls`, but no corresponding tool result messages follow.
**Solution:** LiteLLM automatically adds dummy tool result messages for any missing tool results.
**Example:**
```python
import litellm
litellm.modify_params = True
# Messages with orphaned tool calls
messages = [
{"role": "user", "content": "Search for Python tutorials"},
{
"role": "assistant",
"tool_calls": [
{
"id": "call_abc123",
"type": "function",
"function": {"name": "web_search", "arguments": '{"query": "Python tutorials"}'}
}
]
},
# Missing tool result here!
{"role": "user", "content": "What about JavaScript?"}
]
# LiteLLM automatically adds:
# {
# "role": "tool",
# "tool_call_id": "call_abc123",
# "content": "[System: Tool execution skipped/interrupted by user. No result provided for tool 'web_search'.]"
# }
response = litellm.completion(
model="anthropic/claude-3-5-sonnet-20241022",
messages=messages,
tools=[...]
)
```
**When this happens:**
- User interrupts tool execution
- Client loses tool results due to network issues
- Conversation flow changes before tool completes
- Multi-turn conversations where tools are optional
### Case B: Orphaned Tool Results (Invalid tool_call_id)
**Problem:** A tool message references a `tool_call_id` that doesn't exist in any previous assistant message.
**Solution:** LiteLLM automatically removes these orphaned tool result messages.
**Example:**
```python
import litellm
litellm.modify_params = True
# Messages with orphaned tool result
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi! How can I help?"},
{
"role": "tool",
"tool_call_id": "call_nonexistent", # This tool_call_id doesn't exist!
"content": "Some result"
}
]
# LiteLLM automatically removes the orphaned tool message
response = litellm.completion(
model="anthropic/claude-3-5-sonnet-20241022",
messages=messages
)
```
**When this happens:**
- Message history is manually edited
- Tool results are duplicated or mismatched
- Conversation state is restored incorrectly
- Messages are merged from different conversations
### Case C: Empty Message Content
**Problem:** User or assistant messages have empty or whitespace-only content.
**Solution:** LiteLLM replaces empty content with a system placeholder message.
**Example:**
```python
import litellm
litellm.modify_params = True
# Messages with empty content
messages = [
{"role": "user", "content": ""}, # Empty content
{"role": "assistant", "content": " "}, # Whitespace only
]
# LiteLLM automatically replaces with:
# {"role": "user", "content": "[System: Empty message content sanitised to satisfy protocol]"}
# {"role": "assistant", "content": "[System: Empty message content sanitised to satisfy protocol]"}
response = litellm.completion(
model="anthropic/claude-3-5-sonnet-20241022",
messages=messages
)
```
**When this happens:**
- UI sends empty messages
- Content is stripped during preprocessing
- Placeholder messages in conversation history
- Edge cases in message construction
## Configuration
### Enable Globally
<Tabs>
<TabItem value="sdk" label="SDK">
```python
import litellm
# Enable for all completion calls
litellm.modify_params = True
```
</TabItem>
<TabItem value="proxy" label="PROXY">
```yaml
litellm_settings:
modify_params: true
```
</TabItem>
<TabItem value="env" label="Environment Variable">
```bash
export LITELLM_MODIFY_PARAMS=True
```
</TabItem>
</Tabs>
### Enable Per-Request
```python
import litellm
# Enable only for specific requests
response = litellm.completion(
model="anthropic/claude-3-5-sonnet-20241022",
messages=messages,
modify_params=True # Override global setting
)
```
## Supported Providers
Message sanitization works with all LLM providers that support tool calling:
- ✅ Anthropic (Claude)
- ✅ OpenAI (GPT-4, GPT-3.5)
- ✅ AWS Bedrock (Claude, Titan)
- ✅ Google Vertex AI (Claude, Gemini)
- ✅ Azure OpenAI
- ✅ And all other providers with tool calling support
## Implementation Details
### How It Works
The message sanitization process runs **before** messages are converted to provider-specific formats:
1. **Input:** OpenAI-format messages with potential issues
2. **Sanitization:** Three helper functions process the messages:
- `_sanitize_empty_text_content()` - Fixes empty content
- `_add_missing_tool_results()` - Adds dummy tool results
- `_is_orphaned_tool_result()` - Identifies orphaned results
3. **Output:** Clean, provider-compatible messages
### Code Reference
The sanitization logic is implemented in:
- `litellm/litellm_core_utils/prompt_templates/factory.py`
- Function: `sanitize_messages_for_tool_calling()`
### Logging
When sanitization occurs, LiteLLM logs debug messages:
```python
import litellm
litellm.set_verbose = True # Enable debug logging
# You'll see logs like:
# "_add_missing_tool_results: Found 1 orphaned tool calls. Adding dummy tool results."
# "_is_orphaned_tool_result: Found orphaned tool result with tool_call_id=call_123"
# "_sanitize_empty_text_content: Replaced empty text content in user message"
```
## Best Practices
### 1. Enable for Production Workflows
```python
# Recommended for production
litellm.modify_params = True
# Ensures robust handling of edge cases
response = litellm.completion(
model="anthropic/claude-3-5-sonnet-20241022",
messages=messages,
tools=tools
)
```
### 2. Preserve Tool Results When Possible
While sanitization handles missing tool results, it's better to provide actual results:
```python
# Good: Provide actual tool results
messages = [
{"role": "user", "content": "Search for Python"},
{"role": "assistant", "tool_calls": [...]},
{"role": "tool", "tool_call_id": "call_123", "content": "Actual search results"}
]
# Fallback: Sanitization adds dummy result if missing
messages = [
{"role": "user", "content": "Search for Python"},
{"role": "assistant", "tool_calls": [...]},
# Missing tool result - sanitization adds dummy
]
```
### 3. Monitor Sanitization Events
Use logging to track when sanitization occurs:
```python
import litellm
import logging
# Enable debug logging
litellm.set_verbose = True
logging.basicConfig(level=logging.DEBUG)
# Track sanitization events in your application
response = litellm.completion(
model="anthropic/claude-3-5-sonnet-20241022",
messages=messages
)
```
### 4. Test Edge Cases
Ensure your application handles sanitized messages correctly:
```python
import litellm
litellm.modify_params = True
# Test orphaned tool calls
test_messages = [
{"role": "user", "content": "Test"},
{"role": "assistant", "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "test", "arguments": "{}"}}]},
{"role": "user", "content": "Continue"} # No tool result
]
response = litellm.completion(
model="anthropic/claude-3-5-sonnet-20241022",
messages=test_messages,
tools=[...]
)
# Verify the response handles the dummy tool result appropriately
```
## Related Features
- **[Drop Params](./drop_params.md)** - Drop unsupported parameters for specific providers
- **[Message Trimming](./message_trimming.md)** - Trim messages to fit token limits
- **[Function Calling](./function_call.md)** - Complete guide to tool/function calling
- **[Reasoning Content](../reasoning_content.md)** - Extended thinking with tool calling
## Troubleshooting
### Sanitization Not Working
**Issue:** Messages still cause errors despite `modify_params=True`
**Solution:**
1. Verify `modify_params` is enabled:
```python
import litellm
print(litellm.modify_params) # Should be True
```
2. Check if the issue is provider-specific:
```python
litellm.set_verbose = True # Enable debug logging
```
3. Ensure you're using a recent version of LiteLLM:
```bash
pip install --upgrade litellm
```
### Unexpected Dummy Tool Results
**Issue:** Dummy tool results appear when you expect actual results
**Cause:** Tool result messages are missing or have incorrect `tool_call_id`
**Solution:**
1. Verify tool result messages have correct `tool_call_id`:
```python
# Correct
{"role": "tool", "tool_call_id": "call_123", "content": "result"}
# Incorrect - will be treated as orphaned
{"role": "tool", "tool_call_id": "wrong_id", "content": "result"}
```
2. Ensure tool results immediately follow assistant messages with tool_calls
### Performance Impact
**Issue:** Concerned about performance overhead
**Details:** Message sanitization has minimal performance impact:
- Runs in O(n) time where n = number of messages
- Only processes messages when `modify_params=True`
- Typically adds < 1ms to request processing time
## FAQ
**Q: Does sanitization modify my original messages?**
A: No, sanitization creates a new list of messages. Your original messages remain unchanged.
**Q: Can I disable specific sanitization cases?**
A: Currently, all three cases are handled together when `modify_params=True`. To disable sanitization entirely, set `modify_params=False`.
**Q: What happens to the dummy tool results?**
A: Dummy tool results are sent to the LLM provider along with other messages. The model sees them as regular tool results with informative error messages.
**Q: Does this work with streaming?**
A: Yes, message sanitization works with both streaming and non-streaming requests.
**Q: Is this related to `drop_params`?**
A: No, they're separate features:
- `modify_params` - Modifies/fixes message content and structure
- `drop_params` - Removes unsupported API parameters
Both can be enabled simultaneously.
## See Also
- [Reasoning Content with Tool Calling](../reasoning_content.md)
- [Function Calling Guide](./function_call.md)
- [Bedrock Provider Documentation](../providers/bedrock.md)
- [Anthropic Provider Documentation](../providers/anthropic.md)

1
docs/my-website/sidebars.js
View file

@ -822,7 +822,6 @@ const sidebars = {
"completion/knowledgebase",
"guides/code_interpreter",
"completion/message_trimming",
"completion/message_sanitization",
"completion/model_alias",
"completion/mock_requests",
"completion/predict_outputs",

BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.15-py3-none-any.whl vendored
View file

Binary file not shown.

BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.15.tar.gz vendored
View file

Binary file not shown.

BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.22-py3-none-any.whl vendored
View file

Binary file not shown.

BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.22.tar.gz vendored
View file

Binary file not shown.

117
litellm-proxy-extras/litellm_proxy_extras/migrations/20251219110931_add_deleted_keys_and_deleted_teams_tables/migration.sql
View file

@ -1,117 +0,0 @@
-- CreateTable
CREATE TABLE "LiteLLM_DeletedTeamTable" (
"id" TEXT NOT NULL,
"team_id" TEXT NOT NULL,
"team_alias" TEXT,
"organization_id" TEXT,
"object_permission_id" TEXT,
"admins" TEXT[],
"members" TEXT[],
"members_with_roles" JSONB NOT NULL DEFAULT '{}',
"metadata" JSONB NOT NULL DEFAULT '{}',
"max_budget" DOUBLE PRECISION,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"models" TEXT[],
"max_parallel_requests" INTEGER,
"tpm_limit" BIGINT,
"rpm_limit" BIGINT,
"budget_duration" TEXT,
"budget_reset_at" TIMESTAMP(3),
"blocked" BOOLEAN NOT NULL DEFAULT false,
"model_spend" JSONB NOT NULL DEFAULT '{}',
"model_max_budget" JSONB NOT NULL DEFAULT '{}',
"team_member_permissions" TEXT[] DEFAULT ARRAY[]::TEXT[],
"model_id" INTEGER,
"created_at" TIMESTAMP(3),
"updated_at" TIMESTAMP(3),
"deleted_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"deleted_by" TEXT,
"deleted_by_api_key" TEXT,
"litellm_changed_by" TEXT,
CONSTRAINT "LiteLLM_DeletedTeamTable_pkey" PRIMARY KEY ("id")
);
-- CreateTable
CREATE TABLE "LiteLLM_DeletedVerificationToken" (
"id" TEXT NOT NULL,
"token" TEXT NOT NULL,
"key_name" TEXT,
"key_alias" TEXT,
"soft_budget_cooldown" BOOLEAN NOT NULL DEFAULT false,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"expires" TIMESTAMP(3),
"models" TEXT[],
"aliases" JSONB NOT NULL DEFAULT '{}',
"config" JSONB NOT NULL DEFAULT '{}',
"user_id" TEXT,
"team_id" TEXT,
"permissions" JSONB NOT NULL DEFAULT '{}',
"max_parallel_requests" INTEGER,
"metadata" JSONB NOT NULL DEFAULT '{}',
"blocked" BOOLEAN,
"tpm_limit" BIGINT,
"rpm_limit" BIGINT,
"max_budget" DOUBLE PRECISION,
"budget_duration" TEXT,
"budget_reset_at" TIMESTAMP(3),
"allowed_cache_controls" TEXT[] DEFAULT ARRAY[]::TEXT[],
"allowed_routes" TEXT[] DEFAULT ARRAY[]::TEXT[],
"model_spend" JSONB NOT NULL DEFAULT '{}',
"model_max_budget" JSONB NOT NULL DEFAULT '{}',
"budget_id" TEXT,
"organization_id" TEXT,
"object_permission_id" TEXT,
"created_at" TIMESTAMP(3),
"created_by" TEXT,
"updated_at" TIMESTAMP(3),
"updated_by" TEXT,
"rotation_count" INTEGER DEFAULT 0,
"auto_rotate" BOOLEAN DEFAULT false,
"rotation_interval" TEXT,
"last_rotation_at" TIMESTAMP(3),
"key_rotation_at" TIMESTAMP(3),
"deleted_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"deleted_by" TEXT,
"deleted_by_api_key" TEXT,
"litellm_changed_by" TEXT,
CONSTRAINT "LiteLLM_DeletedVerificationToken_pkey" PRIMARY KEY ("id")
);
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedTeamTable_team_id_idx" ON "LiteLLM_DeletedTeamTable"("team_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedTeamTable_deleted_at_idx" ON "LiteLLM_DeletedTeamTable"("deleted_at");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedTeamTable_organization_id_idx" ON "LiteLLM_DeletedTeamTable"("organization_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedTeamTable_team_alias_idx" ON "LiteLLM_DeletedTeamTable"("team_alias");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedTeamTable_created_at_idx" ON "LiteLLM_DeletedTeamTable"("created_at");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_token_idx" ON "LiteLLM_DeletedVerificationToken"("token");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_deleted_at_idx" ON "LiteLLM_DeletedVerificationToken"("deleted_at");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_user_id_idx" ON "LiteLLM_DeletedVerificationToken"("user_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_team_id_idx" ON "LiteLLM_DeletedVerificationToken"("team_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_organization_id_idx" ON "LiteLLM_DeletedVerificationToken"("organization_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_key_alias_idx" ON "LiteLLM_DeletedVerificationToken"("key_alias");
-- CreateIndex
CREATE INDEX "LiteLLM_DeletedVerificationToken_created_at_idx" ON "LiteLLM_DeletedVerificationToken"("created_at");

99
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -132,49 +132,6 @@ model LiteLLM_TeamTable {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted teams - preserves spend and team information for historical tracking
model LiteLLM_DeletedTeamTable {
id String @id @default(uuid())
team_id String // Original team_id
team_alias String?
organization_id String?
object_permission_id String?
admins String[]
members String[]
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
spend Float @default(0.0)
models String[]
max_parallel_requests Int?
tpm_limit BigInt?
rpm_limit BigInt?
budget_duration String?
budget_reset_at DateTime?
blocked Boolean @default(false)
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
// Original timestamps from team creation/updates
created_at DateTime? @map("created_at")
updated_at DateTime? @map("updated_at")
// Deletion metadata
deleted_at DateTime @default(now()) @map("deleted_at")
deleted_by String? @map("deleted_by") // User who deleted the team
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([team_id])
@@index([deleted_at])
@@index([organization_id])
@@index([team_alias])
@@index([created_at])
}
// Track spend, rate limit, budget Users
model LiteLLM_UserTable {
user_id String @id
@ -302,62 +259,6 @@ model LiteLLM_VerificationToken {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted keys - preserves spend and key information for historical tracking
model LiteLLM_DeletedVerificationToken {
id String @id @default(uuid())
token String // Original token (hashed)
key_name String?
key_alias String?
soft_budget_cooldown Boolean @default(false)
spend Float @default(0.0)
expires DateTime?
models String[]
aliases Json @default("{}")
config Json @default("{}")
user_id String?
team_id String?
permissions Json @default("{}")
max_parallel_requests Int?
metadata Json @default("{}")
blocked Boolean?
tpm_limit BigInt?
rpm_limit BigInt?
max_budget Float?
budget_duration String?
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
budget_id String?
organization_id String?
object_permission_id String?
created_at DateTime? // Original creation timestamp
created_by String? // Original creator
updated_at DateTime? // Last update timestamp before deletion
updated_by String? // Last user who updated before deletion
rotation_count Int? @default(0)
auto_rotate Boolean? @default(false)
rotation_interval String?
last_rotation_at DateTime?
key_rotation_at DateTime?
// Deletion metadata
deleted_at DateTime @default(now()) @map("deleted_at")
deleted_by String? @map("deleted_by") // User who deleted the key
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([token])
@@index([deleted_at])
@@index([user_id])
@@index([team_id])
@@index([organization_id])
@@index([key_alias])
@@index([created_at])
}
model LiteLLM_EndUserTable {
user_id String @id
alias String? // admin-facing alias

22
litellm/_logging.py
View file

@ -133,6 +133,26 @@ ALL_LOGGERS = [
]
def _get_loggers_to_initialize():
"""
Get all loggers that should be initialized with the JSON handler.
Includes third-party integration loggers (like langfuse) if they are
configured as callbacks.
"""
import litellm
loggers = list(ALL_LOGGERS)
# Add langfuse logger if langfuse is being used as a callback
langfuse_callbacks = {"langfuse", "langfuse_otel"}
all_callbacks = set(litellm.success_callback + litellm.failure_callback)
if langfuse_callbacks & all_callbacks:
loggers.append(logging.getLogger("langfuse"))
return loggers
def _initialize_loggers_with_handler(handler: logging.Handler):
"""
Initialize all loggers with a handler
@ -140,7 +160,7 @@ def _initialize_loggers_with_handler(handler: logging.Handler):
- Adds a handler to each logger
- Prevents bubbling to parent/root (critical to prevent duplicate JSON logs)
"""
for lg in ALL_LOGGERS:
for lg in _get_loggers_to_initialize():
lg.handlers.clear() # remove any existing handlers
lg.addHandler(handler) # add JSON formatter handler
lg.propagate = False # prevent bubbling to parent/root

240
litellm/litellm_core_utils/prompt_templates/factory.py
View file

@ -1989,223 +1989,6 @@ def anthropic_process_openai_file_message(
)
def _sanitize_empty_text_content(
message: AllMessageValues,
) -> AllMessageValues:
"""
Case C: Sanitize empty text content
- Replace empty or whitespace-only text content with a placeholder message.
Returns:
The message with sanitized content if needed, otherwise the original message
"""
if message.get("role") in ["user", "assistant"]:
content = message.get("content")
if isinstance(content, str):
if not content or not content.strip():
message = dict(message) # Make a copy
message["content"] = "[System: Empty message content sanitised to satisfy protocol]"
verbose_logger.debug(
f"_sanitize_empty_text_content: Replaced empty text content in {message.get('role')} message"
)
return message
def _add_missing_tool_results(
current_message: AllMessageValues,
messages: List[AllMessageValues],
current_index: int,
) -> List[AllMessageValues]:
"""
Case A: Missing tool_result for tool_use (orphaned tool calls)
- If an assistant message has tool_calls but no corresponding tool result follows,
add a dummy tool result message indicating the user did not provide the result.
Returns:
A list containing the assistant message followed by any dummy tool results needed
"""
result_messages: List[AllMessageValues] = []
tool_calls = current_message.get("tool_calls")
if not tool_calls or len(tool_calls) == 0:
return [current_message]
# Collect all tool_call_ids from this assistant message
expected_tool_call_ids = set()
for tool_call in tool_calls:
tool_call_id = None
if isinstance(tool_call, dict):
tool_call_id = tool_call.get("id")
else:
tool_call_id = getattr(tool_call, "id", None)
if tool_call_id:
expected_tool_call_ids.add(tool_call_id)
found_tool_call_ids = set()
j = current_index + 1
while j < len(messages):
next_msg = messages[j]
next_role = next_msg.get("role")
if next_role == "assistant":
break
if next_role in ["tool", "function"]:
tool_call_id = next_msg.get("tool_call_id")
if tool_call_id:
found_tool_call_ids.add(tool_call_id)
j += 1
# Find missing tool results
missing_tool_call_ids = expected_tool_call_ids - found_tool_call_ids
if missing_tool_call_ids:
verbose_logger.debug(
f"_add_missing_tool_results: Found {len(missing_tool_call_ids)} orphaned tool calls. Adding dummy tool results."
)
result_messages.append(current_message)
for tool_call_id in missing_tool_call_ids:
tool_name = "unknown_tool"
for tool_call in tool_calls:
tc_id = None
if isinstance(tool_call, dict):
tc_id = tool_call.get("id")
else:
tc_id = getattr(tool_call, "id", None)
if tc_id == tool_call_id:
if isinstance(tool_call, dict):
function = tool_call.get("function", {})
if isinstance(function, dict):
tool_name = function.get("name", "unknown_tool")
else:
tool_name = getattr(function, "name", "unknown_tool")
else:
function = getattr(tool_call, "function", None)
if function:
tool_name = getattr(function, "name", "unknown_tool")
break
dummy_tool_result: ChatCompletionToolMessage = {
"role": "tool",
"tool_call_id": tool_call_id,
"content": f"[System: Tool execution skipped/interrupted by user. No result provided for tool '{tool_name}'.]",
}
result_messages.append(dummy_tool_result)
return result_messages
return [current_message]
def _is_orphaned_tool_result(
current_message: AllMessageValues,
sanitized_messages: List[AllMessageValues],
) -> bool:
"""
Case B: Orphaned tool_result (unexpected result)
- Check if a tool message references a tool_call_id that doesn't exist in the previous
assistant message.
Returns:
True if this is an orphaned tool result that should be removed, False otherwise
"""
if current_message.get("role") not in ["tool", "function"]:
return False
tool_call_id = current_message.get("tool_call_id")
if not tool_call_id:
return False
# Look back to find the most recent assistant message with tool_calls
found_matching_tool_call = False
for j in range(len(sanitized_messages) - 1, -1, -1):
prev_msg = sanitized_messages[j]
if prev_msg.get("role") == "assistant":
tool_calls = prev_msg.get("tool_calls")
if tool_calls:
for tool_call in tool_calls:
tc_id = None
if isinstance(tool_call, dict):
tc_id = tool_call.get("id")
else:
tc_id = getattr(tool_call, "id", None)
if tc_id == tool_call_id:
found_matching_tool_call = True
break
break
if not found_matching_tool_call:
verbose_logger.debug(
"_is_orphaned_tool_result: Found orphaned tool result with redacted tool_call_id"
)
return True
return False
def sanitize_messages_for_tool_calling(
messages: List[AllMessageValues],
) -> List[AllMessageValues]:
"""
Sanitize messages for tool calling to handle common issues when modify_params=True:
Case A: Missing tool_result for tool_use (orphaned tool calls)
- If an assistant message has tool_calls but no corresponding tool result follows,
add a dummy tool result message indicating the user did not provide the result.
Case B: Orphaned tool_result (unexpected result)
- If a tool message references a tool_call_id that doesn't exist in the previous
assistant message, remove that tool message.
Case C: Empty text content
- Replace empty or whitespace-only text content with a placeholder message.
This function operates on OpenAI format messages before they are converted to
provider-specific formats.
"""
if not litellm.modify_params:
return messages
sanitized_messages: List[AllMessageValues] = []
i = 0
while i < len(messages):
current_message = messages[i]
# Case C: Sanitize empty text content
current_message = _sanitize_empty_text_content(current_message)
# Case A: Check if assistant message has tool_calls without following tool results
if current_message.get("role") == "assistant":
result_messages = _add_missing_tool_results(current_message, messages, i)
# If dummy tool results were added, extend sanitized_messages and continue
if len(result_messages) > 1:
sanitized_messages.extend(result_messages)
i += 1
continue
# Case B: Check for orphaned tool results
if _is_orphaned_tool_result(current_message, sanitized_messages):
i += 1
continue # Skip this orphaned tool result
# Add the message to sanitized list
sanitized_messages.append(current_message)
i += 1
return sanitized_messages
def anthropic_messages_pt( # noqa: PLR0915
messages: List[AllMessageValues],
model: str,
@ -2225,9 +2008,6 @@ def anthropic_messages_pt( # noqa: PLR0915
5. System messages are a separate param to the Messages API
6. Ensure we only accept role, content. (message.name is not supported)
"""
# Sanitize messages for tool calling issues when modify_params=True
messages = sanitize_messages_for_tool_calling(messages)
# add role=tool support to allow function call result/error submission
user_message_types = {"user", "tool", "function"}
# reformat messages to ensure user/assistant are alternating, if there's either 2 consecutive 'user' messages or 2 consecutive 'assistant' message, merge them.
@ -3453,21 +3233,17 @@ def _convert_to_bedrock_tool_call_invoke(
id = tool["id"]
name = tool["function"].get("name", "")
arguments = tool["function"].get("arguments", "")
arguments_dict = json.loads(arguments) if arguments else {}
# Ensure arguments_dict is always a dict (Bedrock requires toolUse.input to be an object)
# When some providers return arguments: '""' (JSON-encoded empty string), json.loads returns ""
if not isinstance(arguments_dict, dict):
arguments_dict = {}
if not arguments or not arguments.strip():
arguments_input = {}
arguments_dict = {}
else:
# Try to parse the arguments JSON
try:
arguments_input = json.loads(arguments)
except json.JSONDecodeError as e:
verbose_logger.warning(
f"Malformed JSON in tool call arguments for tool '{name}': {str(e)}. "
f"Storing as raw string to allow conversation to continue."
)
arguments_input = arguments
arguments_dict = json.loads(arguments)
bedrock_tool = BedrockToolUseBlock(
input=arguments_input, name=name, toolUseId=id
input=arguments_dict, name=name, toolUseId=id
)
bedrock_content_block = BedrockContentBlock(toolUse=bedrock_tool)
_parts_list.append(bedrock_content_block)

9
litellm/llms/bedrock/chat/converse_transformation.py
View file

@ -1395,16 +1395,9 @@ class AmazonConverseConfig(BaseConfig):
response_tool_name = get_bedrock_tool_name(
response_tool_name=_response_tool_name
)
tool_input = content["toolUse"]["input"]
if isinstance(tool_input, str):
arguments_str = tool_input
else:
# Otherwise, serialize it to JSON
arguments_str = json.dumps(tool_input)
_function_chunk = ChatCompletionToolCallFunctionChunk(
name=response_tool_name,
arguments=arguments_str,
arguments=json.dumps(content["toolUse"]["input"]),
)
_tool_response_chunk = ChatCompletionToolCallChunk(

4
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -304,7 +304,7 @@ class ContextCachingEndpoints(VertexBase):
## CHECK IF CACHED ALREADY
generated_cache_key = local_cache_obj.get_cache_key(
messages=cached_messages, tools=tools
messages=cached_messages, tools=tools, model=model
)
google_cache_name = self.check_cache(
cache_key=generated_cache_key,
@ -433,7 +433,7 @@ class ContextCachingEndpoints(VertexBase):
## CHECK IF CACHED ALREADY
generated_cache_key = local_cache_obj.get_cache_key(
messages=cached_messages, tools=tools
messages=cached_messages, tools=tools, model=model
)
google_cache_name = await self.async_check_cache(
cache_key=generated_cache_key,

30
litellm/proxy/_types.py
View file

@ -1723,21 +1723,6 @@ class LiteLLM_TeamTableCachedObj(LiteLLM_TeamTable):
last_refreshed_at: Optional[float] = None
class LiteLLM_DeletedTeamTable(LiteLLM_TeamTable):
"""
Recording of deleted teams for audit purposes. Mirrors LiteLLM_TeamTable
plus metadata captured at deletion time.
"""
id: Optional[str] = None
deleted_at: Optional[datetime] = None
deleted_by: Optional[str] = None
deleted_by_api_key: Optional[str] = None
litellm_changed_by: Optional[str] = None
model_config = ConfigDict(protected_namespaces=())
class TeamRequest(LiteLLMPydanticObjectBase):
teams: List[str]
@ -2132,21 +2117,6 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
model_config = ConfigDict(protected_namespaces=())
class LiteLLM_DeletedVerificationToken(LiteLLM_VerificationToken):
"""
Recording of deleted keys for audit purposes. Mirrors LiteLLM_VerificationToken
plus metadata captured at deletion time.
"""
id: Optional[str] = None
deleted_at: Optional[datetime] = None
deleted_by: Optional[str] = None
deleted_by_api_key: Optional[str] = None
litellm_changed_by: Optional[str] = None
model_config = ConfigDict(protected_namespaces=())
class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
"""
Combined view of litellm verification token + litellm team table (select values)

138
litellm/proxy/auth/auth_checks.py
View file

@ -202,21 +202,29 @@ async def common_checks(
and general_settings["enforce_user_param"] is True
):
# Get HTTP method from request
http_method = request.method if hasattr(request, 'method') else None
http_method = request.method if hasattr(request, "method") else None
# Check if it's a POST request and if it's an OpenAI route but not MCP
is_post_method = http_method and http_method.upper() == "POST"
is_openai_route = RouteChecks.is_llm_api_route(route=route)
is_mcp_route = route in LiteLLMRoutes.mcp_routes.value or RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
is_mcp_route = (
route in LiteLLMRoutes.mcp_routes.value
or RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
)
)
# Enforce user param only for POST requests on OpenAI routes (excluding MCP routes)
if is_post_method and is_openai_route and not is_mcp_route and "user" not in request_body:
if (
is_post_method
and is_openai_route
and not is_mcp_route
and "user" not in request_body
):
raise Exception(
f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}"
)
# 6.1 [OPTIONAL] If 'reject_clientside_metadata_tags' enabled - reject request if it has client-side 'metadata.tags'
if (
general_settings.get("reject_clientside_metadata_tags", None) is not None
@ -502,53 +510,51 @@ async def get_default_end_user_budget(
) -> Optional[LiteLLM_BudgetTable]:
"""
Fetches the default end user budget from the database if litellm.max_end_user_budget_id is configured.
This budget is applied to end users who don't have an explicit budget_id set.
Results are cached for performance.
Args:
prisma_client: Database client instance
user_api_key_cache: Cache for storing/retrieving budget data
parent_otel_span: Optional OpenTelemetry span for tracing
Returns:
LiteLLM_BudgetTable if configured and found, None otherwise
"""
if prisma_client is None or litellm.max_end_user_budget_id is None:
return None
cache_key = f"default_end_user_budget:{litellm.max_end_user_budget_id}"
# Check cache first
cached_budget = await user_api_key_cache.async_get_cache(key=cache_key)
if cached_budget is not None:
return LiteLLM_BudgetTable(**cached_budget)
# Fetch from database
try:
budget_record = await prisma_client.db.litellm_budgettable.find_unique(
where={"budget_id": litellm.max_end_user_budget_id}
)
if budget_record is None:
verbose_proxy_logger.warning(
f"Default end user budget not found in database: {litellm.max_end_user_budget_id}"
)
return None
# Cache the budget for 60 seconds
await user_api_key_cache.async_set_cache(
key=cache_key,
key=cache_key,
value=budget_record.dict(),
ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
)
return LiteLLM_BudgetTable(**budget_record.dict())
except Exception as e:
verbose_proxy_logger.error(
f"Error fetching default end user budget: {str(e)}"
)
verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}")
return None
@ -560,38 +566,38 @@ async def _apply_default_budget_to_end_user(
) -> LiteLLM_EndUserTable:
"""
Helper function to apply default budget to end user if they don't have a budget assigned.
Args:
end_user_obj: The end user object to potentially apply default budget to
prisma_client: Database client instance
user_api_key_cache: Cache for storing/retrieving data
parent_otel_span: Optional OpenTelemetry span for tracing
Returns:
Updated end user object with default budget applied if applicable
"""
# If end user already has a budget assigned, no need to apply default
if end_user_obj.litellm_budget_table is not None:
return end_user_obj
# If no default budget configured, return as-is
if litellm.max_end_user_budget_id is None:
return end_user_obj
# Fetch and apply default budget
default_budget = await get_default_end_user_budget(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
)
if default_budget is not None:
# Apply default budget to end user object
end_user_obj.litellm_budget_table = default_budget
verbose_proxy_logger.debug(
f"Applied default budget {litellm.max_end_user_budget_id} to end user {end_user_obj.user_id}"
)
return end_user_obj
@ -601,20 +607,20 @@ def _check_end_user_budget(
) -> None:
"""
Check if end user is within their budget limit.
Args:
end_user_obj: The end user object to check
route: The request route
Raises:
litellm.BudgetExceededError: If end user has exceeded their budget
"""
if route in LiteLLMRoutes.info_routes.value:
return
if end_user_obj.litellm_budget_table is None:
return
end_user_budget = end_user_obj.litellm_budget_table.max_budget
if end_user_budget is not None and end_user_obj.spend > end_user_budget:
raise litellm.BudgetExceededError(
@ -635,8 +641,8 @@ async def get_end_user_object(
) -> Optional[LiteLLM_EndUserTable]:
"""
Returns end user object from database or cache.
If end user exists but has no budget_id, applies the default budget
If end user exists but has no budget_id, applies the default budget
(if configured via litellm.max_end_user_budget_id).
Args:
@ -646,7 +652,7 @@ async def get_end_user_object(
route: The request route
parent_otel_span: Optional OpenTelemetry span for tracing
proxy_logging_obj: Optional proxy logging object
Returns:
LiteLLM_EndUserTable if found, None otherwise
"""
@ -655,14 +661,14 @@ async def get_end_user_object(
if end_user_id is None:
return None
_key = "end_user_id:{}".format(end_user_id)
# Check cache first
cached_user_obj = await user_api_key_cache.async_get_cache(key=_key)
if cached_user_obj is not None:
return_obj = LiteLLM_EndUserTable(**cached_user_obj)
# Apply default budget if needed
return_obj = await _apply_default_budget_to_end_user(
end_user_obj=return_obj,
@ -670,10 +676,10 @@ async def get_end_user_object(
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
)
# Check budget limits
_check_end_user_budget(end_user_obj=return_obj, route=route)
return return_obj
# Fetch from database
@ -688,7 +694,7 @@ async def get_end_user_object(
# Convert to LiteLLM_EndUserTable object
_response = LiteLLM_EndUserTable(**response.dict())
# Apply default budget if needed
_response = await _apply_default_budget_to_end_user(
end_user_obj=_response,
@ -696,18 +702,17 @@ async def get_end_user_object(
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
)
# Save to cache (always store as dict for consistency)
await user_api_key_cache.async_set_cache(
key="end_user_id:{}".format(end_user_id),
value=_response.dict()
key="end_user_id:{}".format(end_user_id), value=_response.dict()
)
# Check budget limits
_check_end_user_budget(end_user_obj=_response, route=route)
return _response
except Exception as e:
if isinstance(e, litellm.BudgetExceededError):
raise e
@ -747,7 +752,6 @@ async def get_tag_objects_batch(
tag_objects = {}
uncached_tags = []
# Try to get all tags from cache first
for tag_name in tag_names:
@ -1138,7 +1142,6 @@ async def _cache_management_object(
user_api_key_cache: DualCache,
proxy_logging_obj: Optional[ProxyLogging],
):
await user_api_key_cache.async_set_cache(
key=key,
value=value,
@ -1459,9 +1462,7 @@ async def get_team_object_by_alias(
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(
"Error looking up team by alias: %s", team_alias
)
verbose_proxy_logger.exception("Error looking up team by alias: %s", team_alias)
raise HTTPException(
status_code=500,
detail={
@ -1602,11 +1603,11 @@ class ExperimentalUIJWTToken:
) -> str:
"""
Generate a JWT token for CLI authentication with 24-hour expiration.
Args:
user_info: User information from the database
team_id: Team ID for the user (optional, uses user's team if available)
Returns:
Encrypted JWT token string
"""
@ -1800,7 +1801,7 @@ async def get_org_object(
- Check if org id in proxy Org Table
- if valid, return LiteLLM_OrganizationTable object
- if not, then raise an error
Args:
org_id: Organization ID to look up
prisma_client: Database client
@ -1820,7 +1821,7 @@ async def get_org_object(
cache_key = "org_id:{}".format(org_id)
if include_budget_table:
cache_key = "org_id:{}:with_budget".format(org_id)
# check if in cache
cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key)
if cached_org_obj is not None:
@ -1833,7 +1834,7 @@ async def get_org_object(
query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}}
if include_budget_table:
query_kwargs["include"] = {"litellm_budget_table": True}
response = await prisma_client.db.litellm_organizationtable.find_unique(
**query_kwargs
)
@ -1844,7 +1845,9 @@ async def get_org_object(
# Cache the result
await user_api_key_cache.async_set_cache(
key=cache_key,
value=response.model_dump() if hasattr(response, "model_dump") else response,
value=response.model_dump()
if hasattr(response, "model_dump")
else response,
ttl=DEFAULT_IN_MEMORY_TTL,
)
@ -2218,10 +2221,15 @@ async def _virtual_key_max_budget_alert_check(
and valid_token.spend is not None
and valid_token.spend > 0
):
alert_threshold = valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
alert_threshold = (
valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
)
# Only alert if we've crossed the threshold but haven't exceeded max_budget yet
if valid_token.spend >= alert_threshold and valid_token.spend < valid_token.max_budget:
if (
valid_token.spend >= alert_threshold
and valid_token.spend < valid_token.max_budget
):
verbose_proxy_logger.debug(
"Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s",
valid_token.token,
@ -2274,7 +2282,7 @@ async def _check_team_member_budget(
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if (
team_membership is not None
and team_membership.litellm_budget_table is not None
@ -2282,8 +2290,8 @@ async def _check_team_member_budget(
):
team_member_budget = team_membership.litellm_budget_table.max_budget
team_member_spend = team_membership.spend or 0.0
if team_member_spend > team_member_budget:
if team_member_spend >= team_member_budget:
raise litellm.BudgetExceededError(
current_cost=team_member_spend,
max_budget=team_member_budget,
@ -2343,11 +2351,11 @@ async def _organization_max_budget_check(
):
"""
Check if the organization is over its max budget.
This function checks the organization budget using:
1. First, tries to use valid_token.org_id (if key has organization_id set)
2. Falls back to team_object.organization_id (if key doesn't have org_id but team does)
This ensures organization budget checks work even when keys don't have organization_id
set directly, as long as their team belongs to an organization.
@ -2364,7 +2372,7 @@ async def _organization_max_budget_check(
org_id = valid_token.org_id
elif team_object is not None and team_object.organization_id is not None:
org_id = team_object.organization_id
# If no organization_id found, skip the check
if org_id is None:
return
@ -2655,4 +2663,4 @@ def _can_object_call_vector_stores(
code=status.HTTP_401_UNAUTHORIZED,
)
return True
return True

4
litellm/proxy/common_request_processing.py
View file

@ -49,9 +49,7 @@ if TYPE_CHECKING:
ProxyConfig = _ProxyConfig
else:
ProxyConfig = Any
from litellm.proxy.litellm_pre_call_utils import (
add_litellm_data_to_request,
)
from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
from litellm.types.utils import ModelResponse, ModelResponseStream, Usage

23
litellm/proxy/litellm_pre_call_utils.py
View file

@ -846,9 +846,7 @@ async def add_litellm_data_to_request( # noqa: PLR0915
# Add headers to metadata for guardrails to access (fixes #17477)
# Guardrails use metadata["headers"] to access request headers (e.g., User-Agent)
if _metadata_variable_name in data and isinstance(
data[_metadata_variable_name], dict
):
if _metadata_variable_name in data and isinstance(data[_metadata_variable_name], dict):
data[_metadata_variable_name]["headers"] = _headers
# check for forwardable headers
@ -1316,9 +1314,6 @@ def move_guardrails_to_metadata(
- If guardrails set on API Key metadata then sets guardrails on request metadata
- If guardrails not set on API key, then checks request metadata
Note: We copy (not pop) guardrails from data to metadata to ensure deployment-level
guardrails merged by the router remain in kwargs for async_pre_call_deployment_hook.
"""
# Check key-level guardrails
_add_guardrails_from_key_or_team_metadata(
@ -1331,25 +1326,15 @@ def move_guardrails_to_metadata(
#########################################################################################
# User's might send "guardrails" in the request body, we need to add them to the request metadata.
# Since downstream logic requires "guardrails" to be in the request metadata
#
# IMPORTANT: We copy instead of pop to preserve guardrails in kwargs for
# async_pre_call_deployment_hook (custom_guardrail.py:290) which checks kwargs.get("guardrails").
# This is the event-based approach for deployment-level guardrails.
#########################################################################################
if "guardrails" in data:
request_body_guardrails = data.get("guardrails")
if request_body_guardrails is None:
return
request_body_guardrails = data.pop("guardrails")
if "guardrails" in data[_metadata_variable_name] and isinstance(
data[_metadata_variable_name]["guardrails"], list
):
# Merge unique guardrails
existing = data[_metadata_variable_name]["guardrails"]
for g in request_body_guardrails:
if g not in existing:
existing.append(g)
data[_metadata_variable_name]["guardrails"].extend(request_body_guardrails)
else:
data[_metadata_variable_name]["guardrails"] = list(request_body_guardrails)
data[_metadata_variable_name]["guardrails"] = request_body_guardrails
#########################################################################################
if "guardrail_config" in data:

9
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -343,7 +343,7 @@ def _build_where_conditions(
start_date: str,
end_date: str,
model: Optional[str],
api_key: Optional[str],
api_key: Optional[Union[str, List[str]]],
exclude_entity_ids: Optional[List[str]] = None,
) -> Dict[str, Any]:
"""Build prisma where clause for daily activity queries."""
@ -357,7 +357,10 @@ def _build_where_conditions(
if model:
where_conditions["model"] = model
if api_key:
where_conditions["api_key"] = api_key
if isinstance(api_key, list):
where_conditions["api_key"] = {"in": api_key}
else:
where_conditions["api_key"] = api_key
if entity_id is not None:
if isinstance(entity_id, list):
@ -445,7 +448,7 @@ async def get_daily_activity(
start_date: Optional[str],
end_date: Optional[str],
model: Optional[str],
api_key: Optional[str],
api_key: Optional[Union[str, List[str]]],
page: int,
page_size: int,
exclude_entity_ids: Optional[List[str]] = None,

13
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -412,19 +412,6 @@ async def new_user(
status_code=403,
detail="License is over limit. Please contact support@berri.ai to upgrade your license.",
)
# Only proxy admins can create administrative users
# Check if user_api_key_dict is actually a UserAPIKeyAuth instance (not a Depends object)
# This can happen when the function is called directly in tests
if (
data.user_role in [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]
and isinstance(user_api_key_dict, UserAPIKeyAuth)
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
):
raise HTTPException(
status_code=403,
detail=f"Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). Attempted to create user with role: {data.user_role}. Your role: {user_api_key_dict.user_role}"
)
data_json = data.json() # type: ignore
data_json = _update_internal_new_user_params(data_json, data)

141
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -16,7 +16,7 @@ import secrets
import traceback
import yaml
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, List, Literal, Optional, Tuple, cast
from typing import List, Literal, Optional, Tuple, cast
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status
@ -1791,10 +1791,6 @@ async def delete_key_fn(
if prisma_client is None:
raise Exception("Not connected to DB!")
# Normalize litellm_changed_by: if it's a Header object or not a string, convert to None
if litellm_changed_by is not None and not isinstance(litellm_changed_by, str):
litellm_changed_by = None
## only allow user to delete keys they own
verbose_proxy_logger.debug(
f"user_api_key_dict.user_role: {user_api_key_dict.user_role}"
@ -1807,7 +1803,6 @@ async def delete_key_fn(
tokens=data.keys,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
num_keys_to_be_deleted = len(data.keys)
deleted_keys = data.keys
@ -1817,7 +1812,6 @@ async def delete_key_fn(
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
num_keys_to_be_deleted = len(data.key_aliases)
deleted_keys = data.key_aliases
@ -2439,7 +2433,6 @@ async def delete_verification_tokens(
tokens: List,
user_api_key_cache: DualCache,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]:
"""
Helper that deletes the list of tokens from the database
@ -2476,43 +2469,38 @@ async def delete_verification_tokens(
detail={"error": "No keys found"},
)
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
authorized_keys = _keys_being_deleted
else:
authorized_keys = []
for key in _keys_being_deleted:
if await can_modify_verification_token(
key_info=key,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
):
authorized_keys.append(key)
else:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={
"error": "You are not authorized to delete this key"
},
)
await _persist_deleted_verification_tokens(
keys=authorized_keys,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
# Assuming 'db' is your Prisma Client instance
# check if admin making request - don't filter by user-id
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
deleted_tokens = await prisma_client.delete_data(tokens=tokens)
# else
else:
deletion_tasks = [
prisma_client.delete_data(tokens=[key.token])
for key in authorized_keys
]
await asyncio.gather(*deletion_tasks)
tasks = []
deleted_tokens = []
for key in _keys_being_deleted:
deleted_tokens = [key.token for key in authorized_keys]
if len(deleted_tokens) != len(tokens):
async def _delete_key(key: LiteLLM_VerificationToken):
if await can_modify_verification_token(
key_info=key,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
):
await prisma_client.delete_data(tokens=[key.token])
deleted_tokens.append(key.token)
else:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail={
"error": "You are not authorized to delete this key"
},
)
tasks.append(_delete_key(key))
await asyncio.gather(*tasks)
_num_deleted_tokens = len(deleted_tokens)
if _num_deleted_tokens != len(tokens):
failed_tokens = [
token for token in tokens if token not in deleted_tokens
]
@ -2540,81 +2528,11 @@ async def delete_verification_tokens(
return {"deleted_keys": deleted_tokens}, _keys_being_deleted
def _transform_verification_tokens_to_deleted_records(
keys: List[LiteLLM_VerificationToken],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""Transform verification tokens into deleted token records ready for persistence."""
if not keys:
return []
deleted_at = datetime.now(timezone.utc)
records = []
for key in keys:
key_payload = key.model_dump()
deleted_record = LiteLLM_DeletedVerificationToken(
**key_payload,
deleted_at=deleted_at,
deleted_by=user_api_key_dict.user_id,
deleted_by_api_key=user_api_key_dict.api_key,
litellm_changed_by=litellm_changed_by,
)
record = deleted_record.model_dump()
# Map org_id to organization_id (model uses org_id, but schema expects organization_id)
org_id_value = record.pop("org_id", None)
if org_id_value is not None:
record["organization_id"] = org_id_value
for json_field in ["aliases", "config", "permissions", "metadata", "model_spend", "model_max_budget", "router_settings"]:
if json_field in record and record[json_field] is not None:
record[json_field] = json.dumps(record[json_field])
for rel_key in ("litellm_budget_table", "litellm_organization_table", "object_permission", "id"):
record.pop(rel_key, None)
records.append(record)
return records
async def _save_deleted_verification_token_records(
records: List[Dict[str, Any]],
prisma_client: PrismaClient,
) -> None:
"""Save deleted verification token records to the database."""
if not records:
return
await prisma_client.db.litellm_deletedverificationtoken.create_many(
data=records
)
async def _persist_deleted_verification_tokens(
keys: List[LiteLLM_VerificationToken],
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> None:
"""Persist deleted verification token records by transforming and saving them."""
records = _transform_verification_tokens_to_deleted_records(
keys=keys,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await _save_deleted_verification_token_records(
records=records,
prisma_client=prisma_client,
)
async def delete_key_aliases(
key_aliases: List[str],
user_api_key_cache: DualCache,
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> Tuple[Optional[Dict], List[LiteLLM_VerificationToken]]:
_keys_being_deleted = await prisma_client.db.litellm_verificationtoken.find_many(
where={"key_alias": {"in": key_aliases}}
@ -2625,7 +2543,6 @@ async def delete_key_aliases(
tokens=tokens,
user_api_key_cache=user_api_key_cache,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)

148
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -34,10 +34,8 @@ from litellm.proxy._types import (
LiteLLM_OrganizationTableWithMembers,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_DeletedTeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LiteLLM_VerificationToken,
LitellmTableNames,
LitellmUserRoles,
Member,
@ -2020,28 +2018,6 @@ async def team_member_delete(
## DELETE KEYS CREATED BY USER FOR THIS TEAM
if user_ids_to_delete:
from litellm.proxy.management_endpoints.key_management_endpoints import (
_persist_deleted_verification_tokens,
)
# Fetch keys before deletion to persist them
keys_to_delete: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={
"user_id": {"in": list(user_ids_to_delete)},
"team_id": data.team_id,
}
)
)
if keys_to_delete:
await _persist_deleted_verification_tokens(
keys=keys_to_delete,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=None,
)
await prisma_client.db.litellm_verificationtoken.delete_many(
where={
"user_id": {"in": list(user_ids_to_delete)},
@ -2427,13 +2403,6 @@ async def delete_team(
team_row_pydantic = LiteLLM_TeamTable(**team_row_base.model_dump())
team_rows.append(team_row_pydantic)
await _persist_deleted_team_records(
teams=team_rows,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
if litellm.store_audit_logs is True:
@ -2469,25 +2438,6 @@ async def delete_team(
# End of Audit logging
## DELETE ASSOCIATED KEYS
# Fetch keys before deletion to persist them
from litellm.proxy.management_endpoints.key_management_endpoints import (
_persist_deleted_verification_tokens,
)
keys_to_delete: List[LiteLLM_VerificationToken] = (
await prisma_client.db.litellm_verificationtoken.find_many(
where={"team_id": {"in": data.team_ids}}
)
)
if keys_to_delete:
await _persist_deleted_verification_tokens(
keys=keys_to_delete,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key")
# ## DELETE TEAM MEMBERSHIPS
@ -2516,70 +2466,6 @@ async def delete_team(
return deleted_teams
def _transform_teams_to_deleted_records(
teams: List[LiteLLM_TeamTable],
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""Transform teams into deleted team records ready for persistence."""
if not teams:
return []
deleted_at = datetime.now(timezone.utc)
records = []
for team in teams:
team_payload = team.model_dump()
deleted_record = LiteLLM_DeletedTeamTable(
**team_payload,
deleted_at=deleted_at,
deleted_by=user_api_key_dict.user_id,
deleted_by_api_key=user_api_key_dict.api_key,
litellm_changed_by=litellm_changed_by,
)
record = deleted_record.model_dump()
for json_field in ["members_with_roles", "metadata", "model_spend", "model_max_budget", "router_settings"]:
if json_field in record and record[json_field] is not None:
record[json_field] = json.dumps(record[json_field])
for rel_key in ("litellm_model_table", "object_permission", "id"):
record.pop(rel_key, None)
records.append(record)
return records
async def _save_deleted_team_records(
records: List[Dict[str, Any]],
prisma_client: PrismaClient,
) -> None:
"""Save deleted team records to the database."""
if not records:
return
await prisma_client.db.litellm_deletedteamtable.create_many(
data=records
)
async def _persist_deleted_team_records(
teams: List[LiteLLM_TeamTable],
prisma_client: PrismaClient,
user_api_key_dict: UserAPIKeyAuth,
litellm_changed_by: Optional[str] = None,
) -> None:
"""Persist deleted team records by transforming and saving them."""
records = _transform_teams_to_deleted_records(
teams=teams,
user_api_key_dict=user_api_key_dict,
litellm_changed_by=litellm_changed_by,
)
await _save_deleted_team_records(
records=records,
prisma_client=prisma_client,
)
def validate_membership(
user_api_key_dict: UserAPIKeyAuth, team_table: LiteLLM_TeamTable
):
@ -3715,7 +3601,7 @@ async def get_team_daily_activity(
},
)
## Fetch team aliases
## Fetch team aliases and check team admin status
where_condition = {}
if team_ids_list:
where_condition["team_id"] = {"in": list(team_ids_list)}
@ -3726,6 +3612,36 @@ async def get_team_daily_activity(
t.team_id: {"team_alias": t.team_alias} for t in team_aliases
}
# Check if user is team admin for any requested teams
# If not, filter by user's API keys
user_api_keys: Optional[List[str]] = None
if not _user_has_admin_view(user_api_key_dict) and team_ids_list and team_aliases:
# Check if user is team admin for any of the teams
is_team_admin_for_any = False
for team_alias in team_aliases:
team_obj = LiteLLM_TeamTable(**team_alias.model_dump())
if _is_user_team_admin(
user_api_key_dict=user_api_key_dict, team_obj=team_obj
):
is_team_admin_for_any = True
break
# If user is not a team admin for any team, filter by their API keys
if not is_team_admin_for_any:
# Get all API keys for this user
user_keys = await prisma_client.db.litellm_verificationtoken.find_many(
where={"user_id": user_api_key_dict.user_id}
)
user_api_keys = [key.token for key in user_keys if key.token]
# If user has no API keys, return empty result
if not user_api_keys:
user_api_keys = [""] # Use empty string to ensure no matches
# If api_key parameter is provided, use it; otherwise use user_api_keys if set
final_api_key_filter: Optional[Union[str, List[str]]] = api_key
if final_api_key_filter is None and user_api_keys is not None:
final_api_key_filter = user_api_keys
return await get_daily_activity(
prisma_client=prisma_client,
table_name="litellm_dailyteamspend",
@ -3736,7 +3652,7 @@ async def get_team_daily_activity(
start_date=start_date,
end_date=end_date,
model=model,
api_key=api_key,
api_key=final_api_key_filter,
page=page,
page_size=page_size,
)

1
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -761,6 +761,7 @@ async def handle_bedrock_passthrough_router_model(
proxy_logging_obj=proxy_logging_obj,
)
async def handle_bedrock_count_tokens(
endpoint: str,
request: Request,

2
litellm/proxy/prisma_migration.py
View file

@ -26,5 +26,3 @@ if exit_code != 0:
verbose_proxy_logger.error(
f"'prisma generate' stderr: {result.stderr}"
) # Log stderr
sys.exit(exit_code)

8
litellm/proxy/proxy_cli.py
View file

@ -187,7 +187,6 @@ class ProxyInitializationHelpers:
ssl_certfile_path: str,
ssl_keyfile_path: str,
max_requests_before_restart: Optional[int] = None,
keepalive_timeout: Optional[int] = None,
):
"""
Run litellm with `gunicorn`
@ -268,10 +267,6 @@ class ProxyInitializationHelpers:
"access_log_format": '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s',
}
# Optional: set keepalive timeout if specified by user
if keepalive_timeout is not None:
gunicorn_options["keepalive"] = keepalive_timeout
# Optional: recycle workers after N requests to mitigate memory growth
if max_requests_before_restart is not None:
gunicorn_options["max_requests"] = max_requests_before_restart
@ -494,7 +489,7 @@ class ProxyInitializationHelpers:
"--keepalive_timeout",
default=None,
type=int,
help="Set the keepalive timeout in seconds. For Uvicorn: timeout_keep_alive parameter. For Gunicorn: keepalive parameter. Default: Uvicorn uses ~75s, Gunicorn uses 90s",
help="Set the uvicorn keepalive timeout in seconds (uvicorn timeout_keep_alive parameter)",
envvar="KEEPALIVE_TIMEOUT",
)
@click.option(
@ -864,7 +859,6 @@ def run_server( # noqa: PLR0915
ssl_certfile_path=ssl_certfile_path,
ssl_keyfile_path=ssl_keyfile_path,
max_requests_before_restart=max_requests_before_restart,
keepalive_timeout=keepalive_timeout,
)
elif run_hypercorn is True:
ProxyInitializationHelpers._init_hypercorn_server(

99
litellm/proxy/schema.prisma
View file

@ -132,49 +132,6 @@ model LiteLLM_TeamTable {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted teams - preserves spend and team information for historical tracking
model LiteLLM_DeletedTeamTable {
id String @id @default(uuid())
team_id String // Original team_id
team_alias String?
organization_id String?
object_permission_id String?
admins String[]
members String[]
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
spend Float @default(0.0)
models String[]
max_parallel_requests Int?
tpm_limit BigInt?
rpm_limit BigInt?
budget_duration String?
budget_reset_at DateTime?
blocked Boolean @default(false)
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
// Original timestamps from team creation/updates
created_at DateTime? @map("created_at")
updated_at DateTime? @map("updated_at")
// Deletion metadata
deleted_at DateTime @default(now()) @map("deleted_at")
deleted_by String? @map("deleted_by") // User who deleted the team
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([team_id])
@@index([deleted_at])
@@index([organization_id])
@@index([team_alias])
@@index([created_at])
}
// Track spend, rate limit, budget Users
model LiteLLM_UserTable {
user_id String @id
@ -302,62 +259,6 @@ model LiteLLM_VerificationToken {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted keys - preserves spend and key information for historical tracking
model LiteLLM_DeletedVerificationToken {
id String @id @default(uuid())
token String // Original token (hashed)
key_name String?
key_alias String?
soft_budget_cooldown Boolean @default(false)
spend Float @default(0.0)
expires DateTime?
models String[]
aliases Json @default("{}")
config Json @default("{}")
user_id String?
team_id String?
permissions Json @default("{}")
max_parallel_requests Int?
metadata Json @default("{}")
blocked Boolean?
tpm_limit BigInt?
rpm_limit BigInt?
max_budget Float?
budget_duration String?
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
budget_id String?
organization_id String?
object_permission_id String?
created_at DateTime? // Original creation timestamp
created_by String? // Original creator
updated_at DateTime? // Last update timestamp before deletion
updated_by String? // Last user who updated before deletion
rotation_count Int? @default(0)
auto_rotate Boolean? @default(false)
rotation_interval String?
last_rotation_at DateTime?
key_rotation_at DateTime?
// Deletion metadata
deleted_at DateTime @default(now()) @map("deleted_at")
deleted_by String? @map("deleted_by") // User who deleted the key
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([token])
@@index([deleted_at])
@@index([user_id])
@@index([team_id])
@@index([organization_id])
@@index([key_alias])
@@index([created_at])
}
model LiteLLM_EndUserTable {
user_id String @id
alias String? // admin-facing alias

12
litellm/proxy/video_endpoints/endpoints.py
View file

@ -256,9 +256,7 @@ async def video_status(
# Resolve model_name from model_id if available
# This allows the router to automatically inject litellm_params from the model config
if model_id_from_decoded and llm_router:
resolved_model = llm_router.resolve_model_name_from_model_id(
model_id_from_decoded, custom_llm_provider=provider_from_id
)
resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded)
if resolved_model:
data["model"] = resolved_model
@ -356,9 +354,7 @@ async def video_content(
# Resolve model_name from model_id if available
# This allows the router to automatically inject litellm_params from the model config
if model_id_from_decoded and llm_router:
resolved_model = llm_router.resolve_model_name_from_model_id(
model_id_from_decoded, custom_llm_provider=provider_from_id
)
resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded)
if resolved_model:
data["model"] = resolved_model
# Process request using ProxyBaseLLMRequestProcessing
@ -470,9 +466,7 @@ async def video_remix(
# Resolve model_name from model_id if available
# This allows the router to automatically inject litellm_params from the model config
if model_id_from_decoded and llm_router:
resolved_model = llm_router.resolve_model_name_from_model_id(
model_id_from_decoded, custom_llm_provider=provider_from_id
)
resolved_model = llm_router.resolve_model_name_from_model_id(model_id_from_decoded)
if resolved_model:
data["model"] = resolved_model

44
litellm/router.py
View file

@ -1408,15 +1408,6 @@ class Router:
async for item in model_response:
yield item
except MidStreamFallbackError as e:
# Check if fallbacks are disabled by user
if initial_kwargs.get("disable_fallbacks", False):
verbose_router_logger.info(
"Mid stream fallback disabled by user, re-raising original error"
)
if e.original_exception is not None:
raise e.original_exception
raise e
from litellm.main import stream_chunk_builder
complete_response_object = stream_chunk_builder(
@ -6971,7 +6962,7 @@ class Router:
return candidate_id in self.model_id_to_deployment_index_map
def resolve_model_name_from_model_id(
self, model_id: Optional[str], custom_llm_provider: Optional[str] = None
self, model_id: Optional[str]
) -> Optional[str]:
"""
Resolve model_name from model_id.
@ -6981,15 +6972,12 @@ class Router:
Strategy:
1. First, check if model_id directly matches a model_name or deployment ID
2. If custom_llm_provider is provided, check with provider prefix
3. Search through router's model_list to find a match by litellm_params.model
4. If custom_llm_provider is provided, try to find a wildcard pattern match
5. Return the model_name if found, None otherwise
2. If not, search through router's model_list to find a match by litellm_params.model
3. Return the model_name if found, None otherwise
Args:
model_id: The model_id extracted from decoded video_id
(could be model_name or litellm_params.model value)
custom_llm_provider: The provider name (e.g., "vertex_ai") for wildcard matching
Returns:
model_name if found, None otherwise. If None, the request will fall through
@ -7002,26 +6990,15 @@ class Router:
if model_id in self.model_names or self.has_model_id(model_id):
return model_id
# Strategy 2: Check with provider prefix (e.g., "vertex_ai/veo-3.0-generate-preview")
if custom_llm_provider:
full_model_name = f"{custom_llm_provider}/{model_id}"
if full_model_name in self.model_names or self.has_model_id(full_model_name):
return full_model_name
# Strategy 3: Search through router's model_list to find by litellm_params.model
# Strategy 2: Search through router's model_list to find by litellm_params.model
all_models = self.get_model_list(model_name=None)
if not all_models:
return None
# First pass: exact matches (non-wildcard)
for deployment in all_models:
litellm_params = deployment.get("litellm_params", {})
actual_model = litellm_params.get("model")
# Skip wildcard patterns in first pass
if actual_model and actual_model.endswith("/*"):
continue
# Match by exact match or by checking if actual_model ends with /model_id or :model_id
# e.g., model_id="veo-2.0-generate-001" matches actual_model="vertex_ai/veo-2.0-generate-001"
matches = (
@ -7035,19 +7012,6 @@ class Router:
if model_name:
return model_name
# Strategy 4: Wildcard patterns using PatternMatchRouter
# For video status/content, we need to match model_id like "veo-3.0-generate-preview"
# to wildcard patterns like "vertex_ai/*"
if custom_llm_provider:
full_model_name = f"{custom_llm_provider}/{model_id}"
pattern_deployments = self.pattern_router.route(full_model_name)
if pattern_deployments:
# Return the first matching wildcard model_name
for pattern_deployment in pattern_deployments:
matched_model_name = pattern_deployment.get("model_name")
if matched_model_name:
return matched_model_name
# No match found
return None

2
litellm/types/llms/bedrock.py
View file

@ -62,7 +62,7 @@ class ToolResultBlock(TypedDict, total=False):
class ToolUseBlock(TypedDict):
input: Any # Per boto3 spec: document type can be dict, list, int, float, str, bool, or None
input: dict
name: str
toolUseId: str

42
model_prices_and_context_window.json
View file

@ -10201,48 +10201,6 @@
"mode": "completion",
"output_cost_per_token": 5e-07
},
"deepseek-v3-2-251201": {
"input_cost_per_token": 0.0,
"litellm_provider": "volcengine",
"max_input_tokens": 98304,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 0.0,
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"glm-4-7-251222": {
"input_cost_per_token": 0.0,
"litellm_provider": "volcengine",
"max_input_tokens": 204800,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 0.0,
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"kimi-k2-thinking-251104": {
"input_cost_per_token": 0.0,
"litellm_provider": "volcengine",
"max_input_tokens": 229376,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 0.0,
"supports_assistant_prefill": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"doubao-embedding": {
"input_cost_per_token": 0.0,
"litellm_provider": "volcengine",

38
poetry.lock generated
View file

@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand.
# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand.
[[package]]
name = "aiofiles"
@ -525,36 +525,36 @@ files = [
[[package]]
name = "boto3"
version = "1.40.61"
version = "1.36.0"
description = "The AWS SDK for Python"
optional = true
python-versions = ">=3.9"
python-versions = ">=3.8"
groups = ["main"]
markers = "extra == \"proxy\""
files = [
{file = "boto3-1.40.61-py3-none-any.whl", hash = "sha256:6b9c57b2a922b5d8c17766e29ed792586a818098efe84def27c8f582b33f898c"},
{file = "boto3-1.40.61.tar.gz", hash = "sha256:d6c56277251adf6c2bdd25249feae625abe4966831676689ff23b4694dea5b12"},
{file = "boto3-1.36.0-py3-none-any.whl", hash = "sha256:d0ca7a58ce25701a52232cc8df9d87854824f1f2964b929305722ebc7959d5a9"},
{file = "boto3-1.36.0.tar.gz", hash = "sha256:159898f51c2997a12541c0e02d6e5a8fe2993ddb307b9478fd9a339f98b57e00"},
]
[package.dependencies]
botocore = ">=1.40.61,<1.41.0"
botocore = ">=1.36.0,<1.37.0"
jmespath = ">=0.7.1,<2.0.0"
s3transfer = ">=0.14.0,<0.15.0"
s3transfer = ">=0.11.0,<0.12.0"
[package.extras]
crt = ["botocore[crt] (>=1.21.0,<2.0a0)"]
[[package]]
name = "botocore"
version = "1.40.76"
version = "1.36.26"
description = "Low-level, data-driven core of boto 3."
optional = true
python-versions = ">=3.9"
python-versions = ">=3.8"
groups = ["main"]
markers = "extra == \"proxy\""
files = [
{file = "botocore-1.40.76-py3-none-any.whl", hash = "sha256:fe425d386e48ac64c81cbb4a7181688d813df2e2b4c78b95ebe833c9e868c6f4"},
{file = "botocore-1.40.76.tar.gz", hash = "sha256:2b16024d68b29b973005adfb5039adfe9099ebe772d40a90ca89f2e165c495dc"},
{file = "botocore-1.36.26-py3-none-any.whl", hash = "sha256:4e3f19913887a58502e71ef8d696fe7eaa54de7813ff73390cd5883f837dfa6e"},
{file = "botocore-1.36.26.tar.gz", hash = "sha256:4a63bcef7ecf6146fd3a61dc4f9b33b7473b49bdaf1770e9aaca6eee0c9eab62"},
]
[package.dependencies]
@ -566,7 +566,7 @@ urllib3 = [
]
[package.extras]
crt = ["awscrt (==0.28.4)"]
crt = ["awscrt (==0.23.8)"]
[[package]]
name = "cachetools"
@ -6255,22 +6255,22 @@ files = [
[[package]]
name = "s3transfer"
version = "0.14.0"
version = "0.11.3"
description = "An Amazon S3 Transfer Manager"
optional = true
python-versions = ">=3.9"
python-versions = ">=3.8"
groups = ["main"]
markers = "extra == \"proxy\""
files = [
{file = "s3transfer-0.14.0-py3-none-any.whl", hash = "sha256:ea3b790c7077558ed1f02a3072fb3cb992bbbd253392f4b6e9e8976941c7d456"},
{file = "s3transfer-0.14.0.tar.gz", hash = "sha256:eff12264e7c8b4985074ccce27a3b38a485bb7f7422cc8046fee9be4983e4125"},
{file = "s3transfer-0.11.3-py3-none-any.whl", hash = "sha256:ca855bdeb885174b5ffa95b9913622459d4ad8e331fc98eb01e6d5eb6a30655d"},
{file = "s3transfer-0.11.3.tar.gz", hash = "sha256:edae4977e3a122445660c7c114bba949f9d191bae3b34a096f18a1c8c354527a"},
]
[package.dependencies]
botocore = ">=1.37.4,<2.0a.0"
botocore = ">=1.36.0,<2.0a.0"
[package.extras]
crt = ["botocore[crt] (>=1.37.4,<2.0a.0)"]
crt = ["botocore[crt] (>=1.36.0,<2.0a.0)"]
[[package]]
name = "scikit-learn"
@ -7981,4 +7981,4 @@ utils = ["numpydoc"]
[metadata]
lock-version = "2.1"
python-versions = ">=3.9,<4.0"
content-hash = "f391c702cf58ef2ba7641acdc3ae13d7c8e672faede68c0a624bd2ba0fb46b12"
content-hash = "ea62b77c662ab9fc486e421c576f0868bcde16d62a24703ee1f4916a0465ffb2"

2
pyproject.toml
View file

@ -56,7 +56,7 @@ google-cloud-iam = {version = "^2.19.1", optional = true}
resend = {version = ">=0.8.0", optional = true}
pynacl = {version = "^1.5.0", optional = true}
websockets = {version = "^15.0.1", optional = true}
boto3 = {version = "1.40.61", optional = true}
boto3 = {version = "1.36.0", optional = true}
redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
mcp = {version = "^1.21.2", optional = true, python = ">=3.10"}
litellm-proxy-extras = {version = "0.4.22", optional = true}

4
requirements.txt
View file

@ -10,7 +10,7 @@ uvicorn==0.31.1 # server dep
gunicorn==23.0.0 # server dep
fastuuid==0.13.5 # for uuid4
uvloop==0.21.0 # uvicorn dep, gives us much better performance under load
boto3==1.40.61 # aws bedrock/sagemaker calls
boto3==1.36.0 # aws bedrock/sagemaker calls
redis==5.2.1 # redis caching
prisma==0.11.0 # for db
nodejs-wheel-binaries==24.12.0 ## required by prisma for migrations, prevents runtime download (updated from nodejs-bin for security fixes)
@ -59,7 +59,7 @@ click==8.1.7 # for proxy cli
rich==13.7.1 # for litellm proxy cli
jinja2==3.1.6 # for prompt templates
aiohttp==3.13.3 # for network calls
aioboto3==15.5.0 # for async sagemaker calls
aioboto3==13.4.0 # for async sagemaker calls
tenacity==8.5.0 # for retrying requests, when litellm.num_retries set
pydantic>=2.11,<3 # proxy + openai req. + mcp
jsonschema>=4.23.0,<5.0.0 # validating json schema - aligned with openapi-core + mcp

99
schema.prisma
View file

@ -132,49 +132,6 @@ model LiteLLM_TeamTable {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted teams - preserves spend and team information for historical tracking
model LiteLLM_DeletedTeamTable {
id String @id @default(uuid())
team_id String // Original team_id
team_alias String?
organization_id String?
object_permission_id String?
admins String[]
members String[]
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
spend Float @default(0.0)
models String[]
max_parallel_requests Int?
tpm_limit BigInt?
rpm_limit BigInt?
budget_duration String?
budget_reset_at DateTime?
blocked Boolean @default(false)
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
team_member_permissions String[] @default([])
model_id Int? // id for LiteLLM_ModelTable -> stores team-level model aliases
// Original timestamps from team creation/updates
created_at DateTime? @map("created_at")
updated_at DateTime? @map("updated_at")
// Deletion metadata
deleted_at DateTime @default(now()) @map("deleted_at")
deleted_by String? @map("deleted_by") // User who deleted the team
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([team_id])
@@index([deleted_at])
@@index([organization_id])
@@index([team_alias])
@@index([created_at])
}
// Track spend, rate limit, budget Users
model LiteLLM_UserTable {
user_id String @id
@ -302,62 +259,6 @@ model LiteLLM_VerificationToken {
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
}
// Audit table for deleted keys - preserves spend and key information for historical tracking
model LiteLLM_DeletedVerificationToken {
id String @id @default(uuid())
token String // Original token (hashed)
key_name String?
key_alias String?
soft_budget_cooldown Boolean @default(false)
spend Float @default(0.0)
expires DateTime?
models String[]
aliases Json @default("{}")
config Json @default("{}")
user_id String?
team_id String?
permissions Json @default("{}")
max_parallel_requests Int?
metadata Json @default("{}")
blocked Boolean?
tpm_limit BigInt?
rpm_limit BigInt?
max_budget Float?
budget_duration String?
budget_reset_at DateTime?
allowed_cache_controls String[] @default([])
allowed_routes String[] @default([])
model_spend Json @default("{}")
model_max_budget Json @default("{}")
router_settings Json? @default("{}")
budget_id String?
organization_id String?
object_permission_id String?
created_at DateTime? // Original creation timestamp
created_by String? // Original creator
updated_at DateTime? // Last update timestamp before deletion
updated_by String? // Last user who updated before deletion
rotation_count Int? @default(0)
auto_rotate Boolean? @default(false)
rotation_interval String?
last_rotation_at DateTime?
key_rotation_at DateTime?
// Deletion metadata
deleted_at DateTime @default(now()) @map("deleted_at")
deleted_by String? @map("deleted_by") // User who deleted the key
deleted_by_api_key String? @map("deleted_by_api_key") // API key hash that performed the deletion
litellm_changed_by String? @map("litellm_changed_by") // From litellm-changed-by header if provided
@@index([token])
@@index([deleted_at])
@@index([user_id])
@@index([team_id])
@@index([organization_id])
@@index([key_alias])
@@index([created_at])
}
model LiteLLM_EndUserTable {
user_id String @id
alias String? // admin-facing alias

4
tests/code_coverage_tests/license_cache.json
View file

File diff suppressed because one or more lines are too long

154
tests/llm_translation/test_bedrock_completion.py
View file

@ -3954,157 +3954,3 @@ def test_bedrock_openai_error_handling():
assert exc_info.value.status_code == 422
print("✓ Error handling works correctly")
def test_bedrock_malformed_tool_json_handling():
"""
Test that Bedrock handles malformed JSON in tool call arguments gracefully.
This test covers the issue where:
1. LLM generates malformed JSON in tool call arguments
2. Subsequent requests with conversation history should not crash
3. The toolUse.input field should handle any JSON value type per boto3 spec
Related issue: https://github.com/BerriAI/litellm/issues/[issue_number]
"""
from litellm.litellm_core_utils.prompt_templates.factory import (
_convert_to_bedrock_tool_call_invoke,
)
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
from litellm.types.llms.bedrock import ContentBlock
# Test 1: Malformed JSON in tool call arguments
malformed_tool_calls = [
{
"id": "call_123",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "Paris", "invalid_json', # Malformed JSON
},
}
]
# Should not raise an exception, but store as raw string
result = _convert_to_bedrock_tool_call_invoke(malformed_tool_calls)
assert len(result) == 1
assert result[0]["toolUse"]["name"] == "get_weather"
# The malformed JSON should be stored as a string
assert isinstance(result[0]["toolUse"]["input"], str)
assert result[0]["toolUse"]["input"] == '{"location": "Paris", "invalid_json'
print("✓ Malformed JSON stored as raw string")
# Test 2: Valid JSON should still work normally
valid_tool_calls = [
{
"id": "call_456",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "London"}',
},
}
]
result = _convert_to_bedrock_tool_call_invoke(valid_tool_calls)
assert len(result) == 1
assert result[0]["toolUse"]["name"] == "get_weather"
assert isinstance(result[0]["toolUse"]["input"], dict)
assert result[0]["toolUse"]["input"] == {"location": "London"}
print("✓ Valid JSON parsed correctly")
# Test 3: Empty arguments should create empty dict
empty_tool_calls = [
{
"id": "call_789",
"type": "function",
"function": {
"name": "no_args_function",
"arguments": "",
},
}
]
result = _convert_to_bedrock_tool_call_invoke(empty_tool_calls)
assert len(result) == 1
assert result[0]["toolUse"]["input"] == {}
print("✓ Empty arguments handled correctly")
# Test 4: Bedrock to OpenAI conversion handles string input
converse_config = AmazonConverseConfig()
content_blocks = [
ContentBlock(
toolUse={
"name": "get_weather",
"toolUseId": "call_123",
"input": '{"location": "Paris", "invalid_json', # String input (malformed)
}
)
]
content_str, tools, reasoning = converse_config._translate_message_content(
content_blocks
)
assert len(tools) == 1
assert tools[0]["function"]["name"] == "get_weather"
# Should return the string as-is
assert tools[0]["function"]["arguments"] == '{"location": "Paris", "invalid_json'
print("✓ Bedrock to OpenAI conversion handles string input")
# Test 5: Bedrock to OpenAI conversion handles dict input
content_blocks_dict = [
ContentBlock(
toolUse={
"name": "get_weather",
"toolUseId": "call_456",
"input": {"location": "London"}, # Dict input (normal case)
}
)
]
content_str, tools, reasoning = converse_config._translate_message_content(
content_blocks_dict
)
assert len(tools) == 1
assert tools[0]["function"]["name"] == "get_weather"
# Should serialize dict to JSON string
assert tools[0]["function"]["arguments"] == '{"location": "London"}'
print("✓ Bedrock to OpenAI conversion handles dict input")
# Test 6: Round-trip conversion with malformed JSON
# Test that we can convert OpenAI -> Bedrock -> OpenAI with malformed JSON
malformed_tool_calls_roundtrip = [
{
"id": "call_999",
"type": "function",
"function": {
"name": "test_function",
"arguments": '{"key": "value", "broken', # Malformed
},
}
]
# Step 1: OpenAI to Bedrock (should store as string)
bedrock_blocks = _convert_to_bedrock_tool_call_invoke(malformed_tool_calls_roundtrip)
assert isinstance(bedrock_blocks[0]["toolUse"]["input"], str)
# Step 2: Bedrock back to OpenAI (should preserve the string)
content_blocks_roundtrip = [
ContentBlock(
toolUse={
"name": bedrock_blocks[0]["toolUse"]["name"],
"toolUseId": bedrock_blocks[0]["toolUse"]["toolUseId"],
"input": bedrock_blocks[0]["toolUse"]["input"],
}
)
]
content_str, tools_roundtrip, reasoning = converse_config._translate_message_content(
content_blocks_roundtrip
)
# Should preserve the malformed JSON string through the round trip
assert tools_roundtrip[0]["function"]["arguments"] == '{"key": "value", "broken'
print("✓ Round-trip conversion preserves malformed JSON")
print("✓ All malformed JSON handling tests passed")

2
tests/proxy_admin_ui_tests/test_key_management.py
View file

@ -1061,7 +1061,6 @@ async def test_list_key_helper(prisma_client):
api_key="sk-1234",
user_id="admin",
),
litellm_changed_by=None,
)
@ -1182,7 +1181,6 @@ async def test_list_key_helper_team_filtering(prisma_client):
api_key="sk-1234",
user_id="admin",
),
litellm_changed_by=None,
)

20
tests/proxy_unit_tests/test_key_generate_prisma.py
View file

@ -1166,10 +1166,8 @@ def test_delete_key_auth(prisma_client):
asyncio.run(test())
except Exception as e:
print("Got Exception", e)
# Handle different exception types - ProxyException has .message, others might have .detail or str(e)
error_message = getattr(e, "message", None) or getattr(e, "detail", None) or str(e)
print(f"Error message: {error_message}")
assert "Authentication Error" in error_message or "Invalid proxy server token" in error_message or "not found in db" in error_message
print(e.message)
assert "Authentication Error" in e.message
pass
@ -2710,12 +2708,7 @@ async def test_reset_spend_authentication(prisma_client):
_response = await new_user(
data=NewUserRequest(
tpm_limit=20,
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key=master_key,
user_id="1234",
),
)
)
generate_key = "Bearer " + _response.key
@ -2735,12 +2728,7 @@ async def test_reset_spend_authentication(prisma_client):
data=NewUserRequest(
user_role=LitellmUserRoles.PROXY_ADMIN,
tpm_limit=20,
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key=master_key,
user_id="1234",
),
)
)
generate_key = "Bearer " + _response.key

42
tests/test_litellm/enterprise/test_responses_background_cost.py → tests/test_litellm/integrations/test_responses_background_cost.py
View file

@ -2,28 +2,14 @@
Integration tests for responses API background cost tracking
"""
import asyncio
import os
import sys
from datetime import datetime
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
# Import litellm first to ensure it's in sys.modules before enterprise imports
import litellm # noqa: E402
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse # noqa: E402
# Now import enterprise modules
try:
from litellm_enterprise.proxy.common_utils.check_responses_cost import ( # noqa: E402
CheckResponsesCost,
)
except ImportError as e:
# Skip all tests in this module if enterprise module is not available
pytest.skip(f"Enterprise module not available: {e}", allow_module_level=True)
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
class TestResponsesBackgroundCostTracking:
@ -298,6 +284,10 @@ class TestCheckResponsesCost:
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
):
"""Test CheckResponsesCost initialization"""
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
CheckResponsesCost,
)
checker = CheckResponsesCost(
proxy_logging_obj=mock_proxy_logging_obj,
prisma_client=mock_prisma_client,
@ -313,6 +303,10 @@ class TestCheckResponsesCost:
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
):
"""Test polling when there are no jobs"""
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
CheckResponsesCost,
)
# Mock find_many to return empty list
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[]
@ -340,6 +334,10 @@ class TestCheckResponsesCost:
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
):
"""Test polling with a completed job"""
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
CheckResponsesCost,
)
# Create a mock job
mock_job = MagicMock()
mock_job.id = "job-123"
@ -393,6 +391,10 @@ class TestCheckResponsesCost:
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
):
"""Test polling with a failed job"""
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
CheckResponsesCost,
)
# Create a mock job
mock_job = MagicMock()
mock_job.id = "job-456"
@ -433,6 +435,10 @@ class TestCheckResponsesCost:
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
):
"""Test polling with a job still in progress"""
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
CheckResponsesCost,
)
# Create a mock job
mock_job = MagicMock()
mock_job.id = "job-789"
@ -473,6 +479,10 @@ class TestCheckResponsesCost:
self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
):
"""Test that errors when querying responses are handled gracefully"""
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
CheckResponsesCost,
)
# Create a mock job
mock_job = MagicMock()
mock_job.id = "job-error"

104
tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py
View file

@ -101,69 +101,55 @@ async def test_bedrock_converse_budget_tokens_preserved():
The bug was that the messages -> completion adapter was converting thinking to reasoning_effort
and losing the original budget_tokens value, causing it to use the default (128) instead.
"""
import os
client = AsyncHTTPHandler()
# Mock at httpx level for better CI compatibility
with patch("httpx.AsyncClient.post") as mock_httpx_post:
with patch.object(client, "post") as mock_post:
mock_response = AsyncMock()
mock_response.status_code = 200
mock_response.headers = {}
mock_response.text = "mock response"
mock_response.json.return_value = {
"output": {
"message": {
"role": "assistant",
"content": [{"text": "4"}]
}
},
"stopReason": "end_turn",
"usage": {
"inputTokens": 10,
"outputTokens": 5,
"totalTokens": 15
with patch.object(client, "post") as mock_post:
mock_response = AsyncMock()
mock_response.status_code = 200
mock_response.headers = {}
mock_response.text = "mock response"
mock_response.json.return_value = {
"output": {
"message": {
"role": "assistant",
"content": [{"text": "4"}]
}
},
"stopReason": "end_turn",
"usage": {
"inputTokens": 10,
"outputTokens": 5,
"totalTokens": 15
}
mock_post.return_value = mock_response
mock_httpx_post.return_value = mock_response
try:
await messages.acreate(
client=client,
max_tokens=1024,
messages=[{"role": "user", "content": "What is 2+2?"}],
model="bedrock/converse/us.anthropic.claude-sonnet-4-20250514-v1:0",
thinking={
"budget_tokens": 1024,
"type": "enabled"
},
)
except Exception:
pass # Expected due to mock response format
# Check which mock was called (client.post or httpx.AsyncClient.post)
if mock_post.call_count == 0 and mock_httpx_post.call_count == 0:
# Skip test if neither mock was called (CI environment issue)
if os.getenv("CI") == "true":
pytest.skip("Mock not intercepted in CI environment")
else:
pytest.fail("Expected mock to be called but it wasn't")
# Use whichever mock was actually called
active_mock = mock_post if mock_post.call_count > 0 else mock_httpx_post
call_kwargs = active_mock.call_args.kwargs
json_data = call_kwargs.get("json") or json.loads(call_kwargs.get("data", "{}"))
print("Request json: ", json.dumps(json_data, indent=4, default=str))
additional_fields = json_data.get("additionalModelRequestFields", {})
thinking_config = additional_fields.get("thinking", {})
assert "thinking" in additional_fields, "thinking parameter should be in additionalModelRequestFields"
assert thinking_config.get("type") == "enabled", "thinking.type should be 'enabled'"
assert thinking_config.get("budget_tokens") == 1024, f"thinking.budget_tokens should be 1024, but got {thinking_config.get('budget_tokens')}"
}
mock_post.return_value = mock_response
try:
await messages.acreate(
client=client,
max_tokens=1024,
messages=[{"role": "user", "content": "What is 2+2?"}],
model="bedrock/converse/us.anthropic.claude-sonnet-4-20250514-v1:0",
thinking={
"budget_tokens": 1024,
"type": "enabled"
},
)
except Exception:
pass # Expected due to mock response format
mock_post.assert_called_once()
call_kwargs = mock_post.call_args.kwargs
json_data = call_kwargs.get("json") or json.loads(call_kwargs.get("data", "{}"))
print("Request json: ", json.dumps(json_data, indent=4, default=str))
additional_fields = json_data.get("additionalModelRequestFields", {})
thinking_config = additional_fields.get("thinking", {})
assert "thinking" in additional_fields, "thinking parameter should be in additionalModelRequestFields"
assert thinking_config.get("type") == "enabled", "thinking.type should be 'enabled'"
assert thinking_config.get("budget_tokens") == 1024, f"thinking.budget_tokens should be 1024, but got {thinking_config.get('budget_tokens')}"
def test_openai_model_with_thinking_converts_to_reasoning_effort():

380
tests/test_litellm/llms/anthropic/test_message_sanitization.py
View file

@ -1,380 +0,0 @@
"""
Test message sanitization for Anthropic API when modify_params=True
Tests three cases:
A. Missing tool_result for tool_use (orphaned tool calls)
B. Orphaned tool_result without matching tool_use
C. Empty text content
"""
import pytest
import sys
import os
# Add the parent directory to the path so we can import litellm
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../..")))
import litellm
from litellm.litellm_core_utils.prompt_templates.factory import (
sanitize_messages_for_tool_calling,
anthropic_messages_pt,
)
class TestMessageSanitization:
"""Test message sanitization for tool calling scenarios"""
def setup_method(self):
"""Setup for each test"""
# Save original modify_params value
self.original_modify_params = litellm.modify_params
litellm.modify_params = True
def teardown_method(self):
"""Cleanup after each test"""
# Restore original modify_params value
litellm.modify_params = self.original_modify_params
def test_case_a_orphaned_tool_call_single(self):
"""
Test Case A: Assistant message with tool_calls but no tool result
Should add a dummy tool result message
"""
messages = [
{
"role": "user",
"content": "What is the weather in Nashik?"
},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "toolu_01Kus2cC3ydjBW7UK4GJqBP4",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "Nashik, India"}'
}
}
]
}
]
sanitized = sanitize_messages_for_tool_calling(messages)
# Should have 3 messages: user, assistant, and dummy tool result
assert len(sanitized) == 3
assert sanitized[0]["role"] == "user"
assert sanitized[1]["role"] == "assistant"
assert sanitized[2]["role"] == "tool"
assert sanitized[2]["tool_call_id"] == "toolu_01Kus2cC3ydjBW7UK4GJqBP4"
assert "skipped" in sanitized[2]["content"].lower() or "interrupted" in sanitized[2]["content"].lower()
assert "get_weather" in sanitized[2]["content"]
def test_case_a_orphaned_tool_call_multiple(self):
"""
Test Case A: Assistant message with multiple tool_calls, some missing results
"""
messages = [
{
"role": "user",
"content": "Get weather for Nashik and Mumbai"
},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "Nashik"}'
}
},
{
"id": "call_2",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "Mumbai"}'
}
}
]
},
{
"role": "tool",
"tool_call_id": "call_1",
"content": "Weather in Nashik: 25°C"
}
]
sanitized = sanitize_messages_for_tool_calling(messages)
# Should have 4 messages: user, assistant, tool result for call_1, dummy for call_2
assert len(sanitized) == 4
assert sanitized[0]["role"] == "user"
assert sanitized[1]["role"] == "assistant"
assert sanitized[2]["tool_call_id"] == "call_2" # Dummy added first
assert sanitized[3]["tool_call_id"] == "call_1" # Original tool result
def test_case_b_orphaned_tool_result(self):
"""
Test Case B: Tool result without matching tool_call in previous assistant message
Should remove the orphaned tool result
"""
messages = [
{
"role": "user",
"content": "Hello"
},
{
"role": "assistant",
"content": "Hi there!"
},
{
"role": "tool",
"tool_call_id": "nonexistent_id",
"content": "Some result"
}
]
sanitized = sanitize_messages_for_tool_calling(messages)
# Should have only 2 messages, orphaned tool result removed
assert len(sanitized) == 2
assert sanitized[0]["role"] == "user"
assert sanitized[1]["role"] == "assistant"
def test_case_b_valid_tool_result_preserved(self):
"""
Test Case B: Valid tool result with matching tool_call should be preserved
"""
messages = [
{
"role": "user",
"content": "What's the weather?"
},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_123",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "Boston"}'
}
}
]
},
{
"role": "tool",
"tool_call_id": "call_123",
"content": "Weather: 20°C"
}
]
sanitized = sanitize_messages_for_tool_calling(messages)
# All messages should be preserved
assert len(sanitized) == 3
assert sanitized[2]["role"] == "tool"
assert sanitized[2]["tool_call_id"] == "call_123"
def test_case_c_empty_text_content_user(self):
"""
Test Case C: Empty text content in user message
Should replace with placeholder
"""
messages = [
{
"role": "user",
"content": ""
},
{
"role": "assistant",
"content": "Hello!"
}
]
sanitized = sanitize_messages_for_tool_calling(messages)
assert len(sanitized) == 2
assert sanitized[0]["role"] == "user"
assert sanitized[0]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
def test_case_c_whitespace_only_content(self):
"""
Test Case C: Whitespace-only content
Should replace with placeholder
"""
messages = [
{
"role": "user",
"content": " \n \t "
},
{
"role": "assistant",
"content": " "
}
]
sanitized = sanitize_messages_for_tool_calling(messages)
assert len(sanitized) == 2
assert sanitized[0]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
assert sanitized[1]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
def test_case_c_valid_content_preserved(self):
"""
Test Case C: Valid non-empty content should be preserved
"""
messages = [
{
"role": "user",
"content": "Hello"
},
{
"role": "assistant",
"content": "Hi there!"
}
]
sanitized = sanitize_messages_for_tool_calling(messages)
assert len(sanitized) == 2
assert sanitized[0]["content"] == "Hello"
assert sanitized[1]["content"] == "Hi there!"
def test_combined_cases(self):
"""
Test combination of multiple cases
"""
messages = [
{
"role": "user",
"content": "Get weather"
},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "NYC"}'
}
}
]
},
# Missing tool result for call_1
{
"role": "user",
"content": "" # Empty content
},
{
"role": "assistant",
"content": "Response"
},
{
"role": "tool",
"tool_call_id": "orphaned_id", # Orphaned tool result
"content": "Some data"
}
]
sanitized = sanitize_messages_for_tool_calling(messages)
# Should have: user, assistant, dummy tool result, user (sanitized), assistant
# Orphaned tool result should be removed
assert len(sanitized) == 5
assert sanitized[0]["role"] == "user"
assert sanitized[1]["role"] == "assistant"
assert sanitized[2]["role"] == "tool"
assert sanitized[2]["tool_call_id"] == "call_1" # Dummy added
assert sanitized[3]["role"] == "user"
assert sanitized[3]["content"] == "[System: Empty message content sanitised to satisfy protocol]"
assert sanitized[4]["role"] == "assistant"
def test_modify_params_false_no_sanitization(self):
"""
Test that sanitization is skipped when modify_params=False
"""
litellm.modify_params = False
messages = [
{
"role": "user",
"content": ""
},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_1",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{}'
}
}
]
}
]
sanitized = sanitize_messages_for_tool_calling(messages)
# Messages should be unchanged
assert len(sanitized) == 2
assert sanitized[0]["content"] == ""
assert len(sanitized[1].get("tool_calls", [])) == 1
def test_anthropic_messages_pt_integration(self):
"""
Test that sanitization is integrated into anthropic_messages_pt
"""
litellm.modify_params = True
messages = [
{
"role": "user",
"content": "What is the weather in Nashik?"
},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "toolu_01Kus2cC3ydjBW7UK4GJqBP4",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "Nashik, India"}'
}
}
]
}
]
# This should not raise an error and should add dummy tool result
result = anthropic_messages_pt(
messages=messages,
model="claude-sonnet-4-5",
llm_provider="anthropic"
)
# Should have at least 2 messages (user and assistant)
# The tool result will be merged into user content
assert len(result) >= 2
assert result[0]["role"] == "user"
assert result[1]["role"] == "assistant"
if __name__ == "__main__":
pytest.main([__file__, "-v"])

93
tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py
View file

@ -2610,6 +2610,99 @@ def test_request_metadata_not_provided():
assert "requestMetadata" not in request_data
def test_empty_assistant_message_handling():
"""
Test that empty assistant messages are handled correctly by replacing
empty or whitespace-only content with a placeholder to prevent AWS Bedrock
Converse API 400 Bad Request errors.
"""
from litellm.litellm_core_utils.prompt_templates.factory import (
_bedrock_converse_messages_pt,
)
# Test case 1: Empty string content - test with modify_params=True to prevent merging
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": ""}, # Empty content
{"role": "user", "content": "How are you?"}
]
# Enable modify_params to prevent consecutive user message merging
original_modify_params = litellm.modify_params
litellm.modify_params = True
try:
result = _bedrock_converse_messages_pt(
messages=messages,
model="anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
# Should have 3 messages: user, assistant (with placeholder), user
assert len(result) == 3
assert result[0]["role"] == "user"
assert result[1]["role"] == "assistant"
assert result[2]["role"] == "user"
# Assistant message should have placeholder text instead of empty content
assert len(result[1]["content"]) == 1
assert result[1]["content"][0]["text"] == "Please continue."
# Test case 2: Whitespace-only content
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": " "}, # Whitespace-only content
{"role": "user", "content": "How are you?"}
]
result = _bedrock_converse_messages_pt(
messages=messages,
model="anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
# Assistant message should have placeholder text instead of whitespace
assert len(result[1]["content"]) == 1
assert result[1]["content"][0]["text"] == "Please continue."
# Test case 3: Empty list content
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": [{"type": "text", "text": ""}]}, # Empty text in list
{"role": "user", "content": "How are you?"}
]
result = _bedrock_converse_messages_pt(
messages=messages,
model="anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
# Assistant message should have placeholder text instead of empty text
assert len(result[1]["content"]) == 1
assert result[1]["content"][0]["text"] == "Please continue."
# Test case 4: Normal content should not be affected
messages = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "I'm doing well, thank you!"}, # Normal content
{"role": "user", "content": "How are you?"}
]
result = _bedrock_converse_messages_pt(
messages=messages,
model="anthropic.claude-3-5-sonnet-20240620-v1:0",
llm_provider="bedrock_converse"
)
# Assistant message should keep original content
assert len(result[1]["content"]) == 1
assert result[1]["content"][0]["text"] == "I'm doing well, thank you!"
finally:
# Restore original modify_params setting
litellm.modify_params = original_modify_params
def test_is_nova_lite_2_model():
"""Test the _is_nova_lite_2_model() method for detecting Nova 2 models."""

146
tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py
View file

@ -21,51 +21,43 @@ class TestBedrockFilesIntegration:
file_id = "s3://test-bucket/test-file.jsonl"
expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}'
# Mock AWS credentials
with patch.dict(
"os.environ",
{
"AWS_ACCESS_KEY_ID": "test-access-key",
"AWS_SECRET_ACCESS_KEY": "test-secret-key",
},
):
# Mock the bedrock_files_instance.file_content method
with patch(
"litellm.files.main.bedrock_files_instance.file_content",
new_callable=AsyncMock,
) as mock_file_content:
# Create a mock HttpxBinaryResponseContent response
import httpx
# Mock the bedrock_files_instance.file_content method
with patch(
"litellm.files.main.bedrock_files_instance.file_content",
new_callable=AsyncMock,
) as mock_file_content:
# Create a mock HttpxBinaryResponseContent response
import httpx
mock_response = httpx.Response(
status_code=200,
content=expected_content,
headers={"content-type": "application/octet-stream"},
request=httpx.Request(
method="GET", url="s3://test-bucket/test-file.jsonl"
),
)
mock_file_content.return_value = HttpxBinaryResponseContent(
response=mock_response
)
mock_response = httpx.Response(
status_code=200,
content=expected_content,
headers={"content-type": "application/octet-stream"},
request=httpx.Request(
method="GET", url="s3://test-bucket/test-file.jsonl"
),
)
mock_file_content.return_value = HttpxBinaryResponseContent(
response=mock_response
)
# Call litellm.afile_content
result = await litellm.afile_content(
file_id=file_id,
custom_llm_provider="bedrock",
aws_region_name="us-west-2",
)
# Call litellm.afile_content
result = await litellm.afile_content(
file_id=file_id,
custom_llm_provider="bedrock",
aws_region_name="us-west-2",
)
# Verify the result
assert isinstance(result, HttpxBinaryResponseContent)
assert result.response.content == expected_content
assert result.response.status_code == 200
# Verify the result
assert isinstance(result, HttpxBinaryResponseContent)
assert result.response.content == expected_content
assert result.response.status_code == 200
# Verify the mock was called with correct parameters
mock_file_content.assert_called_once()
call_kwargs = mock_file_content.call_args.kwargs
assert call_kwargs["_is_async"] is True
assert call_kwargs["file_content_request"]["file_id"] == file_id
# Verify the mock was called with correct parameters
mock_file_content.assert_called_once()
call_kwargs = mock_file_content.call_args.kwargs
assert call_kwargs["_is_async"] is True
assert call_kwargs["file_content_request"]["file_id"] == file_id
@pytest.mark.asyncio
async def test_litellm_afile_content_bedrock_provider_with_unified_file_id(self):
@ -80,47 +72,39 @@ class TestBedrockFilesIntegration:
expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}'
# Mock AWS credentials
with patch.dict(
"os.environ",
{
"AWS_ACCESS_KEY_ID": "test-access-key",
"AWS_SECRET_ACCESS_KEY": "test-secret-key",
},
):
# Mock the bedrock_files_instance.file_content method
with patch(
"litellm.files.main.bedrock_files_instance.file_content",
new_callable=AsyncMock,
) as mock_file_content:
# Create a mock HttpxBinaryResponseContent response
import httpx
# Mock the bedrock_files_instance.file_content method
with patch(
"litellm.files.main.bedrock_files_instance.file_content",
new_callable=AsyncMock,
) as mock_file_content:
# Create a mock HttpxBinaryResponseContent response
import httpx
mock_response = httpx.Response(
status_code=200,
content=expected_content,
headers={"content-type": "application/octet-stream"},
request=httpx.Request(method="GET", url=s3_uri),
)
mock_file_content.return_value = HttpxBinaryResponseContent(
response=mock_response
)
mock_response = httpx.Response(
status_code=200,
content=expected_content,
headers={"content-type": "application/octet-stream"},
request=httpx.Request(method="GET", url=s3_uri),
)
mock_file_content.return_value = HttpxBinaryResponseContent(
response=mock_response
)
# Call litellm.afile_content with unified file ID
result = await litellm.afile_content(
file_id=encoded_file_id,
custom_llm_provider="bedrock",
aws_region_name="us-west-2",
)
# Call litellm.afile_content with unified file ID
result = await litellm.afile_content(
file_id=encoded_file_id,
custom_llm_provider="bedrock",
aws_region_name="us-west-2",
)
# Verify the result
assert isinstance(result, HttpxBinaryResponseContent)
assert result.response.content == expected_content
assert result.response.status_code == 200
# Verify the result
assert isinstance(result, HttpxBinaryResponseContent)
assert result.response.content == expected_content
assert result.response.status_code == 200
# Verify the mock was called - the handler should extract S3 URI from unified file ID
mock_file_content.assert_called_once()
call_kwargs = mock_file_content.call_args.kwargs
assert call_kwargs["_is_async"] is True
# The handler extracts S3 URI from the unified file ID
assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id
# Verify the mock was called - the handler should extract S3 URI from unified file ID
mock_file_content.assert_called_once()
call_kwargs = mock_file_content.call_args.kwargs
assert call_kwargs["_is_async"] is True
# The handler extracts S3 URI from the unified file ID
assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id

34
tests/test_litellm/llms/huggingface/embedding/test_handler.py
View file

@ -41,12 +41,8 @@ def mock_embedding_async_http_handler():
class TestHuggingFaceEmbedding:
@pytest.fixture(autouse=True)
def setup(self, mock_embedding_http_handler, mock_embedding_async_http_handler):
# Mock both sync and async versions of get_hf_task functions
self.mock_get_task_patcher = patch("litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model")
self.mock_get_task_async_patcher = patch("litellm.llms.huggingface.embedding.handler.async_get_hf_task_embedding_for_model", new_callable=AsyncMock)
self.mock_get_task = self.mock_get_task_patcher.start()
self.mock_get_task_async = self.mock_get_task_async_patcher.start()
def mock_get_task_side_effect(model, task_type, api_base):
if task_type is not None:
@ -54,7 +50,6 @@ class TestHuggingFaceEmbedding:
return "sentence-similarity"
self.mock_get_task.side_effect = mock_get_task_side_effect
self.mock_get_task_async.side_effect = mock_get_task_side_effect
self.model = "huggingface/BAAI/bge-m3"
self.mock_http = mock_embedding_http_handler
@ -64,7 +59,6 @@ class TestHuggingFaceEmbedding:
yield
self.mock_get_task_patcher.stop()
self.mock_get_task_async_patcher.stop()
def test_input_type_preserved_in_optional_params(self):
input_text = ["hello world"]
@ -87,3 +81,31 @@ class TestHuggingFaceEmbedding:
# Should NOT have sentence-similarity format
assert "source_sentence" not in str(request_data)
assert "sentences" not in str(request_data)
def test_embedding_with_sentence_similarity_task(self):
"""Test embedding when task type is sentence-similarity (requires 2+ sentences)"""
similarity_response = {
"similarities": [[0, 0.9], [1, 0.8]]
}
self.mock_http.return_value.json.return_value = similarity_response
# Test with 2+ sentences (required for sentence-similarity)
input_text = ["This is the source sentence", "This is sentence one", "This is sentence two"]
response = litellm.embedding(
model=self.model,
input=input_text,
# Use the model's natural task type (sentence-similarity)
)
self.mock_http.assert_called_once()
post_call_args = self.mock_http.call_args
request_data = json.loads(post_call_args[1]["data"])
assert "inputs" in request_data
assert "source_sentence" in request_data["inputs"]
assert "sentences" in request_data["inputs"]
assert request_data["inputs"]["source_sentence"] == input_text[0]
assert request_data["inputs"]["sentences"] == input_text[1:]

8
tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
View file

@ -187,9 +187,9 @@ class TestContextCachingEndpoints:
assert returned_params == optional_params
assert returned_cache == "existing_cache_name"
# Verify cache key was generated with tools
# Verify cache key was generated with tools and model
mock_cache_obj.get_cache_key.assert_called_once_with(
messages=cached_messages, tools=self.sample_tools
messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro"
)
@pytest.mark.parametrize(
@ -460,9 +460,9 @@ class TestContextCachingEndpoints:
assert returned_params == optional_params
assert returned_cache == "existing_cache_name"
# Verify cache key was generated with tools
# Verify cache key was generated with tools and model
mock_cache_obj.get_cache_key.assert_called_once_with(
messages=cached_messages, tools=self.sample_tools
messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro"
)
@pytest.mark.asyncio

46
tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py
View file

@ -12,7 +12,53 @@ from litellm.types.llms.openai import HttpxBinaryResponseContent
class TestVertexAIFilesIntegration:
"""Test integration of Vertex AI files with main litellm API"""
@pytest.mark.asyncio
async def test_litellm_afile_content_vertex_ai_provider(self):
"""Test litellm.afile_content with vertex_ai provider"""
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
expected_content = b"test file content"
# Mock the vertex_ai_files_instance.file_content method
with patch(
"litellm.files.main.vertex_ai_files_instance.file_content",
new_callable=AsyncMock,
) as mock_file_content:
# Create a mock HttpxBinaryResponseContent response
import httpx
mock_response = httpx.Response(
status_code=200,
content=expected_content,
headers={"content-type": "application/octet-stream"},
request=httpx.Request(
method="GET", url="gs://test-bucket/test-file.txt"
),
)
mock_file_content.return_value = HttpxBinaryResponseContent(
response=mock_response
)
# Call litellm.afile_content
result = await litellm.afile_content(
file_id=file_id,
custom_llm_provider="vertex_ai",
vertex_project="test-project",
vertex_location="us-central1",
vertex_credentials=None,
)
# Verify the result
assert isinstance(result, HttpxBinaryResponseContent)
assert result.response.content == expected_content
assert result.response.status_code == 200
# Verify the mock was called with correct parameters
mock_file_content.assert_called_once()
call_kwargs = mock_file_content.call_args.kwargs
assert call_kwargs["_is_async"] is True
assert call_kwargs["file_content_request"]["file_id"] == file_id
assert call_kwargs["vertex_project"] == "test-project"
assert call_kwargs["vertex_location"] == "us-central1"
def test_litellm_file_content_vertex_ai_provider(self):
"""Test litellm.file_content with vertex_ai provider (sync)"""

34
tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py
View file

@ -75,6 +75,40 @@ class TestCreateToolFunction:
call_args[0][0]
)
@pytest.mark.asyncio
async def test_leading_digit_parameter(self):
"""Test function with parameter starting with digit (e.g., 2fa-code)."""
operation = {
"parameters": [
{
"name": "2fa-code",
"in": "query",
"required": False,
"schema": {"type": "string"},
}
]
}
func = create_tool_function(
path="/verify",
method="post",
operation=operation,
base_url="https://api.example.com",
)
assert callable(func)
with patch(GET_ASYNC_CLIENT_TARGET) as mock_client:
async_client = _create_mock_client("post", "verified")
mock_client.return_value = async_client
result = await func(**{"2fa-code": "123456"})
assert result == "verified"
# Verify query parameter was included
call_args = async_client.post.call_args
assert call_args[1]["params"]["2fa-code"] == "123456"
@pytest.mark.asyncio
async def test_dot_in_parameter_name(self):
"""Test function with dot in parameter name (e.g., user.name)."""

29
tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py
View file

@ -8,7 +8,7 @@ and following LiteLLM testing patterns and best practices.
# Standard library imports
import os
import sys
from typing import Dict, Any
from typing import Dict
from unittest.mock import Mock, patch
# Add parent directory to path for imports
@ -43,6 +43,33 @@ from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
# ============================================================================
@pytest.fixture(scope="function", autouse=True)
def setup_and_teardown():
"""
Standard LiteLLM fixture that reloads litellm before every function
to speed up testing by removing callbacks being chained.
"""
import importlib
import asyncio
# Reload litellm to ensure clean state
importlib.reload(litellm)
# Set up async loop
loop = asyncio.get_event_loop_policy().new_event_loop()
asyncio.set_event_loop(loop)
# Set up litellm state
litellm.set_verbose = True
litellm.guardrail_name_config_map = {}
yield
# Teardown
loop.close()
asyncio.set_event_loop(None)
@pytest.fixture
def env_setup(monkeypatch):
"""Fixture to set up environment variables for testing."""

83
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -12,7 +12,6 @@ sys.path.insert(
from litellm.proxy._types import (
LiteLLM_UserTableFiltered,
LitellmUserRoles,
NewUserRequest,
ProxyException,
UpdateUserRequest,
@ -307,88 +306,6 @@ async def test_new_user_license_over_limit(mocker):
mock_license_check.is_over_limit.assert_called_once_with(total_users=1000)
@pytest.mark.asyncio
async def test_new_user_non_admin_cannot_create_admin(mocker):
"""
Test that non-admin users cannot create administrative users (PROXY_ADMIN or PROXY_ADMIN_VIEW_ONLY).
This prevents privilege escalation vulnerabilities.
"""
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
# Mock the prisma client
mock_prisma_client = mocker.MagicMock()
# Setup the mock count response (under license limit)
async def mock_count(*args, **kwargs):
return 5 # Low user count, under limit
mock_prisma_client.db.litellm_usertable.count = mock_count
# Mock duplicate checks to pass
async def mock_check_duplicate_user_email(*args, **kwargs):
return None # No duplicate found
async def mock_check_duplicate_user_id(*args, **kwargs):
return None # No duplicate found
mocker.patch(
"litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_email",
mock_check_duplicate_user_email,
)
mocker.patch(
"litellm.proxy.management_endpoints.internal_user_endpoints._check_duplicate_user_id",
mock_check_duplicate_user_id,
)
# Mock the license check to return False (under limit)
mock_license_check = mocker.MagicMock()
mock_license_check.is_over_limit.return_value = False
# Patch the imports in the endpoint
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch("litellm.proxy.proxy_server._license_check", mock_license_check)
# Test Case 1: INTERNAL_USER trying to create PROXY_ADMIN
user_request = NewUserRequest(
user_email="admin@example.com", user_role=LitellmUserRoles.PROXY_ADMIN
)
# Mock user_api_key_dict with non-admin role
mock_user_api_key_dict = UserAPIKeyAuth(
user_id="test_internal_user", user_role=LitellmUserRoles.INTERNAL_USER
)
# Call new_user function and expect ProxyException
with pytest.raises(ProxyException) as exc_info:
await new_user(data=user_request, user_api_key_dict=mock_user_api_key_dict)
# Verify the exception details
assert exc_info.value.code == 403 or exc_info.value.code == "403"
assert "Only proxy admins can create administrative users" in str(exc_info.value.message)
assert "proxy_admin" in str(exc_info.value.message)
assert "proxy_admin_viewer" in str(exc_info.value.message)
assert str(LitellmUserRoles.PROXY_ADMIN) in str(exc_info.value.message)
assert str(LitellmUserRoles.INTERNAL_USER) in str(exc_info.value.message)
# Test Case 2: INTERNAL_USER trying to create PROXY_ADMIN_VIEW_ONLY
user_request_viewer = NewUserRequest(
user_email="admin_viewer@example.com",
user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
)
with pytest.raises(ProxyException) as exc_info2:
await new_user(
data=user_request_viewer, user_api_key_dict=mock_user_api_key_dict
)
# Verify the exception details
assert exc_info2.value.code == 403 or exc_info2.value.code == "403"
assert "Only proxy admins can create administrative users" in str(
exc_info2.value.message
)
assert str(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) in str(exc_info2.value.message)
@pytest.mark.asyncio
async def test_user_info_url_encoding_plus_character(mocker):
"""

416
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -31,13 +31,9 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
_check_team_key_limits,
_common_key_generation_helper,
_list_key_helper,
_persist_deleted_verification_tokens,
_save_deleted_verification_token_records,
_transform_verification_tokens_to_deleted_records,
can_modify_verification_token,
check_org_key_model_specific_limits,
check_team_key_model_specific_limits,
delete_verification_tokens,
generate_key_helper_fn,
prepare_key_update_data,
validate_key_team_change,
@ -2732,364 +2728,64 @@ def test_check_org_key_model_specific_limits_org_model_tpm_overallocation():
)
def test_transform_verification_tokens_to_deleted_records():
from datetime import datetime, timezone
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
key1 = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id="team-456",
key_alias="test-key-1",
spend=100.0,
max_budget=1000.0,
models=["gpt-4"],
aliases={},
config={},
permissions={},
metadata={"test": "value"},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
key2 = LiteLLM_VerificationToken(
token="hashed-token-2",
user_id="user-789",
team_id=None,
key_alias="test-key-2",
spend=50.0,
max_budget=500.0,
models=["gpt-3.5-turbo"],
aliases={"alias": "model"},
config={"config": "value"},
permissions={"permission": True},
metadata={},
model_max_budget={"gpt-4": {"budget_limit": 100.0}},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
records = _transform_verification_tokens_to_deleted_records(
keys=[key1, key2],
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
assert len(records) == 2
assert all("deleted_at" in record for record in records)
assert all("deleted_by" in record for record in records)
assert all("deleted_by_api_key" in record for record in records)
assert all("litellm_changed_by" in record for record in records)
assert all(record["deleted_by"] == "user-123" for record in records)
assert all(record["deleted_by_api_key"] == user_api_key_dict.api_key for record in records)
assert all(record["litellm_changed_by"] == "admin-user" for record in records)
record1 = records[0]
assert record1["token"] == "hashed-token-1"
assert record1["user_id"] == "user-123"
assert record1["team_id"] == "team-456"
assert isinstance(record1["aliases"], str)
assert isinstance(record1["config"], str)
assert isinstance(record1["permissions"], str)
assert isinstance(record1["metadata"], str)
assert "litellm_budget_table" not in record1
assert "litellm_organization_table" not in record1
assert "object_permission" not in record1
assert "id" not in record1
record2 = records[1]
assert record2["token"] == "hashed-token-2"
assert isinstance(record2["model_max_budget"], str)
def test_transform_verification_tokens_to_deleted_records_empty_list():
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
records = _transform_verification_tokens_to_deleted_records(
keys=[],
user_api_key_dict=user_api_key_dict,
)
assert records == []
@pytest.mark.asyncio
async def test_save_deleted_verification_token_records():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many
)
records = [
{
"token": "hashed-token-1",
"user_id": "user-123",
"deleted_at": "2024-01-01T00:00:00Z",
"deleted_by": "admin",
},
{
"token": "hashed-token-2",
"user_id": "user-456",
"deleted_at": "2024-01-01T00:00:00Z",
"deleted_by": "admin",
},
]
await _save_deleted_verification_token_records(
records=records, prisma_client=mock_prisma_client
)
mock_create_many.assert_called_once_with(data=records)
@pytest.mark.asyncio
async def test_save_deleted_verification_token_records_empty_list():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many
)
await _save_deleted_verification_token_records(
records=[], prisma_client=mock_prisma_client
)
mock_create_many.assert_not_called()
@pytest.mark.asyncio
async def test_persist_deleted_verification_tokens():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many
)
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
key = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id="team-456",
key_alias="test-key",
spend=100.0,
max_budget=1000.0,
models=["gpt-4"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
await _persist_deleted_verification_tokens(
keys=[key],
prisma_client=mock_prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
mock_create_many.assert_called_once()
call_args = mock_create_many.call_args
assert "data" in call_args.kwargs
records = call_args.kwargs["data"]
assert len(records) == 1
assert records[0]["token"] == "hashed-token-1"
assert records[0]["deleted_by"] == "user-123"
assert records[0]["litellm_changed_by"] == "admin-user"
@pytest.mark.asyncio
async def test_delete_verification_tokens_persists_deleted_keys(monkeypatch):
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
key1 = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id="team-456",
key_alias="test-key-1",
spend=100.0,
max_budget=1000.0,
models=["gpt-4"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
key2 = LiteLLM_VerificationToken(
token="hashed-token-2",
user_id="user-789",
team_id=None,
key_alias="test-key-2",
spend=50.0,
max_budget=500.0,
models=["gpt-3.5-turbo"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
mock_find_many = AsyncMock(return_value=[key1, key2])
mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many
# delete_data returns {"deleted_keys": ...} from utils.py line 3049
# The function at line 2410 assigns it to deleted_tokens
# Then at line 2444 returns {"deleted_keys": deleted_tokens}
# So if delete_data returns {"deleted_keys": list}, then result would be nested
# But looking at the error, it seems like delete_data might return just the list
# Or the code extracts it. Let's return the list directly since that's what the test expects
mock_delete_data = AsyncMock(return_value=["hashed-token-1", "hashed-token-2"])
mock_prisma_client.delete_data = mock_delete_data
# Mock cache delete_cache method
mock_user_api_key_cache.delete_cache = MagicMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many
)
def mock_hash_token(token):
return token if not token.startswith("sk-") else f"hashed-{token}"
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed",
mock_hash_token,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.hash_token",
mock_hash_token,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
)
result, deleted_keys = await delete_verification_tokens(
tokens=["sk-token-1", "sk-token-2"],
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
mock_create_many.assert_called_once()
call_args = mock_create_many.call_args
assert "data" in call_args.kwargs
records = call_args.kwargs["data"]
assert len(records) == 2
assert all(record["deleted_by"] == "admin-user" for record in records)
assert all(record["litellm_changed_by"] == "admin-user" for record in records)
# delete_data returns the list directly, which gets wrapped in {"deleted_keys": ...}
assert isinstance(result["deleted_keys"], list)
assert set(result["deleted_keys"]) == {"hashed-token-1", "hashed-token-2"}
assert len(deleted_keys) == 2
@pytest.mark.asyncio
async def test_delete_key_fn_persists_deleted_keys(monkeypatch):
from litellm.proxy._types import KeyRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_key_fn,
delete_verification_tokens,
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
key1 = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id="team-456",
key_alias="test-key-1",
spend=100.0,
max_budget=1000.0,
models=["gpt-4"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
model_spend={},
soft_budget_cooldown=False,
allowed_routes=[],
)
async def mock_delete_verification_tokens(*args, **kwargs):
return ({"deleted_keys": ["sk-token-1"]}, [key1])
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.delete_verification_tokens",
mock_delete_verification_tokens,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_api_key_cache",
mock_user_api_key_cache,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_deleted_hook",
AsyncMock(),
)
data = KeyRequest(keys=["sk-token-1"])
result = await delete_key_fn(
data=data,
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
assert result["deleted_keys"] == ["sk-token-1"]
@pytest.mark.asyncio
async def test_can_delete_verification_token_proxy_admin_team_key(monkeypatch):
"""Test that proxy admin can delete any team key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id="test-team-123",
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
api_key="sk-admin",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_delete_verification_token_proxy_admin_personal_key(monkeypatch):
"""Test that proxy admin can delete any personal key."""
key_info = LiteLLM_VerificationToken(
token="test-token",
user_id="other-user",
team_id=None,
)
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
user_id="admin-user",
api_key="sk-admin",
)
mock_prisma_client = AsyncMock()
mock_user_api_key_cache = MagicMock()
result = await can_modify_verification_token(
key_info=key_info,
user_api_key_cache=mock_user_api_key_cache,
user_api_key_dict=user_api_key_dict,
prisma_client=mock_prisma_client,
)
assert result is True
@pytest.mark.asyncio
async def test_can_delete_verification_token_team_admin_own_team(monkeypatch):
"""Test that team admin can delete team keys from their own team."""
key_info = LiteLLM_VerificationToken(
token="test-token",

712
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

@ -20,6 +20,7 @@ from litellm.proxy._types import (
LiteLLM_OrganizationTable,
LiteLLM_OrganizationTableWithMembers,
LiteLLM_TeamTable,
LiteLLM_UserTable,
LitellmUserRoles,
Member,
ProxyErrorTypes,
@ -32,13 +33,8 @@ from litellm.proxy.management_endpoints.team_endpoints import (
from litellm.proxy.management_endpoints.team_endpoints import (
GetTeamMemberPermissionsResponse,
UpdateTeamMemberPermissionsRequest,
_persist_deleted_team_records,
_save_deleted_team_records,
_transform_teams_to_deleted_records,
delete_team,
router,
team_member_add_duplication_check,
team_member_delete,
validate_team_org_change,
)
from litellm.proxy.management_helpers.team_member_permission_checks import (
@ -2264,7 +2260,6 @@ async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_a
# Verification token deletion should be called
mock_db_client.db.litellm_verificationtoken = MagicMock()
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock())
# Execute
@ -2312,7 +2307,6 @@ async def test_team_member_delete_cleans_verification_tokens(mock_db_client, moc
mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(return_value=MagicMock())
mock_db_client.db.litellm_verificationtoken = MagicMock()
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock())
await team_member_delete(
@ -4331,348 +4325,6 @@ async def test_update_team_guardrails_with_org_id():
assert first_call_kwargs["include"]["teams"] is True
def test_transform_teams_to_deleted_records():
from datetime import datetime, timezone
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
team1 = LiteLLM_TeamTable(
team_id="team-1",
team_alias="test-team-1",
members_with_roles=[
Member(user_id="user-1", role="admin"),
Member(user_id="user-2", role="user"),
],
metadata={"test": "value"},
model_max_budget={},
model_spend={},
)
team2 = LiteLLM_TeamTable(
team_id="team-2",
team_alias="test-team-2",
members_with_roles=[],
metadata=None,
model_max_budget={"gpt-4": {"budget_limit": 100.0}},
model_spend={},
)
records = _transform_teams_to_deleted_records(
teams=[team1, team2],
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
assert len(records) == 2
assert all("deleted_at" in record for record in records)
assert all("deleted_by" in record for record in records)
assert all("deleted_by_api_key" in record for record in records)
assert all("litellm_changed_by" in record for record in records)
assert all(record["deleted_by"] == "user-123" for record in records)
# UserAPIKeyAuth hashes the api_key, so we check against the hashed value
assert all(record["deleted_by_api_key"] == user_api_key_dict.api_key for record in records)
assert all(record["litellm_changed_by"] == "admin-user" for record in records)
record1 = records[0]
assert record1["team_id"] == "team-1"
assert isinstance(record1["members_with_roles"], str)
assert isinstance(record1["metadata"], str)
assert "litellm_model_table" not in record1
assert "object_permission" not in record1
assert "id" not in record1
record2 = records[1]
assert record2["team_id"] == "team-2"
# model_max_budget should be converted to JSON string if it exists
if "model_max_budget" in record2:
assert isinstance(record2["model_max_budget"], str)
def test_transform_teams_to_deleted_records_empty_list():
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
records = _transform_teams_to_deleted_records(
teams=[],
user_api_key_dict=user_api_key_dict,
)
assert records == []
@pytest.mark.asyncio
async def test_save_deleted_team_records():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many
records = [
{
"team_id": "team-1",
"team_alias": "test-team-1",
"deleted_at": "2024-01-01T00:00:00Z",
"deleted_by": "admin",
},
{
"team_id": "team-2",
"team_alias": "test-team-2",
"deleted_at": "2024-01-01T00:00:00Z",
"deleted_by": "admin",
},
]
await _save_deleted_team_records(records=records, prisma_client=mock_prisma_client)
mock_create_many.assert_called_once_with(data=records)
@pytest.mark.asyncio
async def test_save_deleted_team_records_empty_list():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many
await _save_deleted_team_records(records=[], prisma_client=mock_prisma_client)
mock_create_many.assert_not_called()
@pytest.mark.asyncio
async def test_persist_deleted_team_records():
mock_prisma_client = AsyncMock()
mock_create_many = AsyncMock()
mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many
user_api_key_dict = UserAPIKeyAuth(
user_id="user-123",
api_key="sk-test",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
team = LiteLLM_TeamTable(
team_id="team-1",
team_alias="test-team",
members_with_roles=[
Member(user_id="user-1", role="admin"),
],
metadata={},
model_max_budget={},
model_spend={},
)
await _persist_deleted_team_records(
teams=[team],
prisma_client=mock_prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_changed_by="admin-user",
)
mock_create_many.assert_called_once()
call_args = mock_create_many.call_args
assert "data" in call_args.kwargs
records = call_args.kwargs["data"]
assert len(records) == 1
assert records[0]["team_id"] == "team-1"
assert records[0]["deleted_by"] == "user-123"
assert records[0]["litellm_changed_by"] == "admin-user"
@pytest.mark.asyncio
async def test_delete_team_persists_deleted_teams(monkeypatch):
from litellm.proxy._types import DeleteTeamRequest
mock_prisma_client = AsyncMock()
mock_user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
team1 = LiteLLM_TeamTable(
team_id="team-1",
team_alias="test-team-1",
members_with_roles=[
Member(user_id="user-1", role="admin"),
],
metadata={},
model_max_budget={},
model_spend={},
)
mock_find_unique = AsyncMock(return_value=team1)
mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique
mock_delete_data = AsyncMock(return_value={"deleted_teams": ["team-1"]})
mock_prisma_client.delete_data = mock_delete_data
mock_create_many_teams = AsyncMock()
mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many_teams
mock_create_many_keys = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many_keys
)
mock_find_many_keys = AsyncMock(return_value=[])
mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.create_audit_log_for_update",
AsyncMock(),
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.litellm_proxy_admin_name",
"admin",
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.team_endpoints.team_member_delete",
AsyncMock(return_value=team1),
)
data = DeleteTeamRequest(team_ids=["team-1"])
result = await delete_team(
data=data,
http_request=MagicMock(),
user_api_key_dict=mock_user_api_key_dict,
litellm_changed_by="admin-user",
)
mock_create_many_teams.assert_called_once()
call_args = mock_create_many_teams.call_args
assert "data" in call_args.kwargs
records = call_args.kwargs["data"]
assert len(records) == 1
assert records[0]["team_id"] == "team-1"
assert records[0]["deleted_by"] == "admin-user"
assert records[0]["litellm_changed_by"] == "admin-user"
@pytest.mark.asyncio
async def test_team_member_delete_persists_deleted_keys(monkeypatch):
from litellm.proxy._types import TeamMemberDeleteRequest
from litellm.proxy.management_endpoints.key_management_endpoints import (
LiteLLM_VerificationToken,
)
mock_prisma_client = AsyncMock()
mock_user_api_key_dict = UserAPIKeyAuth(
user_id="admin-user",
api_key="sk-admin",
user_role=LitellmUserRoles.PROXY_ADMIN.value,
)
team = LiteLLM_TeamTable(
team_id="team-1",
team_alias="test-team",
members_with_roles=[
Member(user_id="user-123", role="admin"),
],
metadata={},
model_max_budget={},
model_spend={},
)
key1 = LiteLLM_VerificationToken(
token="hashed-token-1",
user_id="user-123",
team_id="team-1",
key_alias="test-key-1",
spend=100.0,
max_budget=1000.0,
models=["gpt-4"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
)
key2 = LiteLLM_VerificationToken(
token="hashed-token-2",
user_id="user-123",
team_id="team-1",
key_alias="test-key-2",
spend=50.0,
max_budget=500.0,
models=["gpt-3.5-turbo"],
aliases={},
config={},
permissions={},
metadata={},
model_max_budget={},
)
mock_find_unique_team = AsyncMock(return_value=team)
mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique_team
mock_find_many_user = AsyncMock(
return_value=[
MagicMock(
user_id="user-123",
teams=["team-1"],
model_dump=lambda: {"user_id": "user-123", "teams": ["team-1"]},
)
]
)
mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_user
mock_update_team = AsyncMock()
mock_prisma_client.db.litellm_teamtable.update = mock_update_team
mock_update_user = AsyncMock()
mock_prisma_client.db.litellm_usertable.update = mock_update_user
mock_delete_membership = AsyncMock()
mock_prisma_client.db.litellm_teammembership.delete_many = mock_delete_membership
mock_find_many_keys = AsyncMock(return_value=[key1, key2])
mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys
mock_delete_keys = AsyncMock()
mock_prisma_client.db.litellm_verificationtoken.delete_many = mock_delete_keys
mock_create_many_keys = AsyncMock()
mock_prisma_client.db.litellm_deletedverificationtoken.create_many = (
mock_create_many_keys
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.prisma_client",
mock_prisma_client,
)
monkeypatch.setattr(
"litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin",
lambda **kwargs: True,
)
data = TeamMemberDeleteRequest(team_id="team-1", user_id="user-123")
result = await team_member_delete(
data=data,
user_api_key_dict=mock_user_api_key_dict,
)
mock_create_many_keys.assert_called_once()
call_args = mock_create_many_keys.call_args
assert "data" in call_args.kwargs
records = call_args.kwargs["data"]
assert len(records) == 2
assert all(record["deleted_by"] == "admin-user" for record in records)
assert all(record["team_id"] == "team-1" for record in records)
assert all(record["user_id"] == "user-123" for record in records)
mock_delete_keys.assert_called_once()
@pytest.mark.asyncio
async def test_new_team_negative_max_budget():
"""
@ -4825,6 +4477,187 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth):
assert deserialized_settings == router_settings_data
@pytest.mark.asyncio
async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys(
mock_db_client,
):
"""
Test that non-team-admin users only see their own spend (filtered by their API keys)
when calling /team/daily/activity endpoint.
"""
from litellm.proxy.management_endpoints.team_endpoints import (
get_team_daily_activity,
)
# Create a non-admin user
user_id = "test_user_123"
team_id = "test_team_456"
user_api_key_dict = UserAPIKeyAuth(
user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER
)
# Mock user info
mock_user_info = LiteLLM_UserTable(
user_id=user_id,
teams=[team_id],
max_budget=1000.0,
spend=0.0,
user_email="test@example.com",
user_role="internal_user",
)
# Mock team with user as non-admin member
mock_team_member = Member(user_id=user_id, role="user")
mock_team = MagicMock(spec=LiteLLM_TeamTable)
mock_team.team_id = team_id
mock_team.team_alias = "Test Team"
mock_team.members_with_roles = [mock_team_member]
mock_team.model_dump.return_value = {
"team_id": team_id,
"team_alias": "Test Team",
"members_with_roles": [{"user_id": user_id, "role": "user"}],
}
# Mock user's API keys
user_api_key_1 = MagicMock()
user_api_key_1.token = "user_key_1"
user_api_key_2 = MagicMock()
user_api_key_2.token = "user_key_2"
# Setup mocks
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(
return_value=[mock_team]
)
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[user_api_key_1, user_api_key_2]
)
# Mock get_user_object
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_user_object",
new_callable=AsyncMock,
) as mock_get_user_object:
mock_get_user_object.return_value = mock_user_info
# Mock get_daily_activity to capture the api_key parameter
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_daily_activity",
new_callable=AsyncMock,
) as mock_get_daily_activity:
mock_get_daily_activity.return_value = MagicMock()
# Call the endpoint
await get_team_daily_activity(
team_ids=team_id,
start_date="2024-01-01",
end_date="2024-01-02",
model=None,
api_key=None,
page=1,
page_size=10,
exclude_team_ids=None,
user_api_key_dict=user_api_key_dict,
)
# Verify get_daily_activity was called with user's API keys as filter
mock_get_daily_activity.assert_called_once()
call_kwargs = mock_get_daily_activity.call_args[1]
assert call_kwargs["api_key"] == ["user_key_1", "user_key_2"]
assert call_kwargs["entity_id"] == [team_id]
# Verify user's API keys were fetched
mock_db_client.db.litellm_verificationtoken.find_many.assert_called_once()
api_key_call_kwargs = (
mock_db_client.db.litellm_verificationtoken.find_many.call_args[1]
)
assert api_key_call_kwargs["where"] == {"user_id": user_id}
@pytest.mark.asyncio
async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client):
"""
Test that team admin users see all team spend (no API key filtering)
when calling /team/daily/activity endpoint.
"""
from litellm.proxy.management_endpoints.team_endpoints import (
get_team_daily_activity,
)
# Create a team admin user
user_id = "test_admin_123"
team_id = "test_team_456"
user_api_key_dict = UserAPIKeyAuth(
user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER
)
# Mock user info
mock_user_info = LiteLLM_UserTable(
user_id=user_id,
teams=[team_id],
max_budget=1000.0,
spend=0.0,
user_email="admin@example.com",
user_role="internal_user",
)
# Mock team with user as admin member
mock_team_member = Member(user_id=user_id, role="admin")
mock_team = MagicMock(spec=LiteLLM_TeamTable)
mock_team.team_id = team_id
mock_team.team_alias = "Test Team"
mock_team.members_with_roles = [mock_team_member]
mock_team.model_dump.return_value = {
"team_id": team_id,
"team_alias": "Test Team",
"members_with_roles": [{"user_id": user_id, "role": "admin"}],
}
# Setup mocks
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(
return_value=[mock_team]
)
# Mock get_user_object
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_user_object",
new_callable=AsyncMock,
) as mock_get_user_object:
mock_get_user_object.return_value = mock_user_info
# Mock get_daily_activity to capture the api_key parameter
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_daily_activity",
new_callable=AsyncMock,
) as mock_get_daily_activity:
mock_get_daily_activity.return_value = MagicMock()
# Call the endpoint
await get_team_daily_activity(
team_ids=team_id,
start_date="2024-01-01",
end_date="2024-01-02",
model=None,
api_key=None,
page=1,
page_size=10,
exclude_team_ids=None,
user_api_key_dict=user_api_key_dict,
)
# Verify get_daily_activity was called WITHOUT API key filtering
mock_get_daily_activity.assert_called_once()
call_kwargs = mock_get_daily_activity.call_args[1]
assert call_kwargs["api_key"] is None
assert call_kwargs["entity_id"] == [team_id]
# Verify user's API keys were NOT fetched (since they're admin)
if hasattr(
mock_db_client.db.litellm_verificationtoken, "find_many"
) and mock_db_client.db.litellm_verificationtoken.find_many.called:
# If it was called, that's unexpected for admin users
assert False, "API keys should not be fetched for team admin users"
@pytest.mark.asyncio
async def test_update_team_with_router_settings(mock_db_client, mock_admin_auth):
"""
@ -4901,3 +4734,184 @@ async def test_update_team_with_router_settings(mock_db_client, mock_admin_auth)
# Verify router_settings can be deserialized and matches input
deserialized_settings = json.loads(team_data["router_settings"])
assert deserialized_settings == router_settings_data
@pytest.mark.asyncio
async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys(
mock_db_client,
):
"""
Test that non-team-admin users only see their own spend (filtered by their API keys)
when calling /team/daily/activity endpoint.
"""
from litellm.proxy.management_endpoints.team_endpoints import (
get_team_daily_activity,
)
# Create a non-admin user
user_id = "test_user_123"
team_id = "test_team_456"
user_api_key_dict = UserAPIKeyAuth(
user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER
)
# Mock user info
mock_user_info = LiteLLM_UserTable(
user_id=user_id,
teams=[team_id],
max_budget=1000.0,
spend=0.0,
user_email="test@example.com",
user_role="internal_user",
)
# Mock team with user as non-admin member
mock_team_member = Member(user_id=user_id, role="user")
mock_team = MagicMock(spec=LiteLLM_TeamTable)
mock_team.team_id = team_id
mock_team.team_alias = "Test Team"
mock_team.members_with_roles = [mock_team_member]
mock_team.model_dump.return_value = {
"team_id": team_id,
"team_alias": "Test Team",
"members_with_roles": [{"user_id": user_id, "role": "user"}],
}
# Mock user's API keys
user_api_key_1 = MagicMock()
user_api_key_1.token = "user_key_1"
user_api_key_2 = MagicMock()
user_api_key_2.token = "user_key_2"
# Setup mocks
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(
return_value=[mock_team]
)
mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(
return_value=[user_api_key_1, user_api_key_2]
)
# Mock get_user_object
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_user_object",
new_callable=AsyncMock,
) as mock_get_user_object:
mock_get_user_object.return_value = mock_user_info
# Mock get_daily_activity to capture the api_key parameter
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_daily_activity",
new_callable=AsyncMock,
) as mock_get_daily_activity:
mock_get_daily_activity.return_value = MagicMock()
# Call the endpoint
await get_team_daily_activity(
team_ids=team_id,
start_date="2024-01-01",
end_date="2024-01-02",
model=None,
api_key=None,
page=1,
page_size=10,
exclude_team_ids=None,
user_api_key_dict=user_api_key_dict,
)
# Verify get_daily_activity was called with user's API keys as filter
mock_get_daily_activity.assert_called_once()
call_kwargs = mock_get_daily_activity.call_args[1]
assert call_kwargs["api_key"] == ["user_key_1", "user_key_2"]
assert call_kwargs["entity_id"] == [team_id]
# Verify user's API keys were fetched
mock_db_client.db.litellm_verificationtoken.find_many.assert_called_once()
api_key_call_kwargs = (
mock_db_client.db.litellm_verificationtoken.find_many.call_args[1]
)
assert api_key_call_kwargs["where"] == {"user_id": user_id}
@pytest.mark.asyncio
async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client):
"""
Test that team admin users see all team spend (no API key filtering)
when calling /team/daily/activity endpoint.
"""
from litellm.proxy.management_endpoints.team_endpoints import (
get_team_daily_activity,
)
# Create a team admin user
user_id = "test_admin_123"
team_id = "test_team_456"
user_api_key_dict = UserAPIKeyAuth(
user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER
)
# Mock user info
mock_user_info = LiteLLM_UserTable(
user_id=user_id,
teams=[team_id],
max_budget=1000.0,
spend=0.0,
user_email="admin@example.com",
user_role="internal_user",
)
# Mock team with user as admin member
mock_team_member = Member(user_id=user_id, role="admin")
mock_team = MagicMock(spec=LiteLLM_TeamTable)
mock_team.team_id = team_id
mock_team.team_alias = "Test Team"
mock_team.members_with_roles = [mock_team_member]
mock_team.model_dump.return_value = {
"team_id": team_id,
"team_alias": "Test Team",
"members_with_roles": [{"user_id": user_id, "role": "admin"}],
}
# Setup mocks
mock_db_client.db.litellm_teamtable.find_many = AsyncMock(
return_value=[mock_team]
)
# Mock get_user_object
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_user_object",
new_callable=AsyncMock,
) as mock_get_user_object:
mock_get_user_object.return_value = mock_user_info
# Mock get_daily_activity to capture the api_key parameter
with patch(
"litellm.proxy.management_endpoints.team_endpoints.get_daily_activity",
new_callable=AsyncMock,
) as mock_get_daily_activity:
mock_get_daily_activity.return_value = MagicMock()
# Call the endpoint
await get_team_daily_activity(
team_ids=team_id,
start_date="2024-01-01",
end_date="2024-01-02",
model=None,
api_key=None,
page=1,
page_size=10,
exclude_team_ids=None,
user_api_key_dict=user_api_key_dict,
)
# Verify get_daily_activity was called WITHOUT API key filtering
mock_get_daily_activity.assert_called_once()
call_kwargs = mock_get_daily_activity.call_args[1]
assert call_kwargs["api_key"] is None
assert call_kwargs["entity_id"] == [team_id]
# Verify user's API keys were NOT fetched (since they're admin)
if hasattr(
mock_db_client.db.litellm_verificationtoken, "find_many"
) and mock_db_client.db.litellm_verificationtoken.find_many.called:
# If it was called, that's unexpected for admin users
assert False, "API keys should not be fetched for team admin users"

16
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -1393,23 +1393,21 @@ async def test_embedding_header_forwarding_with_model_group():
version="test-version",
)
# Verify that headers were added to the request metadata
assert "metadata" in updated_data, "Metadata should be added to embedding request"
assert "headers" in updated_data["metadata"], "Headers should be added to embedding request metadata"
# Verify that headers were added to the request data
assert "headers" in updated_data, "Headers should be added to embedding request"
# Verify that only x- prefixed headers (except x-stainless) were forwarded
forwarded_headers = updated_data["metadata"]["headers"]
forwarded_headers = updated_data["headers"]
assert "X-Custom-Header" in forwarded_headers, "X-Custom-Header should be forwarded"
assert forwarded_headers["X-Custom-Header"] == "custom-value"
assert "X-Request-ID" in forwarded_headers, "X-Request-ID should be forwarded"
assert forwarded_headers["X-Request-ID"] == "test-request-123"
# Verify that Authorization header is present in metadata (not filtered out at this level)
# Note: The metadata headers contain all original headers for logging/tracking purposes
assert "Authorization" in forwarded_headers, "Authorization header should be in metadata headers"
# Verify that authorization header was NOT forwarded (sensitive header)
assert "Authorization" not in forwarded_headers, "Authorization header should not be forwarded"
# Verify that Content-Type is present (it's included in metadata headers)
assert "Content-Type" in forwarded_headers, "Content-Type should be in metadata headers"
# Verify that Content-Type was NOT forwarded (doesn't start with x-)
assert "Content-Type" not in forwarded_headers, "Content-Type should not be forwarded"
# Verify original data fields are preserved
assert updated_data["model"] == "local-openai/text-embedding-3-small"

69
tests/test_litellm/proxy/test_proxy_cli.py
View file

@ -483,75 +483,6 @@ class TestProxyInitializationHelpers:
# Verify that uvicorn.run was called again
mock_uvicorn_run.assert_called_once()
@patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers._run_gunicorn_server")
@patch("builtins.print")
def test_gunicorn_keepalive_timeout_flag(self, mock_print, mock_gunicorn):
"""Test that the keepalive_timeout flag is properly passed to Gunicorn"""
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_app = MagicMock()
mock_proxy_config = MagicMock()
mock_key_mgmt = MagicMock()
mock_save_worker_config = MagicMock()
with patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=mock_app,
ProxyConfig=mock_proxy_config,
KeyManagementSettings=mock_key_mgmt,
save_worker_config=mock_save_worker_config,
)
},
):
result = runner.invoke(
run_server, ["--local", "--run_gunicorn", "--keepalive_timeout", "120"]
)
assert result.exit_code == 0
# Verify _run_gunicorn_server was called with keepalive_timeout
mock_gunicorn.assert_called_once()
call_kwargs = mock_gunicorn.call_args.kwargs
assert call_kwargs["keepalive_timeout"] == 120
@patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers._run_gunicorn_server")
@patch("builtins.print")
def test_gunicorn_keepalive_default(self, mock_print, mock_gunicorn):
"""Test that Gunicorn uses default 90s when keepalive_timeout not specified"""
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_app = MagicMock()
mock_proxy_config = MagicMock()
mock_key_mgmt = MagicMock()
mock_save_worker_config = MagicMock()
with patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=mock_app,
ProxyConfig=mock_proxy_config,
KeyManagementSettings=mock_key_mgmt,
save_worker_config=mock_save_worker_config,
)
},
):
result = runner.invoke(run_server, ["--local", "--run_gunicorn"])
assert result.exit_code == 0
# Verify default behavior (keepalive_timeout is None, Gunicorn will use 90)
call_kwargs = mock_gunicorn.call_args.kwargs
assert call_kwargs.get("keepalive_timeout") is None
class TestHealthAppFactory:
"""Test cases for the health app factory module"""

39
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -55,7 +55,7 @@ example_embedding_result = {
def mock_patch_aembedding():
return mock.patch(
"litellm.aembedding",
"litellm.proxy.proxy_server.llm_router.aembedding",
return_value=example_embedding_result,
)
@ -668,6 +668,43 @@ def test_team_info_masking():
assert "public-test-key" not in str(exc_info.value)
@mock_patch_aembedding()
def test_embedding_input_array_of_tokens(mock_aembedding, client_no_auth):
"""
Test to bypass decoding input as array of tokens for selected providers
Ref: https://github.com/BerriAI/litellm/issues/10113
"""
try:
test_data = {
"model": "vllm_embed_model",
"input": [[2046, 13269, 158208]],
}
response = client_no_auth.post("/v1/embeddings", json=test_data)
# DEPRECATED - mock_aembedding.assert_called_once_with is too strict, and will fail when new kwargs are added to embeddings
# mock_aembedding.assert_called_once_with(
# model="vllm_embed_model",
# input=[[2046, 13269, 158208]],
# metadata=mock.ANY,
# proxy_server_request=mock.ANY,
# secret_fields=mock.ANY,
# )
# Assert that aembedding was called, and that input was not modified
mock_aembedding.assert_called_once()
call_args, call_kwargs = mock_aembedding.call_args
assert call_kwargs["model"] == "vllm_embed_model"
assert call_kwargs["input"] == [[2046, 13269, 158208]]
assert response.status_code == 200
result = response.json()
print(len(result["data"][0]["embedding"]))
assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so
except Exception as e:
pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")
@pytest.mark.asyncio
async def test_get_all_team_models():
"""

384
tests/test_litellm/test_router.py
View file

@ -1171,203 +1171,6 @@ async def test_acompletion_streaming_iterator_edge_cases():
print("✓ Edge case tests passed!")
@pytest.mark.asyncio
async def test_acompletion_streaming_disable_fallbacks_midstream():
"""Test that disable_fallbacks=True prevents mid-stream fallback attempts."""
from unittest.mock import MagicMock
from litellm.exceptions import MidStreamFallbackError
# Set up router with fallback configuration
router = litellm.Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key-2"},
},
],
fallbacks=[{"gpt-4": ["gpt-3.5-turbo"]}],
set_verbose=True,
)
messages = [{"role": "user", "content": "Hello"}]
# Test 1: disable_fallbacks=True with original_exception
print("\n=== Test 1: disable_fallbacks=True with original_exception ===")
# Create an original exception to wrap
from litellm.llms.anthropic.common_utils import AnthropicError
original_error = AnthropicError(
status_code=500,
message="An unexpected error occurred while processing the response",
)
# Create MidStreamFallbackError with original_exception
error_with_original = MidStreamFallbackError(
message="Connection lost",
model="gpt-4",
llm_provider="openai",
generated_content="Hello",
original_exception=original_error,
)
class AsyncIteratorWithError:
def __init__(self, items, error_after_index, error):
self.items = items
self.index = 0
self.error_after_index = error_after_index
self.error = error
self.chunks = []
self.model = "gpt-4"
self.custom_llm_provider = "openai"
self.logging_obj = MagicMock()
def __aiter__(self):
return self
async def __anext__(self):
if self.index == self.error_after_index:
raise self.error
if self.index >= len(self.items):
raise StopAsyncIteration
item = self.items[self.index]
self.index += 1
self.chunks.append(item)
return item
# Create properly structured mock chunks using ModelResponse
from litellm.types.utils import Delta, ModelResponse, StreamingChoices
mock_chunk = ModelResponse(
id="chatcmpl-123",
choices=[
StreamingChoices(
index=0, delta=Delta(content="Hello", role="assistant"), finish_reason=None
)
],
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
)
mock_chunks = [mock_chunk]
mock_error_response = AsyncIteratorWithError(
mock_chunks, 1, error_with_original
) # Error after first chunk
initial_kwargs = {"model": "gpt-4", "stream": True, "disable_fallbacks": True}
# Mock the fallback function to ensure it's NOT called
with patch.object(
router,
"async_function_with_fallbacks_common_utils",
return_value=MagicMock(),
) as mock_fallback_utils:
with pytest.raises(AnthropicError, match="An unexpected error occurred"):
result = await router._acompletion_streaming_iterator(
model_response=mock_error_response,
messages=messages,
initial_kwargs=initial_kwargs,
)
async for chunk in result:
pass # Should not reach here; exception should be raised
# Verify fallback was NOT called
mock_fallback_utils.assert_not_called()
print("✓ Original exception raised correctly when disable_fallbacks=True")
# Test 2: disable_fallbacks=True without original_exception
print("\n=== Test 2: disable_fallbacks=True without original_exception ===")
error_without_original = MidStreamFallbackError(
message="Connection lost",
model="gpt-4",
llm_provider="openai",
generated_content="Hello",
original_exception=None,
)
mock_error_response_2 = AsyncIteratorWithError(
mock_chunks, 1, error_without_original
)
with patch.object(
router,
"async_function_with_fallbacks_common_utils",
return_value=MagicMock(),
) as mock_fallback_utils:
with pytest.raises(MidStreamFallbackError, match="Connection lost"):
result = await router._acompletion_streaming_iterator(
model_response=mock_error_response_2,
messages=messages,
initial_kwargs=initial_kwargs,
)
async for chunk in result:
pass # Should not reach here
# Verify fallback was NOT called
mock_fallback_utils.assert_not_called()
print(
"✓ MidStreamFallbackError raised correctly when no original_exception and disable_fallbacks=True"
)
# Test 3: disable_fallbacks=False (default behavior - fallback should work)
print("\n=== Test 3: disable_fallbacks=False (fallback enabled) ===")
error_for_fallback = MidStreamFallbackError(
message="Connection lost",
model="gpt-4",
llm_provider="openai",
generated_content="Hello",
)
mock_error_response_3 = AsyncIteratorWithError(mock_chunks, 1, error_for_fallback)
# Mock successful fallback response
class EmptyAsyncIterator:
def __aiter__(self):
return self
async def __anext__(self):
raise StopAsyncIteration
mock_fallback_response = EmptyAsyncIterator()
initial_kwargs_fallback_enabled = {
"model": "gpt-4",
"stream": True,
"disable_fallbacks": False,
}
with patch.object(
router,
"async_function_with_fallbacks_common_utils",
return_value=mock_fallback_response,
) as mock_fallback_utils:
collected_chunks = []
result = await router._acompletion_streaming_iterator(
model_response=mock_error_response_3,
messages=messages,
initial_kwargs=initial_kwargs_fallback_enabled,
)
async for chunk in result:
collected_chunks.append(chunk)
# Verify fallback WAS called
assert mock_fallback_utils.called
print("✓ Fallback called correctly when disable_fallbacks=False")
print("\n=== All disable_fallbacks tests passed! ===")
@pytest.mark.asyncio
async def test_async_function_with_fallbacks_common_utils():
"""Test the async_function_with_fallbacks_common_utils method"""
@ -2066,190 +1869,3 @@ async def test_aguardrail():
assert result["result"] == "success"
assert result["selected_guardrail"]["id"] == "guardrail-1"
def test_resolve_model_name_from_model_id_wildcard_pattern():
"""
Test that resolve_model_name_from_model_id correctly resolves model names
for wildcard patterns using PatternMatchRouter.
This is critical for video status/content endpoints where model_id extracted
from video_id (e.g., "veo-3.0-generate-preview") needs to match wildcard
patterns like "vertex_ai/*" to inject credentials from the model config.
"""
# Set up router with wildcard pattern
router = litellm.Router(
model_list=[
{
"model_name": "vertex_ai/*",
"litellm_params": {
"model": "vertex_ai/*",
"vertex_project": "test-project",
"vertex_location": "us-central1",
},
},
{
"model_name": "specific-model",
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"vertex_project": "specific-project",
"vertex_location": "us-east1",
},
},
],
)
# Test Case 1: Wildcard pattern matching with custom_llm_provider
# This simulates video_id like "vertex_ai:veo-3.0-generate-preview:..."
result = router.resolve_model_name_from_model_id(
model_id="veo-3.0-generate-preview",
custom_llm_provider="vertex_ai",
)
assert result == "vertex_ai/*", f"Expected 'vertex_ai/*', got '{result}'"
# Test Case 2: Different model name should also match wildcard
result = router.resolve_model_name_from_model_id(
model_id="gemini-2.0-flash",
custom_llm_provider="vertex_ai",
)
assert result == "vertex_ai/*", f"Expected 'vertex_ai/*', got '{result}'"
# Test Case 3: Without custom_llm_provider, should not match wildcard
result = router.resolve_model_name_from_model_id(
model_id="veo-3.0-generate-preview",
custom_llm_provider=None,
)
assert result is None, f"Expected None without provider, got '{result}'"
# Test Case 4: Exact model_name match should take precedence
result = router.resolve_model_name_from_model_id(
model_id="specific-model",
custom_llm_provider="vertex_ai",
)
assert result == "specific-model", f"Expected 'specific-model', got '{result}'"
def test_resolve_model_name_from_model_id_exact_match():
"""
Test that resolve_model_name_from_model_id correctly resolves exact model names.
"""
router = litellm.Router(
model_list=[
{
"model_name": "my-gpt-model",
"litellm_params": {
"model": "azure/gpt-4",
"api_key": "test-key",
},
},
{
"model_name": "veo-model",
"litellm_params": {
"model": "vertex_ai/veo-2.0-generate-001",
"vertex_project": "test-project",
},
},
],
)
# Test Case 1: Direct model_name match
result = router.resolve_model_name_from_model_id(model_id="my-gpt-model")
assert result == "my-gpt-model", f"Expected 'my-gpt-model', got '{result}'"
# Test Case 2: Match by litellm_params.model suffix
result = router.resolve_model_name_from_model_id(model_id="veo-2.0-generate-001")
assert result == "veo-model", f"Expected 'veo-model', got '{result}'"
# Test Case 3: Non-existent model should return None
result = router.resolve_model_name_from_model_id(model_id="non-existent-model")
assert result is None, f"Expected None, got '{result}'"
def test_resolve_model_name_from_model_id_provider_prefix():
"""
Test that resolve_model_name_from_model_id handles provider prefix correctly.
"""
router = litellm.Router(
model_list=[
{
"model_name": "vertex_ai/gemini-pro",
"litellm_params": {
"model": "vertex_ai/gemini-pro",
"vertex_project": "test-project",
},
},
],
)
# Test Case 1: Full model name with provider prefix as model_name
result = router.resolve_model_name_from_model_id(
model_id="vertex_ai/gemini-pro",
custom_llm_provider=None,
)
assert result == "vertex_ai/gemini-pro", f"Expected 'vertex_ai/gemini-pro', got '{result}'"
# Test Case 2: Model ID with provider prefix constructed from custom_llm_provider
result = router.resolve_model_name_from_model_id(
model_id="gemini-pro",
custom_llm_provider="vertex_ai",
)
assert result == "vertex_ai/gemini-pro", f"Expected 'vertex_ai/gemini-pro', got '{result}'"
def test_resolve_model_name_from_model_id_multiple_wildcards():
"""
Test that resolve_model_name_from_model_id works with multiple wildcard patterns.
"""
router = litellm.Router(
model_list=[
{
"model_name": "vertex_ai/*",
"litellm_params": {
"model": "vertex_ai/*",
"vertex_project": "vertex-project",
},
},
{
"model_name": "openai/*",
"litellm_params": {
"model": "openai/*",
"api_key": "openai-key",
},
},
{
"model_name": "anthropic/*",
"litellm_params": {
"model": "anthropic/*",
"api_key": "anthropic-key",
},
},
],
)
# Test Case 1: Match vertex_ai wildcard
result = router.resolve_model_name_from_model_id(
model_id="veo-3.0-generate-preview",
custom_llm_provider="vertex_ai",
)
assert result == "vertex_ai/*", f"Expected 'vertex_ai/*', got '{result}'"
# Test Case 2: Match openai wildcard
result = router.resolve_model_name_from_model_id(
model_id="gpt-4o",
custom_llm_provider="openai",
)
assert result == "openai/*", f"Expected 'openai/*', got '{result}'"
# Test Case 3: Match anthropic wildcard
result = router.resolve_model_name_from_model_id(
model_id="claude-3-opus",
custom_llm_provider="anthropic",
)
assert result == "anthropic/*", f"Expected 'anthropic/*', got '{result}'"
# Test Case 4: Non-matching provider should return None
result = router.resolve_model_name_from_model_id(
model_id="some-model",
custom_llm_provider="bedrock",
)
assert result is None, f"Expected None for non-matching provider, got '{result}'"