tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-09 04:21:42 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'BerriAI:main' into main

Browse Source
This commit is contained in:
abbas jafari
2025-11-26 10:22:13 +01:00
committed by GitHub
parent aec0ab777b 38bac312cd
commit fdee1e2946
81 changed files with 8524 additions and 551 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/blog/anthropic_opus_4_5_and_advanced_features/index.md
+1281
View File
File diff suppressed because it is too large Load Diff
docs/my-website/docs/completion/image_generation_chat.md
+2 -2
View File
@@ -224,8 +224,8 @@ asyncio.run(generate_image())
| Provider | Model |
|----------|--------|
| Google AI Studio | `gemini/gemini-2.0-flash-preview-image-generation`, `gemini/gemini-2.5-flash-image-preview` |
| Vertex AI | `vertex_ai/gemini-2.0-flash-preview-image-generation`, `vertex_ai/gemini-2.5-flash-image-preview` |
| Google AI Studio | `gemini/gemini-2.0-flash-preview-image-generation`, `gemini/gemini-2.5-flash-image-preview`, `gemini/gemini-3-pro-image-preview` |
| Vertex AI | `vertex_ai/gemini-2.0-flash-preview-image-generation`, `vertex_ai/gemini-2.5-flash-image-preview`, `vertex_ai/gemini-3-pro-image-preview` |
## Spec
docs/my-website/docs/providers/anthropic_effort.md
+279
View File
@@ -0,0 +1,279 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Anthropic Effort Parameter
Control how many tokens Claude uses when responding with the `effort` parameter, trading off between response thoroughness and token efficiency.
## Overview
The `effort` parameter allows you to control how eager Claude is about spending tokens when responding to requests. This gives you the ability to trade off between response thoroughness and token efficiency, all with a single model.
**Note**: The effort parameter is currently in beta and only supported by Claude Opus 4.5. You must include the beta header `effort-2025-11-24` when using this feature (LiteLLM automatically adds this header when `output_config` with `effort` is detected).
## How Effort Works
By default, Claude uses maximum effort—spending as many tokens as needed for the best possible outcome. By lowering the effort level, you can instruct Claude to be more conservative with token usage, optimizing for speed and cost while accepting some reduction in capability.
**Tip**: Setting `effort` to `"high"` produces exactly the same behavior as omitting the `effort` parameter entirely.
The effort parameter affects **all tokens** in the response, including:
- Text responses and explanations
- Tool calls and function arguments
- Extended thinking (when enabled)
This approach has two major advantages:
1. It doesn't require thinking to be enabled in order to use it.
2. It can affect all token spend including tool calls. For example, lower effort would mean Claude makes fewer tool calls.
This gives a much greater degree of control over efficiency.
## Effort Levels
| Level | Description | Typical use case |
|-------|-------------|------------------|
| `high` | Maximum capability—Claude uses as many tokens as needed for the best possible outcome. Equivalent to not setting the parameter. | Complex reasoning, difficult coding problems, agentic tasks |
| `medium` | Balanced approach with moderate token savings. | Agentic tasks that require a balance of speed, cost, and performance |
| `low` | Most efficient—significant token savings with some capability reduction. | Simpler tasks that need the best speed and lowest costs, such as subagents |
## Quick Start
### Using LiteLLM SDK
<Tabs>
<TabItem value="python" label="Python">
```python
import litellm
response = litellm.completion(
model="anthropic/claude-opus-4-5-20251101",
messages=[{
"role": "user",
"content": "Analyze the trade-offs between microservices and monolithic architectures"
}],
output_config={
"effort": "medium"
}
)
print(response.choices[0].message.content)
```
</TabItem>
<TabItem value="typescript" label="TypeScript">
```typescript
import Anthropic from "@anthropic-ai/sdk";
const client = new Anthropic({
apiKey: process.env.ANTHROPIC_API_KEY,
});
const response = await client.messages.create({
model: "claude-opus-4-5-20251101",
max_tokens: 4096,
messages: [{
role: "user",
content: "Analyze the trade-offs between microservices and monolithic architectures"
}],
output_config: {
effort: "medium"
}
});
console.log(response.content[0].text);
```
</TabItem>
</Tabs>
### Using LiteLLM Proxy
```bash
curl http://localhost:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer $LITELLM_API_KEY" \
-d '{
"model": "anthropic/claude-opus-4-5-20251101",
"messages": [{
"role": "user",
"content": "Analyze the trade-offs between microservices and monolithic architectures"
}],
"output_config": {
"effort": "medium"
}
}'
```
### Direct Anthropic API Call
```bash
curl https://api.anthropic.com/v1/messages \
--header "x-api-key: $ANTHROPIC_API_KEY" \
--header "anthropic-version: 2023-06-01" \
--header "anthropic-beta: effort-2025-11-24" \
--header "content-type: application/json" \
--data '{
"model": "claude-opus-4-5-20251101",
"max_tokens": 4096,
"messages": [{
"role": "user",
"content": "Analyze the trade-offs between microservices and monolithic architectures"
}],
"output_config": {
"effort": "medium"
}
}'
```
## Model Compatibility
The effort parameter is currently only supported by:
- **Claude Opus 4.5** (`claude-opus-4-5-20251101`)
## When Should I Adjust the Effort Parameter?
- Use **high effort** (the default) when you need Claude's best work—complex reasoning, nuanced analysis, difficult coding problems, or any task where quality is the top priority.
- Use **medium effort** as a balanced option when you want solid performance without the full token expenditure of high effort.
- Use **low effort** when you're optimizing for speed (because Claude answers with fewer tokens) or cost—for example, simple classification tasks, quick lookups, or high-volume use cases where marginal quality improvements don't justify additional latency or spend.
## Effort with Tool Use
When using tools, the effort parameter affects both the explanations around tool calls and the tool calls themselves. Lower effort levels tend to:
- Combine multiple operations into fewer tool calls
- Make fewer tool calls
- Proceed directly to action
Example with tools:
```python
import litellm
response = litellm.completion(
model="anthropic/claude-opus-4-5-20251101",
messages=[{
"role": "user",
"content": "Check the weather in multiple cities"
}],
tools=[{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather for a location",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
},
"required": ["location"]
}
}
}],
output_config={
"effort": "low" # Will make fewer tool calls
}
)
```
## Effort with Extended Thinking
The effort parameter works seamlessly with extended thinking. When both are enabled, effort controls the token budget across all response types:
```python
import litellm
response = litellm.completion(
model="anthropic/claude-opus-4-5-20251101",
messages=[{
"role": "user",
"content": "Solve this complex problem"
}],
thinking={
"type": "enabled",
"budget_tokens": 5000
},
output_config={
"effort": "medium" # Affects both thinking and response tokens
}
)
```
## Best Practices
1. **Start with the default (high)** for new tasks, then experiment with lower effort levels if you're looking to optimize costs.
2. **Use medium effort for production agentic workflows** where you need a balance of quality and efficiency.
3. **Reserve low effort for high-volume, simple tasks** like classification, routing, or data extraction where speed matters more than nuanced responses.
4. **Monitor token usage** to understand the actual savings from different effort levels for your specific use cases.
5. **Test with your specific prompts** as the impact of effort levels can vary based on task complexity.
## Provider Support
The effort parameter is supported across all Anthropic-compatible providers:
- **Standard Anthropic**: ✅ Supported (Claude Opus 4.5)
- **Azure Anthropic**: ✅ Supported (Claude Opus 4.5)
- **Vertex AI Anthropic**: ✅ Supported (Claude Opus 4.5)
LiteLLM automatically handles the beta header injection for all providers.
## Usage and Pricing
Token usage with different effort levels is tracked in the standard usage object. Lower effort levels result in fewer output tokens, which directly reduces costs:
```python
response = litellm.completion(
model="anthropic/claude-opus-4-5-20251101",
messages=[{"role": "user", "content": "Analyze this"}],
output_config={"effort": "low"}
)
print(f"Output tokens: {response.usage.completion_tokens}")
print(f"Total tokens: {response.usage.total_tokens}")
```
## Troubleshooting
### Beta header not being added
LiteLLM automatically adds the `effort-2025-11-24` beta header when `output_config` with `effort` is detected. If you're not seeing the header:
1. Ensure you're using `output_config` with an `effort` field
2. Verify the model is Claude Opus 4.5
3. Check that LiteLLM version supports this feature
### Invalid effort value error
Only three values are accepted: `"high"`, `"medium"`, `"low"`. Any other value will raise a validation error:
```python
# ❌ This will raise an error
output_config={"effort": "very_low"}
# ✅ Use one of the valid values
output_config={"effort": "low"}
```
### Model not supported
Currently, only Claude Opus 4.5 supports the effort parameter. Using it with other models may result in the parameter being ignored or an error.
## Related Features
- [Extended Thinking](/docs/providers/anthropic_extended_thinking) - Control Claude's reasoning process
- [Tool Use](/docs/providers/anthropic_tools) - Enable Claude to use tools and functions
- [Programmatic Tool Calling](/docs/providers/anthropic_programmatic_tool_calling) - Let Claude write code that calls tools
- [Prompt Caching](/docs/providers/anthropic_prompt_caching) - Cache prompts to reduce costs
## Additional Resources
- [Anthropic Effort Documentation](https://docs.anthropic.com/en/docs/build-with-claude/effort)
- [LiteLLM Anthropic Provider Guide](/docs/providers/anthropic)
- [Cost Optimization Best Practices](/docs/guides/cost_optimization)
docs/my-website/docs/providers/anthropic_programmatic_tool_calling.md
+430
View File
@@ -0,0 +1,430 @@
# Anthropic Programmatic Tool Calling
Programmatic tool calling allows Claude to write code that calls your tools programmatically within a code execution container, rather than requiring round trips through the model for each tool invocation. This reduces latency for multi-tool workflows and decreases token consumption by allowing Claude to filter or process data before it reaches the model's context window.
:::info
Programmatic tool calling is currently in public beta. LiteLLM automatically adds the required `advanced-tool-use-2025-11-20` beta header when it detects tools with the `allowed_callers` field.
This feature requires the code execution tool to be enabled.
:::
## Model Compatibility
Programmatic tool calling is available on the following models:
| Model | Tool Version |
|-------|--------------|
| Claude Opus 4.5 (`claude-opus-4-5-20251101`) | `code_execution_20250825` |
| Claude Sonnet 4.5 (`claude-sonnet-4-5-20250929`) | `code_execution_20250825` |
## Quick Start
Here's a simple example where Claude programmatically queries a database multiple times and aggregates results:
```python
import litellm
response = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[
{
"role": "user",
"content": "Query sales data for the West, East, and Central regions, then tell me which region had the highest revenue"
}
],
tools=[
{
"type": "code_execution_20250825",
"name": "code_execution"
},
{
"type": "function",
"function": {
"name": "query_database",
"description": "Execute a SQL query against the sales database. Returns a list of rows as JSON objects.",
"parameters": {
"type": "object",
"properties": {
"sql": {
"type": "string",
"description": "SQL query to execute"
}
},
"required": ["sql"]
}
},
"allowed_callers": ["code_execution_20250825"]
}
]
)
print(response)
```
## How It Works
When you configure a tool to be callable from code execution and Claude decides to use that tool:
1. Claude writes Python code that invokes the tool as a function, potentially including multiple tool calls and pre/post-processing logic
2. Claude runs this code in a sandboxed container via code execution
3. When a tool function is called, code execution pauses and the API returns a `tool_use` block with a `caller` field
4. You provide the tool result, and code execution continues (intermediate results are not loaded into Claude's context window)
5. Once all code execution completes, Claude receives the final output and continues working on the task
This approach is particularly useful for:
- **Large data processing**: Filter or aggregate tool results before they reach Claude's context
- **Multi-step workflows**: Save tokens and latency by calling tools serially or in a loop without sampling Claude in-between tool calls
- **Conditional logic**: Make decisions based on intermediate tool results
## The `allowed_callers` Field
The `allowed_callers` field specifies which contexts can invoke a tool:
```python
{
"type": "function",
"function": {
"name": "query_database",
"description": "Execute a SQL query against the database",
"parameters": {...}
},
"allowed_callers": ["code_execution_20250825"]
}
```
**Possible values:**
- `["direct"]` - Only Claude can call this tool directly (default if omitted)
- `["code_execution_20250825"]` - Only callable from within code execution
- `["direct", "code_execution_20250825"]` - Callable both directly and from code execution
:::tip
We recommend choosing either `["direct"]` or `["code_execution_20250825"]` for each tool rather than enabling both, as this provides clearer guidance to Claude for how best to use the tool.
:::
## The `caller` Field in Responses
Every tool use block includes a `caller` field indicating how it was invoked:
**Direct invocation (traditional tool use):**
```python
{
"type": "tool_use",
"id": "toolu_abc123",
"name": "query_database",
"input": {"sql": "<sql>"},
"caller": {"type": "direct"}
}
```
**Programmatic invocation:**
```python
{
"type": "tool_use",
"id": "toolu_xyz789",
"name": "query_database",
"input": {"sql": "<sql>"},
"caller": {
"type": "code_execution_20250825",
"tool_id": "srvtoolu_abc123"
}
}
```
The `tool_id` references the code execution tool that made the programmatic call.
## Container Lifecycle
Programmatic tool calling uses code execution containers:
- **Container creation**: A new container is created for each session unless you reuse an existing one
- **Expiration**: Containers expire after approximately 4.5 minutes of inactivity (subject to change)
- **Container ID**: Pass the `container` parameter to reuse an existing container
- **Reuse**: Pass the container ID to maintain state across requests
```python
# First request - creates a new container
response1 = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[{"role": "user", "content": "Query the database"}],
tools=[...]
)
# Get container ID from response (if available in response metadata)
container_id = response1.get("container", {}).get("id")
# Second request - reuse the same container
response2 = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[...],
tools=[...],
container=container_id # Reuse container
)
```
:::warning
When a tool is called programmatically and the container is waiting for your tool result, you must respond before the container expires. Monitor the `expires_at` field. If the container expires, Claude may treat the tool call as timed out and retry it.
:::
## Example Workflow
### Step 1: Initial Request
```python
import litellm
response = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[{
"role": "user",
"content": "Query customer purchase history from the last quarter and identify our top 5 customers by revenue"
}],
tools=[
{
"type": "code_execution_20250825",
"name": "code_execution"
},
{
"type": "function",
"function": {
"name": "query_database",
"description": "Execute a SQL query against the sales database. Returns a list of rows as JSON objects.",
"parameters": {
"type": "object",
"properties": {
"sql": {"type": "string", "description": "SQL query to execute"}
},
"required": ["sql"]
}
},
"allowed_callers": ["code_execution_20250825"]
}
]
)
```
### Step 2: API Response with Tool Call
Claude writes code that calls your tool. The response includes:
```python
{
"role": "assistant",
"content": [
{
"type": "text",
"text": "I'll query the purchase history and analyze the results."
},
{
"type": "server_tool_use",
"id": "srvtoolu_abc123",
"name": "code_execution",
"input": {
"code": "results = await query_database('<sql>')\ntop_customers = sorted(results, key=lambda x: x['revenue'], reverse=True)[:5]"
}
},
{
"type": "tool_use",
"id": "toolu_def456",
"name": "query_database",
"input": {"sql": "<sql>"},
"caller": {
"type": "code_execution_20250825",
"tool_id": "srvtoolu_abc123"
}
}
],
"stop_reason": "tool_use"
}
```
### Step 3: Provide Tool Result
```python
# Add assistant's response and tool result to conversation
messages = [
{"role": "user", "content": "Query customer purchase history..."},
{
"role": "assistant",
"content": response.choices[0].message.content,
"tool_calls": response.choices[0].message.tool_calls
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_def456",
"content": '[{"customer_id": "C1", "revenue": 45000}, ...]'
}
]
}
]
# Continue the conversation
response2 = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=messages,
tools=[...]
)
```
### Step 4: Final Response
Once code execution completes, Claude provides the final response:
```python
{
"content": [
{
"type": "code_execution_tool_result",
"tool_use_id": "srvtoolu_abc123",
"content": {
"type": "code_execution_result",
"stdout": "Top 5 customers by revenue:\n1. Customer C1: $45,000\n...",
"stderr": "",
"return_code": 0
}
},
{
"type": "text",
"text": "I've analyzed the purchase history from last quarter. Your top 5 customers generated $167,500 in total revenue..."
}
],
"stop_reason": "end_turn"
}
```
## Advanced Patterns
### Batch Processing with Loops
Claude can write code that processes multiple items efficiently:
```python
# Claude writes code like this:
regions = ["West", "East", "Central", "North", "South"]
results = {}
for region in regions:
data = await query_database(f"SELECT SUM(revenue) FROM sales WHERE region='{region}'")
results[region] = data[0]["total"]
top_region = max(results.items(), key=lambda x: x[1])
print(f"Top region: {top_region[0]} with ${top_region[1]:,}")
```
This pattern:
- Reduces model round-trips from N (one per region) to 1
- Processes large result sets programmatically before returning to Claude
- Saves tokens by only returning aggregated conclusions
### Early Termination
Claude can stop processing as soon as success criteria are met:
```python
endpoints = ["us-east", "eu-west", "apac"]
for endpoint in endpoints:
status = await check_health(endpoint)
if status == "healthy":
print(f"Found healthy endpoint: {endpoint}")
break # Stop early
```
### Data Filtering
```python
logs = await fetch_logs(server_id)
errors = [log for log in logs if "ERROR" in log]
print(f"Found {len(errors)} errors")
for error in errors[-10:]: # Only return last 10 errors
print(error)
```
## Best Practices
### Tool Design
- **Provide detailed output descriptions**: Since Claude deserializes tool results in code, clearly document the format (JSON structure, field types, etc.)
- **Return structured data**: JSON or other easily parseable formats work best for programmatic processing
- **Keep responses concise**: Return only necessary data to minimize processing overhead
### When to Use Programmatic Calling
**Good use cases:**
- Processing large datasets where you only need aggregates or summaries
- Multi-step workflows with 3+ dependent tool calls
- Operations requiring filtering, sorting, or transformation of tool results
- Tasks where intermediate data shouldn't influence Claude's reasoning
- Parallel operations across many items (e.g., checking 50 endpoints)
**Less ideal use cases:**
- Single tool calls with simple responses
- Tools that need immediate user feedback
- Very fast operations where code execution overhead would outweigh the benefit
## Token Efficiency
Programmatic tool calling can significantly reduce token consumption:
- **Tool results from programmatic calls are not added to Claude's context** - only the final code output is
- **Intermediate processing happens in code** - filtering, aggregation, etc. don't consume model tokens
- **Multiple tool calls in one code execution** - reduces overhead compared to separate model turns
For example, calling 10 tools directly uses ~10x the tokens of calling them programmatically and returning a summary.
## Provider Support
LiteLLM supports programmatic tool calling across all Anthropic-compatible providers:
- **Standard Anthropic API** (`anthropic/claude-sonnet-4-5-20250929`)
- **Azure Anthropic** (`azure/claude-sonnet-4-5-20250929`)
- **Vertex AI Anthropic** (`vertex_ai/claude-sonnet-4-5-20250929`)
The beta header is automatically added when LiteLLM detects tools with `allowed_callers` field.
## Limitations
### Feature Incompatibilities
- **Structured outputs**: Tools with `strict: true` are not supported with programmatic calling
- **Tool choice**: You cannot force programmatic calling of a specific tool via `tool_choice`
- **Parallel tool use**: `disable_parallel_tool_use: true` is not supported with programmatic calling
### Tool Restrictions
The following tools cannot currently be called programmatically:
- Web search
- Web fetch
- Tools provided by an MCP connector
## Troubleshooting
### Common Issues
**"Tool not allowed" error**
- Verify your tool definition includes `"allowed_callers": ["code_execution_20250825"]`
- Check that you're using a compatible model (Claude Sonnet 4.5 or Opus 4.5)
**Container expiration**
- Ensure you respond to tool calls within the container's lifetime (~4.5 minutes)
- Consider implementing faster tool execution
**Beta header not added**
- LiteLLM automatically adds the beta header when it detects `allowed_callers`
- If you're manually setting headers, ensure you include `advanced-tool-use-2025-11-20`
## Related Features
- [Anthropic Tool Search](./anthropic_tool_search.md) - Dynamically discover and load tools on-demand
- [Anthropic Provider](./anthropic.md) - General Anthropic provider documentation
docs/my-website/docs/providers/anthropic_tool_input_examples.md
+438
View File
@@ -0,0 +1,438 @@
# Anthropic Tool Input Examples
Provide concrete examples of valid tool inputs to help Claude understand how to use your tools more effectively. This is particularly useful for complex tools with nested objects, optional parameters, or format-sensitive inputs.
:::info
Tool input examples is a beta feature. LiteLLM automatically adds the required `advanced-tool-use-2025-11-20` beta header when it detects tools with the `input_examples` field.
:::
## When to Use Input Examples
Input examples are most helpful for:
- **Complex nested objects**: Tools with deeply nested parameter structures
- **Optional parameters**: Showing when optional parameters should be included
- **Format-sensitive inputs**: Demonstrating expected formats (dates, addresses, etc.)
- **Enum values**: Illustrating valid enum choices in context
- **Edge cases**: Showing how to handle special cases
:::tip
**Prioritize descriptions first!** Clear, detailed tool descriptions are more important than examples. Use `input_examples` as a supplement for complex tools where descriptions alone may not be sufficient.
:::
## Quick Start
Add an `input_examples` field to your tool definition with an array of example input objects:
```python
import litellm
response = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[
{"role": "user", "content": "What's the weather like in San Francisco?"}
],
tools=[
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state, e.g. San Francisco, CA"
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
"description": "The unit of temperature"
}
},
"required": ["location"]
}
},
"input_examples": [
{
"location": "San Francisco, CA",
"unit": "fahrenheit"
},
{
"location": "Tokyo, Japan",
"unit": "celsius"
},
{
"location": "New York, NY" # 'unit' is optional
}
]
}
]
)
print(response)
```
## How It Works
When you provide `input_examples`:
1. **LiteLLM detects** the `input_examples` field in your tool definition
2. **Beta header added automatically**: The `advanced-tool-use-2025-11-20` header is injected
3. **Examples included in prompt**: Anthropic includes the examples alongside your tool schema
4. **Claude learns patterns**: The model uses examples to understand proper tool usage
5. **Better tool calls**: Claude makes more accurate tool calls with correct parameter formats
## Example Formats
### Simple Tool with Examples
```python
{
"type": "function",
"function": {
"name": "send_email",
"description": "Send an email to a recipient",
"parameters": {
"type": "object",
"properties": {
"to": {"type": "string", "description": "Email address"},
"subject": {"type": "string"},
"body": {"type": "string"}
},
"required": ["to", "subject", "body"]
}
},
"input_examples": [
{
"to": "user@example.com",
"subject": "Meeting Reminder",
"body": "Don't forget our meeting tomorrow at 2 PM."
},
{
"to": "team@company.com",
"subject": "Weekly Update",
"body": "Here's this week's progress report..."
}
]
}
```
### Complex Nested Objects
```python
{
"type": "function",
"function": {
"name": "create_calendar_event",
"description": "Create a new calendar event",
"parameters": {
"type": "object",
"properties": {
"title": {"type": "string"},
"start": {
"type": "object",
"properties": {
"date": {"type": "string"},
"time": {"type": "string"}
}
},
"attendees": {
"type": "array",
"items": {
"type": "object",
"properties": {
"email": {"type": "string"},
"optional": {"type": "boolean"}
}
}
}
},
"required": ["title", "start"]
}
},
"input_examples": [
{
"title": "Team Standup",
"start": {
"date": "2025-01-15",
"time": "09:00"
},
"attendees": [
{"email": "alice@example.com", "optional": False},
{"email": "bob@example.com", "optional": True}
]
},
{
"title": "Lunch Break",
"start": {
"date": "2025-01-15",
"time": "12:00"
}
# No attendees - showing optional field
}
]
}
```
### Format-Sensitive Parameters
```python
{
"type": "function",
"function": {
"name": "search_flights",
"description": "Search for available flights",
"parameters": {
"type": "object",
"properties": {
"origin": {"type": "string", "description": "Airport code"},
"destination": {"type": "string", "description": "Airport code"},
"date": {"type": "string", "description": "Date in YYYY-MM-DD format"},
"passengers": {"type": "integer"}
},
"required": ["origin", "destination", "date"]
}
},
"input_examples": [
{
"origin": "SFO",
"destination": "JFK",
"date": "2025-03-15",
"passengers": 2
},
{
"origin": "LAX",
"destination": "ORD",
"date": "2025-04-20",
"passengers": 1
}
]
}
```
## Requirements and Limitations
### Schema Validation
- Each example **must be valid** according to the tool's `input_schema`
- Invalid examples will return a **400 error** from Anthropic
- Validation happens server-side (LiteLLM passes examples through)
### Server-Side Tools Not Supported
Input examples are **only supported for user-defined tools**. The following server-side tools do NOT support `input_examples`:
- `web_search` (web search tool)
- `code_execution` (code execution tool)
- `computer_use` (computer use tool)
- `bash_tool` (bash execution tool)
- `text_editor` (text editor tool)
### Token Costs
Examples add to your prompt tokens:
- **Simple examples**: ~20-50 tokens per example
- **Complex nested objects**: ~100-200 tokens per example
- **Trade-off**: Higher token cost for better tool call accuracy
### Model Compatibility
Input examples work with all Claude models that support the `advanced-tool-use-2025-11-20` beta header:
- Claude Opus 4.5 (`claude-opus-4-5-20251101`)
- Claude Sonnet 4.5 (`claude-sonnet-4-5-20250929`)
- Claude Opus 4.1 (`claude-opus-4-1-20250805`)
:::note
On Google Cloud's Vertex AI and Amazon Bedrock, only Claude Opus 4.5 supports tool input examples.
:::
## Best Practices
### 1. Show Diverse Examples
Include examples that demonstrate different use cases:
```python
"input_examples": [
{"location": "San Francisco, CA", "unit": "fahrenheit"}, # US city
{"location": "Tokyo, Japan", "unit": "celsius"}, # International
{"location": "New York, NY"} # Optional param omitted
]
```
### 2. Demonstrate Optional Parameters
Show when optional parameters should and shouldn't be included:
```python
"input_examples": [
{
"query": "machine learning",
"filters": {"year": 2024, "category": "research"} # With optional filters
},
{
"query": "artificial intelligence" # Without optional filters
}
]
```
### 3. Illustrate Format Requirements
Make format expectations clear through examples:
```python
"input_examples": [
{
"phone": "+1-555-123-4567", # Shows expected phone format
"date": "2025-01-15", # Shows date format (YYYY-MM-DD)
"time": "14:30" # Shows time format (HH:MM)
}
]
```
### 4. Keep Examples Realistic
Use realistic, production-like examples rather than placeholder data:
```python
# ✅ Good - realistic examples
"input_examples": [
{"email": "alice@company.com", "role": "admin"},
{"email": "bob@company.com", "role": "user"}
]
# ❌ Bad - placeholder examples
"input_examples": [
{"email": "test@test.com", "role": "role1"},
{"email": "example@example.com", "role": "role2"}
]
```
### 5. Limit Example Count
Provide 2-5 examples per tool:
- **Too few** (1): May not show enough variation
- **Just right** (2-5): Demonstrates patterns without bloating tokens
- **Too many** (10+): Wastes tokens, diminishing returns
## Integration with Other Features
Input examples work seamlessly with other Anthropic tool features:
### With Tool Search
```python
{
"type": "function",
"function": {
"name": "query_database",
"description": "Execute a SQL query",
"parameters": {...}
},
"defer_loading": True, # Tool search
"input_examples": [ # Input examples
{"sql": "SELECT * FROM users WHERE id = 1"}
]
}
```
### With Programmatic Tool Calling
```python
{
"type": "function",
"function": {
"name": "fetch_data",
"description": "Fetch data from API",
"parameters": {...}
},
"allowed_callers": ["code_execution_20250825"], # Programmatic calling
"input_examples": [ # Input examples
{"endpoint": "/api/users", "method": "GET"}
]
}
```
### All Features Combined
```python
{
"type": "function",
"function": {
"name": "advanced_tool",
"description": "A complex tool",
"parameters": {...}
},
"defer_loading": True, # Tool search
"allowed_callers": ["code_execution_20250825"], # Programmatic calling
"input_examples": [ # Input examples
{"param1": "value1", "param2": "value2"}
]
}
```
## Provider Support
LiteLLM supports input examples across all Anthropic-compatible providers:
- **Standard Anthropic API** (`anthropic/claude-sonnet-4-5-20250929`)
- **Azure Anthropic** (`azure/claude-sonnet-4-5-20250929`)
- **Vertex AI Anthropic** (`vertex_ai/claude-sonnet-4-5-20250929`)
The beta header is automatically added when LiteLLM detects tools with `input_examples` field.
## Troubleshooting
### "Invalid request" error with examples
**Problem**: Receiving 400 error when using input examples
**Solution**: Ensure each example is valid according to your `input_schema`:
```python
# Check that:
# 1. All required fields are present in examples
# 2. Field types match the schema
# 3. Enum values are valid
# 4. Nested objects follow the schema structure
```
### Examples not improving tool calls
**Problem**: Adding examples doesn't seem to help
**Solution**:
1. **Check descriptions first**: Ensure tool descriptions are detailed and clear
2. **Review example quality**: Make sure examples are realistic and diverse
3. **Verify schema**: Confirm examples actually match your schema
4. **Add more variation**: Include examples showing different use cases
### Token usage too high
**Problem**: Input examples consuming too many tokens
**Solution**:
1. **Reduce example count**: Use 2-3 examples instead of 5+
2. **Simplify examples**: Remove unnecessary fields from examples
3. **Consider descriptions**: If descriptions are clear, examples may not be needed
## When NOT to Use Input Examples
Skip input examples if:
- **Tool is simple**: Single parameter tools with clear descriptions
- **Schema is self-explanatory**: Well-structured schema with good descriptions
- **Token budget is tight**: Examples add 20-200 tokens each
- **Server-side tools**: web_search, code_execution, etc. don't support examples
## Related Features
- [Anthropic Tool Search](./anthropic_tool_search.md) - Dynamically discover and load tools on-demand
- [Anthropic Programmatic Tool Calling](./anthropic_programmatic_tool_calling.md) - Call tools from code execution
- [Anthropic Provider](./anthropic.md) - General Anthropic provider documentation
docs/my-website/docs/providers/anthropic_tool_search.md
+397
View File
@@ -0,0 +1,397 @@
# Anthropic Tool Search
Tool search enables Claude to dynamically discover and load tools on-demand from large tool catalogs (10,000+ tools). Instead of loading all tool definitions into the context window upfront, Claude searches your tool catalog and loads only the tools it needs.
## Benefits
- **Context efficiency**: Avoid consuming massive portions of your context window with tool definitions
- **Better tool selection**: Claude's tool selection accuracy degrades with more than 30-50 tools. Tool search maintains accuracy even with thousands of tools
- **On-demand loading**: Tools are only loaded when Claude needs them
## Supported Models
Tool search is available on:
- Claude Opus 4.5
- Claude Sonnet 4.5
## Supported Platforms
- Anthropic API (direct)
- Azure Anthropic (Microsoft Foundry)
- Google Cloud Vertex AI
- Amazon Bedrock (invoke API only, not converse API)
## Tool Search Variants
LiteLLM supports both tool search variants:
### 1. Regex Tool Search (`tool_search_tool_regex_20251119`)
Claude constructs regex patterns to search for tools.
### 2. BM25 Tool Search (`tool_search_tool_bm25_20251119`)
Claude uses natural language queries to search for tools using the BM25 algorithm.
## Quick Start
### Basic Example with Regex Tool Search
```python
import litellm
response = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[
{"role": "user", "content": "What is the weather in San Francisco?"}
],
tools=[
# Tool search tool (regex variant)
{
"type": "tool_search_tool_regex_20251119",
"name": "tool_search_tool_regex"
},
# Deferred tool - will be loaded on-demand
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the weather at a specific location",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"]
}
},
"required": ["location"]
}
},
"defer_loading": True # Mark for deferred loading
},
# Another deferred tool
{
"type": "function",
"function": {
"name": "search_files",
"description": "Search through files in the workspace",
"parameters": {
"type": "object",
"properties": {
"query": {"type": "string"},
"file_types": {
"type": "array",
"items": {"type": "string"}
}
},
"required": ["query"]
}
},
"defer_loading": True
}
]
)
print(response.choices[0].message.content)
```
### BM25 Tool Search Example
```python
import litellm
response = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[
{"role": "user", "content": "Search for Python files containing 'authentication'"}
],
tools=[
# Tool search tool (BM25 variant)
{
"type": "tool_search_tool_bm25_20251119",
"name": "tool_search_tool_bm25"
},
# Deferred tools...
{
"type": "function",
"function": {
"name": "search_codebase",
"description": "Search through codebase files by content and filename",
"parameters": {
"type": "object",
"properties": {
"query": {"type": "string"},
"file_pattern": {"type": "string"}
},
"required": ["query"]
}
},
"defer_loading": True
}
]
)
```
## Using with Azure Anthropic
```python
import litellm
response = litellm.completion(
model="azure_anthropic/claude-sonnet-4-5",
api_base="https://<your-resource>.services.ai.azure.com/anthropic",
api_key="your-azure-api-key",
messages=[
{"role": "user", "content": "What's the weather like?"}
],
tools=[
{
"type": "tool_search_tool_regex_20251119",
"name": "tool_search_tool_regex"
},
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get current weather",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
},
"required": ["location"]
}
},
"defer_loading": True
}
]
)
```
## Using with Vertex AI
```python
import litellm
response = litellm.completion(
model="vertex_ai/claude-sonnet-4-5",
vertex_project="your-project-id",
vertex_location="us-central1",
messages=[
{"role": "user", "content": "Search my documents"}
],
tools=[
{
"type": "tool_search_tool_bm25_20251119",
"name": "tool_search_tool_bm25"
},
# Your deferred tools...
]
)
```
## Streaming Support
Tool search works with streaming:
```python
import litellm
response = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[
{"role": "user", "content": "Get the weather"}
],
tools=[
{
"type": "tool_search_tool_regex_20251119",
"name": "tool_search_tool_regex"
},
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather information",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
},
"required": ["location"]
}
},
"defer_loading": True
}
],
stream=True
)
for chunk in response:
if chunk.choices[0].delta.content:
print(chunk.choices[0].delta.content, end="")
```
## LiteLLM Proxy
Tool search works automatically through the LiteLLM proxy:
### Proxy Config
```yaml
model_list:
- model_name: claude-sonnet
litellm_params:
model: anthropic/claude-sonnet-4-5-20250929
api_key: os.environ/ANTHROPIC_API_KEY
```
### Client Request
```python
import openai
client = openai.OpenAI(
api_key="your-litellm-proxy-key",
base_url="http://0.0.0.0:4000"
)
response = client.chat.completions.create(
model="claude-sonnet",
messages=[
{"role": "user", "content": "What's the weather?"}
],
tools=[
{
"type": "tool_search_tool_regex_20251119",
"name": "tool_search_tool_regex"
},
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather information",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
},
"required": ["location"]
}
},
"defer_loading": True
}
]
)
```
## Important Notes
### Beta Header
LiteLLM automatically adds the `advanced-tool-use-2025-11-20` beta header when tool search tools are detected. You don't need to manually specify it.
### Deferred Loading
- Tools with `defer_loading: true` are only loaded when Claude discovers them via search
- At least one tool must be non-deferred (the tool search tool itself)
- Keep your 3-5 most frequently used tools as non-deferred for optimal performance
### Tool Descriptions
Write clear, descriptive tool names and descriptions that match how users describe tasks. The search algorithm uses:
- Tool names
- Tool descriptions
- Argument names
- Argument descriptions
### Usage Tracking
Tool search requests are tracked in the usage object:
```python
response = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[{"role": "user", "content": "Search for tools"}],
tools=[...]
)
# Check tool search usage
if response.usage.server_tool_use:
print(f"Tool search requests: {response.usage.server_tool_use.tool_search_requests}")
```
## Error Handling
### All Tools Deferred
```python
# ❌ This will fail - at least one tool must be non-deferred
tools = [
{
"type": "function",
"function": {...},
"defer_loading": True
}
]
# ✅ Correct - tool search tool is non-deferred
tools = [
{
"type": "tool_search_tool_regex_20251119",
"name": "tool_search_tool_regex"
},
{
"type": "function",
"function": {...},
"defer_loading": True
}
]
```
### Missing Tool Definition
If Claude references a tool that isn't in your deferred tools list, you'll get an error. Make sure all tools that might be discovered are included in the tools parameter with `defer_loading: true`.
## Best Practices
1. **Keep frequently used tools non-deferred**: Your 3-5 most common tools should not have `defer_loading: true`
2. **Use semantic descriptions**: Tool descriptions should use natural language that matches user queries
3. **Choose the right variant**:
- Use **regex** for exact pattern matching (faster)
- Use **BM25** for natural language semantic search
4. **Monitor usage**: Track `tool_search_requests` in the usage object to understand search patterns
5. **Optimize tool catalog**: Remove unused tools and consolidate similar functionality
## When to Use Tool Search
**Good use cases:**
- 10+ tools available in your system
- Tool definitions consuming >10K tokens
- Experiencing tool selection accuracy issues
- Building systems with multiple tool categories
- Tool library growing over time
**When traditional tool calling is better:**
- Less than 10 tools total
- All tools are frequently used
- Very small tool definitions (\<100 tokens total)
## Limitations
- Not compatible with tool use examples
- Requires Claude Opus 4.5 or Sonnet 4.5
- On Bedrock, only available via invoke API (not converse API)
- Maximum 10,000 tools in catalog
- Returns 3-5 most relevant tools per search
## Additional Resources
- [Anthropic Tool Search Documentation](https://docs.anthropic.com/en/docs/build-with-claude/tool-use/tool-search)
- [LiteLLM Tool Calling Guide](https://docs.litellm.ai/docs/completion/function_call)
docs/my-website/docs/providers/bedrock.md
+1 -201
View File
@@ -7,7 +7,7 @@ ALL Bedrock models (Anthropic, Meta, Deepseek, Mistral, Amazon, etc.) are Suppor
| Property | Details |
|-------|-------|
| Description | Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs). |
| Provider Route on LiteLLM | `bedrock/`, [`bedrock/converse/`](#set-converse--invoke-route), [`bedrock/invoke/`](#set-invoke-route), [`bedrock/converse_like/`](#calling-via-internal-proxy), [`bedrock/llama/`](#deepseek-not-r1), [`bedrock/deepseek_r1/`](#deepseek-r1), [`bedrock/qwen3/`](#qwen3-imported-models) |
| Provider Route on LiteLLM | `bedrock/`, [`bedrock/converse/`](#set-converse--invoke-route), [`bedrock/invoke/`](#set-invoke-route), [`bedrock/converse_like/`](#calling-via-internal-proxy), [`bedrock/llama/`](#deepseek-not-r1), [`bedrock/deepseek_r1/`](#deepseek-r1), [`bedrock/qwen3/`](#qwen3-imported-models), [`bedrock/openai/`](./bedrock_imported.md#openai-compatible-imported-models-qwen-25-vl-etc) |
| Provider Doc | [Amazon Bedrock ↗](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) |
| Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/embeddings`, `/images/generations` |
| Rerank Endpoint | `/rerank` |
@@ -1598,206 +1598,6 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
</Tabs>
## Bedrock Imported Models (Deepseek, Deepseek R1)
### Deepseek R1
This is a separate route, as the chat template is different.
| Property | Details |
|----------|---------|
| Provider Route | `bedrock/deepseek_r1/{model_arn}` |
| Provider Documentation | [Bedrock Imported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html), [Deepseek Bedrock Imported Model](https://aws.amazon.com/blogs/machine-learning/deploy-deepseek-r1-distilled-llama-models-with-amazon-bedrock-custom-model-import/) |
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
import os
response = completion(
model="bedrock/deepseek_r1/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n", # bedrock/deepseek_r1/{your-model-arn}
messages=[{"role": "user", "content": "Tell me a joke"}],
)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml
model_list:
- model_name: DeepSeek-R1-Distill-Llama-70B
litellm_params:
model: bedrock/deepseek_r1/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n
```
**2. Start proxy**
```bash
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "DeepSeek-R1-Distill-Llama-70B", # 👈 the 'model_name' in config
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
</TabItem>
</Tabs>
### Deepseek (not R1)
| Property | Details |
|----------|---------|
| Provider Route | `bedrock/llama/{model_arn}` |
| Provider Documentation | [Bedrock Imported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html), [Deepseek Bedrock Imported Model](https://aws.amazon.com/blogs/machine-learning/deploy-deepseek-r1-distilled-llama-models-with-amazon-bedrock-custom-model-import/) |
Use this route to call Bedrock Imported Models that follow the `llama` Invoke Request / Response spec
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
import os
response = completion(
model="bedrock/llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n", # bedrock/llama/{your-model-arn}
messages=[{"role": "user", "content": "Tell me a joke"}],
)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml
model_list:
- model_name: DeepSeek-R1-Distill-Llama-70B
litellm_params:
model: bedrock/llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n
```
**2. Start proxy**
```bash
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "DeepSeek-R1-Distill-Llama-70B", # 👈 the 'model_name' in config
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
</TabItem>
</Tabs>
### Qwen3 Imported Models
| Property | Details |
|----------|---------|
| Provider Route | `bedrock/qwen3/{model_arn}` |
| Provider Documentation | [Bedrock Imported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html), [Qwen3 Models](https://aws.amazon.com/about-aws/whats-new/2025/09/qwen3-models-fully-managed-amazon-bedrock/) |
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
import os
response = completion(
model="bedrock/qwen3/arn:aws:bedrock:us-east-1:086734376398:imported-model/your-qwen3-model", # bedrock/qwen3/{your-model-arn}
messages=[{"role": "user", "content": "Tell me a joke"}],
max_tokens=100,
temperature=0.7
)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml
model_list:
- model_name: Qwen3-32B
litellm_params:
model: bedrock/qwen3/arn:aws:bedrock:us-east-1:086734376398:imported-model/your-qwen3-model
```
**2. Start proxy**
```bash
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "Qwen3-32B", # 👈 the 'model_name' in config
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
</TabItem>
</Tabs>
### OpenAI GPT OSS
| Property | Details |
docs/my-website/docs/providers/bedrock_imported.md
+369
View File
@@ -0,0 +1,369 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Bedrock Imported Models
Bedrock Imported Models (Deepseek, Deepseek R1, Qwen, OpenAI-compatible models)
### Deepseek R1
This is a separate route, as the chat template is different.
| Property | Details |
|----------|---------|
| Provider Route | `bedrock/deepseek_r1/{model_arn}` |
| Provider Documentation | [Bedrock Imported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html), [Deepseek Bedrock Imported Model](https://aws.amazon.com/blogs/machine-learning/deploy-deepseek-r1-distilled-llama-models-with-amazon-bedrock-custom-model-import/) |
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
import os
response = completion(
model="bedrock/deepseek_r1/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n", # bedrock/deepseek_r1/{your-model-arn}
messages=[{"role": "user", "content": "Tell me a joke"}],
)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml
model_list:
- model_name: DeepSeek-R1-Distill-Llama-70B
litellm_params:
model: bedrock/deepseek_r1/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n
```
**2. Start proxy**
```bash
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "DeepSeek-R1-Distill-Llama-70B", # 👈 the 'model_name' in config
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
</TabItem>
</Tabs>
### Deepseek (not R1)
| Property | Details |
|----------|---------|
| Provider Route | `bedrock/llama/{model_arn}` |
| Provider Documentation | [Bedrock Imported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html), [Deepseek Bedrock Imported Model](https://aws.amazon.com/blogs/machine-learning/deploy-deepseek-r1-distilled-llama-models-with-amazon-bedrock-custom-model-import/) |
Use this route to call Bedrock Imported Models that follow the `llama` Invoke Request / Response spec
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
import os
response = completion(
model="bedrock/llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n", # bedrock/llama/{your-model-arn}
messages=[{"role": "user", "content": "Tell me a joke"}],
)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml
model_list:
- model_name: DeepSeek-R1-Distill-Llama-70B
litellm_params:
model: bedrock/llama/arn:aws:bedrock:us-east-1:086734376398:imported-model/r4c4kewx2s0n
```
**2. Start proxy**
```bash
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "DeepSeek-R1-Distill-Llama-70B", # 👈 the 'model_name' in config
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
</TabItem>
</Tabs>
### Qwen3 Imported Models
| Property | Details |
|----------|---------|
| Provider Route | `bedrock/qwen3/{model_arn}` |
| Provider Documentation | [Bedrock Imported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html), [Qwen3 Models](https://aws.amazon.com/about-aws/whats-new/2025/09/qwen3-models-fully-managed-amazon-bedrock/) |
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
import os
response = completion(
model="bedrock/qwen3/arn:aws:bedrock:us-east-1:086734376398:imported-model/your-qwen3-model", # bedrock/qwen3/{your-model-arn}
messages=[{"role": "user", "content": "Tell me a joke"}],
max_tokens=100,
temperature=0.7
)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml
model_list:
- model_name: Qwen3-32B
litellm_params:
model: bedrock/qwen3/arn:aws:bedrock:us-east-1:086734376398:imported-model/your-qwen3-model
```
**2. Start proxy**
```bash
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "Qwen3-32B", # 👈 the 'model_name' in config
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
</TabItem>
</Tabs>
### OpenAI-Compatible Imported Models (Qwen 2.5 VL, etc.)
Use this route for Bedrock imported models that follow the **OpenAI Chat Completions API spec**. This includes models like Qwen 2.5 VL that accept OpenAI-formatted messages with support for vision (images), tool calling, and other OpenAI features.
| Property | Details |
|----------|---------|
| Provider Route | `bedrock/openai/{model_arn}` |
| Provider Documentation | [Bedrock Imported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html) |
| Supported Features | Vision (images), tool calling, streaming, system messages |
#### LiteLLMSDK Usage
**Basic Usage**
```python
from litellm import completion
response = completion(
model="bedrock/openai/arn:aws:bedrock:us-east-1:046319184608:imported-model/0m2lasirsp6z", # bedrock/openai/{your-model-arn}
messages=[{"role": "user", "content": "Tell me a joke"}],
max_tokens=300,
temperature=0.5
)
```
**With Vision (Images)**
```python
import base64
from litellm import completion
# Load and encode image
with open("image.jpg", "rb") as f:
image_base64 = base64.b64encode(f.read()).decode("utf-8")
response = completion(
model="bedrock/openai/arn:aws:bedrock:us-east-1:046319184608:imported-model/0m2lasirsp6z",
messages=[
{
"role": "system",
"content": "You are a helpful assistant that can analyze images."
},
{
"role": "user",
"content": [
{"type": "text", "text": "What's in this image?"},
{
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}
}
]
}
],
max_tokens=300,
temperature=0.5
)
```
**Comparing Multiple Images**
```python
import base64
from litellm import completion
# Load images
with open("image1.jpg", "rb") as f:
image1_base64 = base64.b64encode(f.read()).decode("utf-8")
with open("image2.jpg", "rb") as f:
image2_base64 = base64.b64encode(f.read()).decode("utf-8")
response = completion(
model="bedrock/openai/arn:aws:bedrock:us-east-1:046319184608:imported-model/0m2lasirsp6z",
messages=[
{
"role": "system",
"content": "You are a helpful assistant that can analyze images."
},
{
"role": "user",
"content": [
{"type": "text", "text": "Spot the difference between these two images?"},
{
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{image1_base64}"}
},
{
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{image2_base64}"}
}
]
}
],
max_tokens=300,
temperature=0.5
)
```
#### LiteLLM Proxy Usage (AI Gateway)
**1. Add to config**
```yaml
model_list:
- model_name: qwen-25vl-72b
litellm_params:
model: bedrock/openai/arn:aws:bedrock:us-east-1:046319184608:imported-model/0m2lasirsp6z
```
**2. Start proxy**
```bash
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
Basic text request:
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "qwen-25vl-72b",
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
"max_tokens": 300
}'
```
With vision (image):
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "qwen-25vl-72b",
"messages": [
{
"role": "system",
"content": "You are a helpful assistant that can analyze images."
},
{
"role": "user",
"content": [
{"type": "text", "text": "What is in this image?"},
{
"type": "image_url",
"image_url": {"url": "data:image/jpeg;base64,/9j/4AAQSkZ..."}
}
]
}
],
"max_tokens": 300,
"temperature": 0.5
}'
```
docs/my-website/docs/proxy/guardrails/pillar_security.md
+3 -1
View File
@@ -60,6 +60,8 @@ litellm_settings:
set_verbose: true # Enable detailed logging
```
**Note:** Virtual key context is **automatically passed** as headers - no additional configuration needed!
### 3. Start the Proxy
```bash
@@ -210,7 +212,7 @@ export PILLAR_API_KEY="your_api_key_here"
export PILLAR_API_BASE="https://api.pillar.security"
export PILLAR_ON_FLAGGED_ACTION="monitor"
export PILLAR_FALLBACK_ON_ERROR="allow"
export PILLAR_TIMEOUT="30.0"
export PILLAR_TIMEOUT="5.0"
```
### Session Tracking
docs/my-website/docs/rag_ingest.md
+227
View File
@@ -0,0 +1,227 @@
# /rag/ingest
All-in-one document ingestion pipeline: **Upload → Chunk → Embed → Vector Store**
| Feature | Supported |
|---------|-----------|
| Cost Tracking | ❌ |
| Logging | ✅ |
| Supported Providers | `openai`, `bedrock` |
## Quick Start
### OpenAI
```bash showLineNumbers title="Ingest to OpenAI vector store"
curl -X POST "http://localhost:4000/v1/rag/ingest" \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d "{
\"file\": {
\"filename\": \"document.txt\",
\"content\": \"$(base64 -i document.txt)\",
\"content_type\": \"text/plain\"
},
\"ingest_options\": {
\"vector_store\": {
\"custom_llm_provider\": \"openai\"
}
}
}"
```
### Bedrock
```bash showLineNumbers title="Ingest to Bedrock Knowledge Base"
curl -X POST "http://localhost:4000/v1/rag/ingest" \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d "{
\"file\": {
\"filename\": \"document.txt\",
\"content\": \"$(base64 -i document.txt)\",
\"content_type\": \"text/plain\"
},
\"ingest_options\": {
\"vector_store\": {
\"custom_llm_provider\": \"bedrock\"
}
}
}"
```
## Response
```json
{
"id": "ingest_abc123",
"status": "completed",
"vector_store_id": "vs_xyz789",
"file_id": "file_123"
}
```
## Query the Vector Store
After ingestion, query with `/vector_stores/{vector_store_id}/search`:
```bash showLineNumbers title="Search the vector store"
curl -X POST "http://localhost:4000/v1/vector_stores/vs_xyz789/search" \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"query": "What is the main topic?",
"max_num_results": 5
}'
```
## End-to-End Example
### OpenAI
#### 1. Ingest Document
```bash showLineNumbers title="Step 1: Ingest"
curl -X POST "http://localhost:4000/v1/rag/ingest" \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d "{
\"file\": {
\"filename\": \"test_document.txt\",
\"content\": \"$(base64 -i test_document.txt)\",
\"content_type\": \"text/plain\"
},
\"ingest_options\": {
\"name\": \"test-basic-ingest\",
\"vector_store\": {
\"custom_llm_provider\": \"openai\"
}
}
}"
```
Response:
```json
{
"id": "ingest_d834f544-fc5e-4751-902d-fb0bcc183b85",
"status": "completed",
"vector_store_id": "vs_692658d337c4819183f2ad8488d12fc9",
"file_id": "file-M2pJJiWH56cfUP4Fe7rJay"
}
```
#### 2. Query
```bash showLineNumbers title="Step 2: Query"
curl -X POST "http://localhost:4000/v1/vector_stores/vs_692658d337c4819183f2ad8488d12fc9/search" \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"query": "What is LiteLLM?",
"custom_llm_provider": "openai"
}'
```
Response:
```json
{
"object": "vector_store.search_results.page",
"search_query": ["What is LiteLLM?"],
"data": [
{
"file_id": "file-M2pJJiWH56cfUP4Fe7rJay",
"filename": "test_document.txt",
"score": 0.4004629778869299,
"attributes": {},
"content": [
{
"type": "text",
"text": "Test document abc123 for RAG ingestion.\nThis is a sample document to test the RAG ingest API.\nLiteLLM provides a unified interface for vector stores."
}
]
}
],
"has_more": false,
"next_page": null
}
```
## Request Parameters
### Top-Level
| Parameter | Type | Required | Description |
|-----------|------|----------|-------------|
| `file` | object | One of file/file_url/file_id required | Base64-encoded file |
| `file.filename` | string | Yes | Filename with extension |
| `file.content` | string | Yes | Base64-encoded content |
| `file.content_type` | string | Yes | MIME type (e.g., `text/plain`) |
| `file_url` | string | One of file/file_url/file_id required | URL to fetch file from |
| `file_id` | string | One of file/file_url/file_id required | Existing file ID |
| `ingest_options` | object | Yes | Pipeline configuration |
### ingest_options
| Parameter | Type | Required | Description |
|-----------|------|----------|-------------|
| `vector_store` | object | Yes | Vector store configuration |
| `name` | string | No | Pipeline name for logging |
### vector_store (OpenAI)
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `custom_llm_provider` | string | - | `"openai"` |
| `vector_store_id` | string | auto-create | Existing vector store ID |
### vector_store (Bedrock)
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `custom_llm_provider` | string | - | `"bedrock"` |
| `vector_store_id` | string | auto-create | Existing Knowledge Base ID |
| `wait_for_ingestion` | boolean | `false` | Wait for indexing to complete |
| `ingestion_timeout` | integer | `300` | Timeout in seconds (if waiting) |
| `s3_bucket` | string | auto-create | S3 bucket for documents |
| `s3_prefix` | string | `"data/"` | S3 key prefix |
| `embedding_model` | string | `amazon.titan-embed-text-v2:0` | Bedrock embedding model |
| `aws_region_name` | string | `us-west-2` | AWS region |
:::info Bedrock Auto-Creation
When `vector_store_id` is omitted, LiteLLM automatically creates:
- S3 bucket for document storage
- OpenSearch Serverless collection
- IAM role with required permissions
- Bedrock Knowledge Base
- Data Source
:::
## Input Examples
### File (Base64)
```json title="Request body"
{
"file": {
"filename": "document.txt",
"content": "<base64-encoded-content>",
"content_type": "text/plain"
},
"ingest_options": {
"vector_store": {"custom_llm_provider": "openai"}
}
}
```
### File URL
```bash showLineNumbers title="Ingest from URL"
curl -X POST "http://localhost:4000/v1/rag/ingest" \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"file_url": "https://example.com/document.pdf",
"ingest_options": {"vector_store": {"custom_llm_provider": "openai"}}
}'
```
docs/my-website/sidebars.js
+2
View File
@@ -412,6 +412,7 @@ const sidebars = {
"proxy/pass_through"
]
},
"rag_ingest",
"realtime",
"rerank",
"response_api",
@@ -530,6 +531,7 @@ const sidebars = {
items: [
"providers/bedrock",
"providers/bedrock_embedding",
"providers/bedrock_imported",
"providers/bedrock_image_gen",
"providers/bedrock_rerank",
"providers/bedrock_agentcore",
litellm/__init__.py
+7
View File
@@ -1225,6 +1225,9 @@ from .llms.bedrock.chat.invoke_transformations.amazon_titan_transformation impor
from .llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
AmazonInvokeConfig,
)
from .llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import (
AmazonBedrockOpenAIConfig,
)
from .llms.bedrock.image.amazon_stability1_transformation import AmazonStabilityConfig
from .llms.bedrock.image.amazon_stability3_transformation import AmazonStability3Config
@@ -1431,6 +1434,7 @@ from .skills.main import (
)
from .containers.main import *
from .ocr.main import *
from .rag.main import *
from .search.main import *
from .realtime_api.main import _arealtime
from .fine_tuning.main import *
@@ -1467,6 +1471,9 @@ from .vector_stores.vector_store_registry import (
vector_store_registry: Optional[VectorStoreRegistry] = None
vector_store_index_registry: Optional[VectorStoreIndexRegistry] = None
### RAG ###
from . import rag
### CUSTOM LLMs ###
from .types.llms.custom_llm import CustomLLMItem
from .types.utils import GenericStreamingChunk
litellm/completion_extras/litellm_responses_transformation/transformation.py
+62
View File
@@ -234,6 +234,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
cast(List[Dict[str, Any]], value)
)
)
elif key == "response_format":
# Convert response_format to text.format
text_format = self._transform_response_format_to_text_format(value)
if text_format:
responses_api_request["text"] = text_format # type: ignore
elif key in ResponsesAPIOptionalRequestParams.__annotations__.keys():
responses_api_request[key] = value # type: ignore
elif key == "metadata":
@@ -666,6 +671,63 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return Reasoning(effort="minimal")
return None
def _transform_response_format_to_text_format(
self, response_format: Union[Dict[str, Any], Any]
) -> Optional[Dict[str, Any]]:
"""
Transform Chat Completion response_format parameter to Responses API text.format parameter.
Chat Completion response_format structure:
{
"type": "json_schema",
"json_schema": {
"name": "schema_name",
"schema": {...},
"strict": True
}
}
Responses API text parameter structure:
{
"format": {
"type": "json_schema",
"name": "schema_name",
"schema": {...},
"strict": True
}
}
"""
if not response_format:
return None
if isinstance(response_format, dict):
format_type = response_format.get("type")
if format_type == "json_schema":
json_schema = response_format.get("json_schema", {})
return {
"format": {
"type": "json_schema",
"name": json_schema.get("name", "response_schema"),
"schema": json_schema.get("schema", {}),
"strict": json_schema.get("strict", False),
}
}
elif format_type == "json_object":
return {
"format": {
"type": "json_object"
}
}
elif format_type == "text":
return {
"format": {
"type": "text"
}
}
return None
def _map_responses_status_to_finish_reason(self, status: Optional[str]) -> str:
"""Map responses API status to chat completion finish_reason"""
if not status:
litellm/constants.py
+4
View File
@@ -1211,3 +1211,7 @@ SENTRY_PII_DENYLIST = [
COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY = int(
os.getenv("COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY", 1000)
)
########################### RAG Text Splitter Constants ###########################
DEFAULT_CHUNK_SIZE = int(os.getenv("DEFAULT_CHUNK_SIZE", 1000))
DEFAULT_CHUNK_OVERLAP = int(os.getenv("DEFAULT_CHUNK_OVERLAP", 200))
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
+11 -1
View File
@@ -18,6 +18,7 @@ from litellm.types.utils import (
ModelResponseStream,
PromptTokensDetailsWrapper,
Usage,
ServerToolUse
)
from litellm.utils import print_verbose, token_counter
@@ -418,7 +419,8 @@ class ChunkProcessor:
## anthropic prompt caching information ##
cache_creation_input_tokens: Optional[int] = None
cache_read_input_tokens: Optional[int] = None
server_tool_use: Optional[ServerToolUse] = None
web_search_requests: Optional[int] = None
completion_tokens_details: Optional[CompletionTokensDetails] = None
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
@@ -462,6 +464,8 @@ class ChunkProcessor:
completion_tokens_details = usage_chunk_dict[
"completion_tokens_details"
]
if hasattr(usage_chunk, 'server_tool_use') and usage_chunk.server_tool_use is not None:
server_tool_use = usage_chunk.server_tool_use
if (
usage_chunk_dict["prompt_tokens_details"] is not None
and getattr(
@@ -483,6 +487,7 @@ class ChunkProcessor:
completion_tokens=completion_tokens,
cache_creation_input_tokens=cache_creation_input_tokens,
cache_read_input_tokens=cache_read_input_tokens,
server_tool_use=server_tool_use,
web_search_requests=web_search_requests,
completion_tokens_details=completion_tokens_details,
prompt_tokens_details=prompt_tokens_details,
@@ -513,6 +518,9 @@ class ChunkProcessor:
"cache_read_input_tokens"
]
server_tool_use: Optional[ServerToolUse] = calculated_usage_per_chunk[
"server_tool_use"
]
web_search_requests: Optional[int] = calculated_usage_per_chunk[
"web_search_requests"
]
@@ -576,6 +584,8 @@ class ChunkProcessor:
if prompt_tokens_details is not None:
returned_usage.prompt_tokens_details = prompt_tokens_details
if server_tool_use is not None:
returned_usage.server_tool_use = server_tool_use
if web_search_requests is not None:
if returned_usage.prompt_tokens_details is None:
returned_usage.prompt_tokens_details = PromptTokensDetailsWrapper(
litellm/llms/anthropic/chat/handler.py
+53 -28
View File
@@ -42,6 +42,7 @@ from litellm.types.llms.openai import (
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
ChatCompletionToolCallChunk,
ChatCompletionToolCallFunctionChunk,
)
from litellm.types.utils import (
Delta,
@@ -550,15 +551,18 @@ class ModelResponseIterator:
if "text" in content_block["delta"]:
text = content_block["delta"]["text"]
elif "partial_json" in content_block["delta"]:
tool_use = {
"id": None,
"type": "function",
"function": {
"name": None,
"arguments": content_block["delta"]["partial_json"],
tool_use = cast(
ChatCompletionToolCallChunk,
{
"id": None,
"type": "function",
"function": {
"name": None,
"arguments": content_block["delta"]["partial_json"],
},
"index": self.tool_index,
},
"index": self.tool_index,
}
)
elif "citation" in content_block["delta"]:
provider_specific_fields["citation"] = content_block["delta"]["citation"]
elif (
@@ -569,7 +573,7 @@ class ModelResponseIterator:
ChatCompletionThinkingBlock(
type="thinking",
thinking=content_block["delta"].get("thinking") or "",
signature=content_block["delta"].get("signature"),
signature=str(content_block["delta"].get("signature") or ""),
)
]
provider_specific_fields["thinking_blocks"] = thinking_blocks
@@ -625,7 +629,7 @@ class ModelResponseIterator:
return content_block_start
def chunk_parser(self, chunk: dict) -> ModelResponseStream:
def chunk_parser(self, chunk: dict) -> ModelResponseStream: # noqa: PLR0915
try:
type_chunk = chunk.get("type", "") or ""
@@ -672,15 +676,32 @@ class ModelResponseIterator:
text = content_block_start["content_block"]["text"]
elif content_block_start["content_block"]["type"] == "tool_use":
self.tool_index += 1
tool_use = {
"id": content_block_start["content_block"]["id"],
"type": "function",
"function": {
"name": content_block_start["content_block"]["name"],
"arguments": "",
},
"index": self.tool_index,
}
tool_use = ChatCompletionToolCallChunk(
id=content_block_start["content_block"]["id"],
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=content_block_start["content_block"]["name"],
arguments="",
),
index=self.tool_index,
)
# Include caller information if present (for programmatic tool calling)
if "caller" in content_block_start["content_block"]:
caller_data = content_block_start["content_block"]["caller"]
if caller_data:
tool_use["caller"] = cast(Dict[str, Any], caller_data) # type: ignore[typeddict-item]
elif content_block_start["content_block"]["type"] == "server_tool_use":
# Handle server tool use (for tool search)
self.tool_index += 1
tool_use = ChatCompletionToolCallChunk(
id=content_block_start["content_block"]["id"],
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=content_block_start["content_block"]["name"],
arguments="",
),
index=self.tool_index,
)
elif (
content_block_start["content_block"]["type"] == "redacted_thinking"
):
@@ -696,17 +717,21 @@ class ModelResponseIterator:
# check if tool call content block
is_empty = self.check_empty_tool_call_args()
if is_empty:
tool_use = {
"id": None,
"type": "function",
"function": {
"name": None,
"arguments": "{}",
},
"index": self.tool_index,
}
tool_use = ChatCompletionToolCallChunk(
id=None, # type: ignore[typeddict-item]
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=None, # type: ignore[typeddict-item]
arguments="{}",
),
index=self.tool_index,
)
# Reset response_format tool tracking when block stops
self.is_response_format_tool = False
elif type_chunk == "tool_result":
# Handle tool_result blocks (for tool search results with tool_reference)
# These are automatically handled by Anthropic API, we just pass them through
pass
elif type_chunk == "message_delta":
finish_reason, usage = self._handle_message_delta(chunk)
elif type_chunk == "message_start":
litellm/llms/anthropic/chat/transformation.py
+240 -26
View File
@@ -54,7 +54,10 @@ from litellm.types.utils import (
CompletionTokensDetailsWrapper,
)
from litellm.types.utils import Message as LitellmMessage
from litellm.types.utils import PromptTokensDetailsWrapper, ServerToolUse
from litellm.types.utils import (
PromptTokensDetailsWrapper,
ServerToolUse,
)
from litellm.utils import (
ModelResponse,
Usage,
@@ -187,7 +190,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
return _tool_choice
def _map_tool_helper(
def _map_tool_helper( # noqa: PLR0915
self, tool: ChatCompletionToolParam
) -> Tuple[Optional[AllAnthropicToolsValues], Optional[AnthropicMcpServerTool]]:
returned_tool: Optional[AllAnthropicToolsValues] = None
@@ -250,9 +253,10 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
returned_tool = _computer_tool
elif any(tool["type"].startswith(t) for t in ANTHROPIC_HOSTED_TOOLS):
function_name = tool.get("name", tool.get("function", {}).get("name"))
if function_name is None or not isinstance(function_name, str):
function_name_obj = tool.get("name", tool.get("function", {}).get("name"))
if function_name_obj is None or not isinstance(function_name_obj, str):
raise ValueError("Missing required parameter: name")
function_name = function_name_obj
additional_tool_params = {}
for k, v in tool.items():
@@ -268,6 +272,30 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
mcp_server = self._map_openai_mcp_server_tool(
cast(OpenAIMcpServerTool, tool)
)
elif tool["type"] == "tool_search_tool_regex_20251119":
# Tool search tool using regex
from litellm.types.llms.anthropic import AnthropicToolSearchToolRegex
tool_name_obj = tool.get("name", "tool_search_tool_regex")
if not isinstance(tool_name_obj, str):
raise ValueError("Tool search tool must have a valid name")
tool_name = tool_name_obj
returned_tool = AnthropicToolSearchToolRegex(
type="tool_search_tool_regex_20251119",
name=tool_name,
)
elif tool["type"] == "tool_search_tool_bm25_20251119":
# Tool search tool using BM25
from litellm.types.llms.anthropic import AnthropicToolSearchToolBM25
tool_name_obj = tool.get("name", "tool_search_tool_bm25")
if not isinstance(tool_name_obj, str):
raise ValueError("Tool search tool must have a valid name")
tool_name = tool_name_obj
returned_tool = AnthropicToolSearchToolBM25(
type="tool_search_tool_bm25_20251119",
name=tool_name,
)
if returned_tool is None and mcp_server is None:
raise ValueError(f"Unsupported tool type: {tool['type']}")
@@ -275,14 +303,67 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
_cache_control = tool.get("cache_control", None)
_cache_control_function = tool.get("function", {}).get("cache_control", None)
if returned_tool is not None:
if _cache_control is not None:
returned_tool["cache_control"] = _cache_control
elif _cache_control_function is not None and isinstance(
_cache_control_function, dict
):
returned_tool["cache_control"] = ChatCompletionCachedContent(
**_cache_control_function # type: ignore
)
# Only set cache_control on tools that support it (not tool search tools)
tool_type = returned_tool.get("type", "")
if tool_type not in ("tool_search_tool_regex_20251119", "tool_search_tool_bm25_20251119"):
if _cache_control is not None:
returned_tool["cache_control"] = _cache_control # type: ignore[typeddict-item]
elif _cache_control_function is not None and isinstance(
_cache_control_function, dict
):
returned_tool["cache_control"] = ChatCompletionCachedContent( # type: ignore[typeddict-item]
**_cache_control_function # type: ignore
)
## check if defer_loading is set in the tool
_defer_loading = tool.get("defer_loading", None)
_defer_loading_function = tool.get("function", {}).get("defer_loading", None)
if returned_tool is not None:
# Only set defer_loading on tools that support it (not tool search tools or computer tools)
tool_type = returned_tool.get("type", "")
if tool_type not in ("tool_search_tool_regex_20251119", "tool_search_tool_bm25_20251119", "computer_20241022", "computer_20250124"):
if _defer_loading is not None:
if not isinstance(_defer_loading, bool):
raise ValueError("defer_loading must be a boolean")
returned_tool["defer_loading"] = _defer_loading # type: ignore[typeddict-item]
elif _defer_loading_function is not None:
if not isinstance(_defer_loading_function, bool):
raise ValueError("defer_loading must be a boolean")
returned_tool["defer_loading"] = _defer_loading_function # type: ignore[typeddict-item]
## check if allowed_callers is set in the tool
_allowed_callers = tool.get("allowed_callers", None)
_allowed_callers_function = tool.get("function", {}).get("allowed_callers", None)
if returned_tool is not None:
# Only set allowed_callers on tools that support it (not tool search tools or computer tools)
tool_type = returned_tool.get("type", "")
if tool_type not in ("tool_search_tool_regex_20251119", "tool_search_tool_bm25_20251119", "computer_20241022", "computer_20250124"):
if _allowed_callers is not None:
if not isinstance(_allowed_callers, list) or not all(
isinstance(item, str) for item in _allowed_callers
):
raise ValueError("allowed_callers must be a list of strings")
returned_tool["allowed_callers"] = _allowed_callers # type: ignore[typeddict-item]
elif _allowed_callers_function is not None:
if not isinstance(_allowed_callers_function, list) or not all(
isinstance(item, str) for item in _allowed_callers_function
):
raise ValueError("allowed_callers must be a list of strings")
returned_tool["allowed_callers"] = _allowed_callers_function # type: ignore[typeddict-item]
## check if input_examples is set in the tool
_input_examples = tool.get("input_examples", None)
_input_examples_function = tool.get("function", {}).get("input_examples", None)
if returned_tool is not None:
# Only set input_examples on user-defined tools (type "custom" or no type)
tool_type = returned_tool.get("type", "")
if tool_type == "custom" or (tool_type == "" and "name" in returned_tool):
if _input_examples is not None and isinstance(_input_examples, list):
returned_tool["input_examples"] = _input_examples # type: ignore[typeddict-item]
elif _input_examples_function is not None and isinstance(
_input_examples_function, list
):
returned_tool["input_examples"] = _input_examples_function # type: ignore[typeddict-item]
return returned_tool, mcp_server
@@ -334,6 +415,82 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
mcp_servers.append(mcp_server_tool)
return anthropic_tools, mcp_servers
def _detect_tool_search_tools(self, tools: Optional[List]) -> bool:
"""Check if tool search tools are present in the tools list."""
if not tools:
return False
for tool in tools:
tool_type = tool.get("type", "")
if tool_type in ["tool_search_tool_regex_20251119", "tool_search_tool_bm25_20251119"]:
return True
return False
def _separate_deferred_tools(
self, tools: List
) -> Tuple[List, List]:
"""
Separate tools into deferred and non-deferred lists.
Returns:
Tuple of (non_deferred_tools, deferred_tools)
"""
non_deferred = []
deferred = []
for tool in tools:
if tool.get("defer_loading", False):
deferred.append(tool)
else:
non_deferred.append(tool)
return non_deferred, deferred
def _expand_tool_references(
self,
content: List,
deferred_tools: List,
) -> List:
"""
Expand tool_reference blocks to full tool definitions.
When Anthropic's tool search returns results, it includes tool_reference blocks
that reference tools by name. This method expands those references to full
tool definitions from the deferred_tools catalog.
Args:
content: Response content that may contain tool_reference blocks
deferred_tools: List of deferred tools that can be referenced
Returns:
Content with tool_reference blocks expanded to full tool definitions
"""
if not deferred_tools:
return content
# Create a mapping of tool names to tool definitions
tool_map = {}
for tool in deferred_tools:
tool_name = tool.get("name") or tool.get("function", {}).get("name")
if tool_name:
tool_map[tool_name] = tool
# Expand tool references in content
expanded_content = []
for item in content:
if isinstance(item, dict) and item.get("type") == "tool_reference":
tool_name = item.get("tool_name")
if tool_name and tool_name in tool_map:
# Replace reference with full tool definition
expanded_content.append(tool_map[tool_name])
else:
# Keep the reference if we can't find the tool
expanded_content.append(item)
else:
expanded_content.append(item)
return expanded_content
def _map_stop_sequences(
self, stop: Optional[Union[str, List[str]]]
) -> Optional[List[str]]:
@@ -822,6 +979,17 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
"messages": anthropic_messages,
**optional_params,
}
## Handle output_config (Anthropic-specific parameter)
if "output_config" in optional_params:
output_config = optional_params.get("output_config")
if output_config and isinstance(output_config, dict):
effort = output_config.get("effort")
if effort and effort not in ["high", "medium", "low"]:
raise ValueError(
f"Invalid effort value: {effort}. Must be one of: 'high', 'medium', 'low'"
)
data["output_config"] = output_config
return data
@@ -870,18 +1038,40 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
text_content += content["text"]
## TOOL CALLING
elif content["type"] == "tool_use":
tool_calls.append(
ChatCompletionToolCallChunk(
id=content["id"],
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=content["name"],
arguments=json.dumps(content["input"]),
),
index=idx,
)
tool_call = ChatCompletionToolCallChunk(
id=content["id"],
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=content["name"],
arguments=json.dumps(content["input"]),
),
index=idx,
)
# Include caller information if present (for programmatic tool calling)
if "caller" in content:
tool_call["caller"] = cast(Dict[str, Any], content["caller"]) # type: ignore[typeddict-item]
tool_calls.append(tool_call)
## SERVER TOOL USE (for tool search)
elif content["type"] == "server_tool_use":
# Server tool use blocks are for tool search - treat as tool calls
tool_call = ChatCompletionToolCallChunk(
id=content["id"],
type="function",
function=ChatCompletionToolCallFunctionChunk(
name=content["name"],
arguments=json.dumps(content.get("input", {})),
),
index=idx,
)
# Include caller information if present (for programmatic tool calling)
if "caller" in content:
tool_call["caller"] = cast(Dict[str, Any], content["caller"]) # type: ignore[typeddict-item]
tool_calls.append(tool_call)
## TOOL SEARCH TOOL RESULT (skip - this is metadata about tool discovery)
elif content["type"] == "tool_search_tool_result":
# This block contains tool_references that were discovered
# We don't need to include this in the response as it's internal metadata
pass
elif content.get("thinking", None) is not None:
if thinking_blocks is None:
thinking_blocks = []
@@ -916,7 +1106,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
return text_content, citations, thinking_blocks, reasoning_content, tool_calls
def calculate_usage(
self, usage_object: dict, reasoning_content: Optional[str]
self, usage_object: dict, reasoning_content: Optional[str], completion_response: Optional[dict] = None
) -> Usage:
# NOTE: Sometimes the usage object has None set explicitly for token counts, meaning .get() & key access returns None, and we need to account for this
prompt_tokens = usage_object.get("input_tokens", 0) or 0
@@ -926,6 +1116,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
cache_read_input_tokens: int = 0
cache_creation_token_details: Optional[CacheCreationTokenDetails] = None
web_search_requests: Optional[int] = None
tool_search_requests: Optional[int] = None
if (
"cache_creation_input_tokens" in _usage
and _usage["cache_creation_input_tokens"] is not None
@@ -946,6 +1137,25 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
web_search_requests = cast(
int, _usage["server_tool_use"]["web_search_requests"]
)
if (
"tool_search_requests" in _usage["server_tool_use"]
and _usage["server_tool_use"]["tool_search_requests"] is not None
):
tool_search_requests = cast(
int, _usage["server_tool_use"]["tool_search_requests"]
)
# Count tool_search_requests from content blocks if not in usage
# Anthropic doesn't always include tool_search_requests in the usage object
if tool_search_requests is None and completion_response is not None:
tool_search_count = 0
for content in completion_response.get("content", []):
if content.get("type") == "server_tool_use":
tool_name = content.get("name", "")
if "tool_search" in tool_name:
tool_search_count += 1
if tool_search_count > 0:
tool_search_requests = tool_search_count
if "cache_creation" in _usage and _usage["cache_creation"] is not None:
cache_creation_token_details = CacheCreationTokenDetails(
@@ -982,8 +1192,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
cache_read_input_tokens=cache_read_input_tokens,
completion_tokens_details=completion_token_details,
server_tool_use=(
ServerToolUse(web_search_requests=web_search_requests)
if web_search_requests is not None
ServerToolUse(
web_search_requests=web_search_requests,
tool_search_requests=tool_search_requests,
)
if (web_search_requests is not None or tool_search_requests is not None)
else None
),
)
@@ -1077,6 +1290,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
usage = self.calculate_usage(
usage_object=completion_response["usage"],
reasoning_content=reasoning_content,
completion_response=completion_response,
)
setattr(model_response, "usage", usage) # type: ignore
litellm/llms/anthropic/common_utils.py
+101
View File
@@ -88,6 +88,86 @@ class AnthropicModelInfo(BaseLLMModelInfo):
return True
return False
def is_tool_search_used(self, tools: Optional[List]) -> bool:
"""
Check if tool search tools are present in the tools list.
"""
if not tools:
return False
for tool in tools:
tool_type = tool.get("type", "")
if tool_type in ["tool_search_tool_regex_20251119", "tool_search_tool_bm25_20251119"]:
return True
return False
def is_programmatic_tool_calling_used(self, tools: Optional[List]) -> bool:
"""
Check if programmatic tool calling is being used (tools with allowed_callers field).
Returns True if any tool has allowed_callers containing 'code_execution_20250825'.
"""
if not tools:
return False
for tool in tools:
# Check top-level allowed_callers
allowed_callers = tool.get("allowed_callers", None)
if allowed_callers and isinstance(allowed_callers, list):
if "code_execution_20250825" in allowed_callers:
return True
# Check function.allowed_callers for OpenAI format tools
function = tool.get("function", {})
if isinstance(function, dict):
function_allowed_callers = function.get("allowed_callers", None)
if function_allowed_callers and isinstance(function_allowed_callers, list):
if "code_execution_20250825" in function_allowed_callers:
return True
return False
def is_input_examples_used(self, tools: Optional[List]) -> bool:
"""
Check if input_examples is being used in any tools.
Returns True if any tool has input_examples field.
"""
if not tools:
return False
for tool in tools:
# Check top-level input_examples
input_examples = tool.get("input_examples", None)
if input_examples and isinstance(input_examples, list) and len(input_examples) > 0:
return True
# Check function.input_examples for OpenAI format tools
function = tool.get("function", {})
if isinstance(function, dict):
function_input_examples = function.get("input_examples", None)
if function_input_examples and isinstance(function_input_examples, list) and len(function_input_examples) > 0:
return True
return False
def is_effort_used(self, optional_params: Optional[dict]) -> bool:
"""
Check if effort parameter is being used via output_config.
Returns True if output_config with effort field is present.
"""
if not optional_params:
return False
output_config = optional_params.get("output_config")
if output_config and isinstance(output_config, dict):
effort = output_config.get("effort")
if effort and isinstance(effort, str):
return True
return False
def _get_user_anthropic_beta_headers(
self, anthropic_beta_header: Optional[str]
) -> Optional[List[str]]:
@@ -122,6 +202,10 @@ class AnthropicModelInfo(BaseLLMModelInfo):
pdf_used: bool = False,
file_id_used: bool = False,
mcp_server_used: bool = False,
tool_search_used: bool = False,
programmatic_tool_calling_used: bool = False,
input_examples_used: bool = False,
effort_used: bool = False,
is_vertex_request: bool = False,
user_anthropic_beta_headers: Optional[List[str]] = None,
) -> dict:
@@ -138,6 +222,15 @@ class AnthropicModelInfo(BaseLLMModelInfo):
betas.add("code-execution-2025-05-22")
if mcp_server_used:
betas.add("mcp-client-2025-04-04")
# Tool search, programmatic tool calling, and input_examples all use the same beta header
if tool_search_used or programmatic_tool_calling_used or input_examples_used:
from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
betas.add(ANTHROPIC_TOOL_SEARCH_BETA_HEADER)
# Effort parameter uses a separate beta header
if effort_used:
from litellm.types.llms.anthropic import ANTHROPIC_EFFORT_BETA_HEADER
betas.add(ANTHROPIC_EFFORT_BETA_HEADER)
headers = {
"anthropic-version": anthropic_version or "2023-06-01",
@@ -182,6 +275,10 @@ class AnthropicModelInfo(BaseLLMModelInfo):
)
pdf_used = self.is_pdf_used(messages=messages)
file_id_used = self.is_file_id_used(messages=messages)
tool_search_used = self.is_tool_search_used(tools=tools)
programmatic_tool_calling_used = self.is_programmatic_tool_calling_used(tools=tools)
input_examples_used = self.is_input_examples_used(tools=tools)
effort_used = self.is_effort_used(optional_params=optional_params)
user_anthropic_beta_headers = self._get_user_anthropic_beta_headers(
anthropic_beta_header=headers.get("anthropic-beta")
)
@@ -194,6 +291,10 @@ class AnthropicModelInfo(BaseLLMModelInfo):
is_vertex_request=optional_params.get("is_vertex_request", False),
user_anthropic_beta_headers=user_anthropic_beta_headers,
mcp_server_used=mcp_server_used,
tool_search_used=tool_search_used,
programmatic_tool_calling_used=programmatic_tool_calling_used,
input_examples_used=input_examples_used,
effort_used=effort_used,
)
headers = {**headers, **anthropic_headers}
litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py
+1 -1
View File
@@ -645,7 +645,7 @@ class LiteLLMAnthropicMessagesAdapter:
type="tool_use",
id=choice.delta.tool_calls[0].id or str(uuid.uuid4()),
name=choice.delta.tool_calls[0].function.name or "",
input={},
input={}, # type: ignore[typeddict-item]
)
elif isinstance(choice, StreamingChoices) and hasattr(
choice.delta, "thinking_blocks"
litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py
+186
View File
@@ -0,0 +1,186 @@
"""
Transformation for Bedrock imported models that use OpenAI Chat Completions format.
Use this for models imported into Bedrock that accept the OpenAI API format.
Model format: bedrock/openai/<model-id>
Example: bedrock/openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/abc123
"""
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
import httpx
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
from litellm.types.llms.openai import AllMessageValues
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
LiteLLMLoggingObj = _LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM):
"""
Configuration for Bedrock imported models that use OpenAI Chat Completions format.
This class handles the transformation of requests and responses for Bedrock
imported models that accept the OpenAI API format directly.
Inherits from OpenAIGPTConfig to leverage standard OpenAI parameter handling
and response transformation, while adding Bedrock-specific URL generation
and AWS request signing.
Usage:
model = "bedrock/openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/abc123"
"""
def __init__(self, **kwargs):
OpenAIGPTConfig.__init__(self, **kwargs)
BaseAWSLLM.__init__(self, **kwargs)
@property
def custom_llm_provider(self) -> Optional[str]:
return "bedrock"
def _get_openai_model_id(self, model: str) -> str:
"""
Extract the actual model ID from the LiteLLM model name.
Input format: bedrock/openai/<model-id>
Returns: <model-id>
"""
# Remove bedrock/ prefix if present
if model.startswith("bedrock/"):
model = model[8:]
# Remove openai/ prefix
if model.startswith("openai/"):
model = model[7:]
return model
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
"""
Get the complete URL for the Bedrock invoke endpoint.
Uses the standard Bedrock invoke endpoint format.
"""
model_id = self._get_openai_model_id(model)
# Get AWS region
aws_region_name = self._get_aws_region_name(
optional_params=optional_params, model=model
)
# Get runtime endpoint
aws_bedrock_runtime_endpoint = optional_params.get(
"aws_bedrock_runtime_endpoint", None
)
endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint(
api_base=api_base,
aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
aws_region_name=aws_region_name,
)
# Build the invoke URL
if stream:
endpoint_url = f"{endpoint_url}/model/{model_id}/invoke-with-response-stream"
else:
endpoint_url = f"{endpoint_url}/model/{model_id}/invoke"
return endpoint_url
def sign_request(
self,
headers: dict,
optional_params: dict,
request_data: dict,
api_base: str,
api_key: Optional[str] = None,
model: Optional[str] = None,
stream: Optional[bool] = None,
fake_stream: Optional[bool] = None,
) -> Tuple[dict, Optional[bytes]]:
"""
Sign the request using AWS Signature Version 4.
"""
return self._sign_request(
service_name="bedrock",
headers=headers,
optional_params=optional_params,
request_data=request_data,
api_base=api_base,
api_key=api_key,
model=model,
stream=stream,
fake_stream=fake_stream,
)
def transform_request(
self,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
headers: dict,
) -> dict:
"""
Transform the request to OpenAI Chat Completions format for Bedrock imported models.
Removes AWS-specific params and stream param (handled separately in URL),
then delegates to parent class for standard OpenAI request transformation.
"""
# Remove stream from optional_params as it's handled separately in URL
optional_params.pop("stream", None)
# Remove AWS-specific params that shouldn't be in the request body
inference_params = {
k: v
for k, v in optional_params.items()
if k not in self.aws_authentication_params
}
# Use parent class transform_request for OpenAI format
return super().transform_request(
model=self._get_openai_model_id(model),
messages=messages,
optional_params=inference_params,
litellm_params=litellm_params,
headers=headers,
)
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
"""
Validate the environment and return headers.
For Bedrock, we don't need Bearer token auth since we use AWS SigV4.
"""
return headers
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BedrockError:
"""Return the appropriate error class for Bedrock."""
return BedrockError(status_code=status_code, message=error_message)
litellm/llms/bedrock/common_utils.py
+16 -2
View File
@@ -403,6 +403,9 @@ class BedrockModelInfo(BaseLLMModelInfo):
if model.startswith("invoke/"):
model = model.split("/", 1)[1]
if model.startswith("openai/"):
model = model.split("/", 1)[1]
return model
@staticmethod
@@ -446,12 +449,12 @@ class BedrockModelInfo(BaseLLMModelInfo):
@staticmethod
def get_bedrock_route(
model: str,
) -> Literal["converse", "invoke", "converse_like", "agent", "agentcore", "async_invoke"]:
) -> Literal["converse", "invoke", "converse_like", "agent", "agentcore", "async_invoke", "openai"]:
"""
Get the bedrock route for the given model.
"""
route_mappings: Dict[
str, Literal["invoke", "converse_like", "converse", "agent", "agentcore", "async_invoke"]
str, Literal["invoke", "converse_like", "converse", "agent", "agentcore", "async_invoke", "openai"]
] = {
"invoke/": "invoke",
"converse_like/": "converse_like",
@@ -459,6 +462,7 @@ class BedrockModelInfo(BaseLLMModelInfo):
"agent/": "agent",
"agentcore/": "agentcore",
"async_invoke/": "async_invoke",
"openai/": "openai",
}
# Check explicit routes first
@@ -517,6 +521,14 @@ class BedrockModelInfo(BaseLLMModelInfo):
"""
return "async_invoke/" in model
@staticmethod
def _explicit_openai_route(model: str) -> bool:
"""
Check if the model is an explicit openai route.
Used for Bedrock imported models that use OpenAI Chat Completions format.
"""
return "openai/" in model
@staticmethod
def get_bedrock_provider_config_for_messages_api(
model: str,
@@ -566,6 +578,8 @@ def get_bedrock_chat_config(model: str):
# Handle explicit routes first
if bedrock_route == "converse" or bedrock_route == "converse_like":
return litellm.AmazonConverseConfig()
elif bedrock_route == "openai":
return litellm.AmazonBedrockOpenAIConfig()
elif bedrock_route == "agent":
from litellm.llms.bedrock.chat.invoke_agent.transformation import (
AmazonInvokeAgentConfig,
litellm/llms/bedrock/image/image_handler.py
+2
View File
@@ -1,3 +1,5 @@
from __future__ import annotations
import json
from typing import TYPE_CHECKING, Any, Optional, Union
litellm/llms/gemini/image_generation/transformation.py
+2 -1
View File
@@ -25,6 +25,7 @@ FLASH_IMAGE_PREVIEW_MODEL_IDENTIFIERS = (
"2.0-flash-preview-image",
"2.0-flash-preview-image-generation",
"2.5-flash-image-preview",
"3-pro-image-preview",
)
class GoogleImageGenConfig(BaseImageGenerationConfig):
DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta"
@@ -75,7 +76,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
"1792x1024": "16:9",
"1024x1792": "9:16",
"1280x896": "4:3",
"896x1280": "3:4"
"896x1280": "3:4",
}
return aspect_ratio_map.get(size, "1:1")
litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py
+2
View File
@@ -65,6 +65,8 @@ class VertexAIPartnerModelsTokenCounter(VertexBase):
# Use custom api_base if provided, otherwise construct default
if api_base:
base_url = api_base
elif vertex_location == "global":
base_url = "https://aiplatform.googleapis.com"
else:
base_url = f"https://{vertex_location}-aiplatform.googleapis.com"
litellm/model_prices_and_context_window_backup.json
+52
View File
@@ -24605,6 +24605,58 @@
"supports_tool_choice": true,
"supports_vision": true
},
"vertex_ai/claude-opus-4-5": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "vertex_ai-anthropic_models",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": true,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
},
"vertex_ai/claude-opus-4-5@20251101": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "vertex_ai-anthropic_models",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": true,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
},
"vertex_ai/claude-sonnet-4-5": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+2 -2
View File
@@ -1976,7 +1976,7 @@ class MCPServerManager:
verbose_logger.debug(
f"Adding server to registry: {server.server_id} ({server.server_name})"
)
self.add_update_server(server)
await self.add_update_server(server)
verbose_logger.debug(
f"Registry now contains {len(self.get_registry())} servers"
@@ -2270,7 +2270,7 @@ class MCPServerManager:
server.status = "unhealthy"
## try adding server to registry to get error
try:
self.add_update_server(server)
await self.add_update_server(server)
except Exception as e:
server.health_check_error = str(e)
server.health_check_error = "Server is not in in memory registry yet. This could be a temporary sync issue."
litellm/proxy/common_request_processing.py
+2
View File
@@ -330,6 +330,7 @@ class ProxyBaseLLMRequestProcessing:
"avideo_remix",
"acreate_container",
"alist_containers",
"aingest",
"aretrieve_container",
"adelete_container",
"acreate_skill",
@@ -453,6 +454,7 @@ class ProxyBaseLLMRequestProcessing:
"avideo_remix",
"acreate_container",
"alist_containers",
"aingest",
"aretrieve_container",
"adelete_container",
"acreate_skill",
litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py
+43 -8
View File
@@ -90,6 +90,10 @@ class PillarGuardrail(CustomGuardrail):
fallback_on_error: Action when API errors occur ('allow' or 'block')
timeout: Timeout for API calls in seconds
**kwargs: Additional arguments passed to parent class
Note:
LiteLLM virtual key context (user_id, team_id, key_alias, etc.) is always
automatically passed as X-LiteLLM-* headers to enable application/user tracking.
"""
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
self.api_key = api_key or os.environ.get("PILLAR_API_KEY")
@@ -222,7 +226,7 @@ class PillarGuardrail(CustomGuardrail):
return data
verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook")
result = await self.run_pillar_guardrail(data)
result = await self.run_pillar_guardrail(data, user_api_key_dict)
# Add guardrail name to response headers
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
@@ -265,7 +269,7 @@ class PillarGuardrail(CustomGuardrail):
return data
verbose_proxy_logger.debug("Pillar Guardrail: During-call moderation hook")
result = await self.run_pillar_guardrail(data)
result = await self.run_pillar_guardrail(data, user_api_key_dict)
# Add guardrail name to response headers
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
@@ -315,7 +319,7 @@ class PillarGuardrail(CustomGuardrail):
post_call_data["messages"] = data.get("messages", []) + response_messages
# Reuse the existing guardrail logic - zero duplication!
await self.run_pillar_guardrail(post_call_data)
await self.run_pillar_guardrail(post_call_data, user_api_key_dict)
# Add guardrail name to response headers
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
@@ -326,12 +330,13 @@ class PillarGuardrail(CustomGuardrail):
# CORE LOGIC METHOD
# =========================================================================
async def run_pillar_guardrail(self, data: dict) -> dict:
async def run_pillar_guardrail(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> dict:
"""
Core method to run the Pillar guardrail scan.
Args:
data: Request data containing messages and metadata
user_api_key_dict: User API key authentication info containing key context
Returns:
Original data if safe or in monitor mode
@@ -345,7 +350,7 @@ class PillarGuardrail(CustomGuardrail):
return data
try:
headers = self._prepare_headers()
headers = self._prepare_headers(user_api_key_dict)
payload = self._prepare_payload(data)
response = await self._call_pillar_api(
@@ -403,8 +408,16 @@ class PillarGuardrail(CustomGuardrail):
},
)
def _prepare_headers(self) -> Dict[str, str]:
"""Prepare headers for the Pillar API request."""
def _prepare_headers(self, user_api_key_dict: UserAPIKeyAuth) -> Dict[str, str]:
"""
Prepare headers for the Pillar API request.
Args:
user_api_key_dict: User API key authentication info containing key context
Returns:
Dictionary of headers to send to Pillar API
"""
if not self.api_key:
msg = (
"Couldn't get Pillar API key, either set the `PILLAR_API_KEY` in the environment or "
@@ -415,7 +428,7 @@ class PillarGuardrail(CustomGuardrail):
headers: Dict[str, str] = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
}
# Add Pillar-specific headers based on configuration
self._set_bool_header(headers, "plr_scanners", self.include_scanners)
@@ -423,6 +436,20 @@ class PillarGuardrail(CustomGuardrail):
self._set_bool_header(headers, "plr_async", self.async_mode)
self._set_bool_header(headers, "plr_persist", self.persist_session)
# Always add LiteLLM virtual key context headers (metadata excluded for security)
context_mapping = {
"X-LiteLLM-Key-Name": user_api_key_dict.key_name,
"X-LiteLLM-Key-Alias": user_api_key_dict.key_alias,
"X-LiteLLM-User-Id": user_api_key_dict.user_id,
"X-LiteLLM-User-Email": user_api_key_dict.user_email,
"X-LiteLLM-Team-Id": user_api_key_dict.team_id,
"X-LiteLLM-Team-Name": user_api_key_dict.team_alias,
"X-LiteLLM-Org-Id": user_api_key_dict.org_id,
}
for header_name, value in context_mapping.items():
if value:
headers[header_name] = str(value)
return headers
def _set_bool_header(self, headers: Dict[str, str], header_name: str, value: Optional[bool]) -> None:
@@ -517,6 +544,14 @@ class PillarGuardrail(CustomGuardrail):
"""
Prepare the payload for the Pillar API request following the /api/v1/protect contract.
This method supports multi-modal content (images, files, audio, video, etc.) as messages
are passed through without modification. The messages array can contain any OpenAI-compatible
message structure including:
- Text content (string)
- Multi-modal content blocks (image_url, image_file, audio, video, document, file)
- Attachments
- Tool calls
Args:
data: Request data
litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py
+13 -6
View File
@@ -1,13 +1,18 @@
import os
import re
import asyncio
import base64
import os
import re
from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Type, Union
from fastapi import HTTPException
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client, httpxSpecialProvider
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import (
Choices,
@@ -15,7 +20,7 @@ from litellm.types.utils import (
EmbeddingResponse,
ImageResponse,
ModelResponse,
ModelResponseStream
ModelResponseStream,
)
if TYPE_CHECKING:
@@ -267,8 +272,10 @@ class PromptSecurityGuardrail(CustomGuardrail):
content = msg.get('content', '')
# Handle both string and list content types
if isinstance(content, str):
if content.startswith('### '): return False
if '"follow_ups": [' in content: return False
if content.startswith('### '):
return False
if '"follow_ups": [' in content:
return False
return True
messages = list(filter(lambda msg: good_msg(msg), messages))
litellm/proxy/guardrails/guardrail_hooks/tool_permission.py
+10 -15
View File
@@ -23,6 +23,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
ToolResult,
)
from litellm.types.utils import (
CallTypesLiteral,
ChatCompletionMessageToolCall,
Choices,
LLMResponseTypes,
@@ -202,16 +203,21 @@ class ToolPermissionGuardrail(CustomGuardrail):
return {}
def _collect_argument_paths(
self, value: Any, current_path: str, collected: Dict[str, List[Any]]
self, value: Any, current_path: str, collected: Dict[str, List[Any]], depth: int = 0
) -> None:
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
if depth > DEFAULT_MAX_RECURSE_DEPTH:
return
if isinstance(value, dict):
for key, sub_value in value.items():
next_path = f"{current_path}.{key}" if current_path else key
self._collect_argument_paths(sub_value, next_path, collected)
self._collect_argument_paths(sub_value, next_path, collected, depth + 1)
elif isinstance(value, list):
list_path = f"{current_path}[]" if current_path else "[]"
for item in value:
self._collect_argument_paths(item, list_path, collected)
self._collect_argument_paths(item, list_path, collected, depth + 1)
else:
if not current_path:
return
@@ -437,18 +443,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
call_type: CallTypesLiteral,
) -> Union[Exception, str, dict, None]:
""" """
verbose_proxy_logger.debug("Tool Permission Guardrail Pre-Call Hook")
litellm/proxy/proxy_config.yaml
+2 -17
View File
@@ -1,22 +1,7 @@
model_list:
- model_name: aws/anthropic/bedrock-claude-3-5-sonnet-v1
- model_name: qwen-25vl-72b
litellm_params:
model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0
aws_region_name: us-east-1
custom_llm_provider: bedrock
- model_name: aws/anthropic/bedrock-claude-3-5-sonnet-v1
litellm_params:
model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0
aws_region_name: us-west-2
custom_llm_provider: bedrock
- model_name: bedrock/*
litellm_params:
model: bedrock/*
custom_llm_provider: bedrock
aws_region_name: us-west-2
- model_name: runwayml/*
litellm_params:
model: runwayml/*
model: bedrock/openai/arn:aws:bedrock:us-east-1:046319184608:imported-model/0m2lasirsp6z
litellm/proxy/proxy_server.py
+4 -1
View File
@@ -362,6 +362,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
)
from litellm.proxy.prompts.prompt_endpoints import router as prompts_router
from litellm.proxy.public_endpoints import router as public_endpoints_router
from litellm.proxy.rag_endpoints.endpoints import router as rag_router
from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router
from litellm.proxy.response_api_endpoints.endpoints import router as response_router
from litellm.proxy.route_llm_request import route_request
@@ -5483,6 +5484,7 @@ async def audio_transcriptions(
file_object = io.BytesIO(file_content)
file_object.name = file.filename
data["file"] = file_object
try:
### CALL HOOKS ### - modify incoming data / reject request before calling the model
data = await proxy_logging_obj.pre_call_hook(
@@ -5500,7 +5502,7 @@ async def audio_transcriptions(
)
response = await llm_call
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
raise e
finally:
file_object.close() # close the file read in by io library
@@ -10156,6 +10158,7 @@ app.include_router(batches_router)
app.include_router(public_endpoints_router)
app.include_router(rerank_router)
app.include_router(ocr_router)
app.include_router(rag_router)
app.include_router(video_router)
app.include_router(container_router)
app.include_router(search_router)
litellm/proxy/rag_endpoints/__init__.py
+6
View File
@@ -0,0 +1,6 @@
"""RAG Endpoints for LiteLLM Proxy."""
from litellm.proxy.rag_endpoints.endpoints import router
__all__ = ["router"]
litellm/proxy/rag_endpoints/endpoints.py
+200
View File
@@ -0,0 +1,200 @@
"""
RAG Ingest Endpoints for LiteLLM Proxy.
Provides an all-in-one API for document ingestion:
Upload -> (OCR) -> Chunk -> Embed -> Vector Store
"""
import base64
from typing import Any, Dict, Optional, Tuple
import orjson
from fastapi import APIRouter, Depends, HTTPException, Request, Response
from fastapi.responses import ORJSONResponse
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
_safe_get_request_headers,
get_form_data,
)
router = APIRouter()
async def parse_rag_ingest_request(
request: Request,
) -> Tuple[Dict[str, Any], Optional[Tuple[str, bytes, str]], Optional[str], Optional[str]]:
"""
Parse RAG ingest request.
Supports:
- Form: file + request JSON in form field
- JSON body for URL-based ingestion
Returns:
Tuple of (ingest_options, file_data, file_url, file_id)
"""
headers = _safe_get_request_headers(request)
content_type = headers.get("content-type", "")
file_data = None
file_url = None
file_id = None
ingest_options: Dict[str, Any] = {}
if "multipart/form-data" in content_type:
# Form upload
form_data = await get_form_data(request)
# Get file
file_obj = form_data.get("file")
if file_obj is not None and hasattr(file_obj, "read"):
file_content = await file_obj.read()
file_data = (file_obj.filename, file_content, file_obj.content_type)
# Parse JSON from 'request' form field (contains full request body as JSON)
request_json_str = form_data.get("request")
if request_json_str:
request_data = orjson.loads(request_json_str)
ingest_options = request_data.get("ingest_options", {})
file_url = request_data.get("file_url")
file_id = request_data.get("file_id")
else:
# JSON body
data = await _read_request_body(request)
ingest_options = data.get("ingest_options", {})
file_url = data.get("file_url")
file_id = data.get("file_id")
# Handle base64-encoded file in JSON body
file_obj = data.get("file")
if file_obj and isinstance(file_obj, dict):
filename = file_obj.get("filename")
content_b64 = file_obj.get("content")
content_type = file_obj.get("content_type", "application/octet-stream")
if filename and content_b64:
try:
file_content = base64.b64decode(content_b64)
file_data = (filename, file_content, content_type)
except Exception as e:
raise HTTPException(
status_code=400,
detail={"error": f"Invalid base64 content: {e}"},
)
# Validate
if file_data is None and file_url is None and file_id is None:
raise HTTPException(
status_code=400,
detail={"error": "Must provide file, file_url, or file_id"},
)
if "vector_store" not in ingest_options:
raise HTTPException(
status_code=400,
detail={"error": "ingest_options must contain 'vector_store' configuration"},
)
return ingest_options, file_data, file_url, file_id
@router.post(
"/v1/rag/ingest",
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse,
tags=["rag"],
)
@router.post(
"/rag/ingest",
dependencies=[Depends(user_api_key_auth)],
response_class=ORJSONResponse,
tags=["rag"],
)
async def rag_ingest(
request: Request,
fastapi_response: Response,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
RAG Ingest endpoint - all-in-one document ingestion pipeline.
Supports form upload (for files) or JSON body (for URLs).
## Form upload (for files):
```bash
curl -X POST "http://localhost:4000/v1/rag/ingest" \\
-H "Authorization: Bearer sk-1234" \\
-F file="@document.pdf" \\
-F 'ingest_options={"vector_store": {"custom_llm_provider": "openai"}}'
```
## JSON body (for URLs):
```bash
curl -X POST "http://localhost:4000/v1/rag/ingest" \\
-H "Authorization: Bearer sk-1234" \\
-H "Content-Type: application/json" \\
-d '{
"file_url": "https://example.com/document.pdf",
"ingest_options": {"vector_store": {"custom_llm_provider": "openai"}}
}'
```
## Bedrock:
```bash
curl -X POST "http://localhost:4000/v1/rag/ingest" \\
-H "Authorization: Bearer sk-1234" \\
-F file="@document.pdf" \\
-F 'ingest_options={"vector_store": {"custom_llm_provider": "bedrock"}}'
```
"""
from litellm.proxy.proxy_server import (
add_litellm_data_to_request,
general_settings,
llm_router,
proxy_config,
version,
)
try:
# Parse request
ingest_options, file_data, file_url, file_id = await parse_rag_ingest_request(request)
# Add litellm data
request_data: Dict[str, Any] = {}
request_data = await add_litellm_data_to_request(
data=request_data,
request=request,
general_settings=general_settings,
user_api_key_dict=user_api_key_dict,
version=version,
proxy_config=proxy_config,
)
verbose_proxy_logger.debug(f"RAG Ingest - options: {ingest_options}")
# Call ingest
response = await litellm.aingest(
ingest_options=ingest_options,
file_data=file_data,
file_url=file_url,
file_id=file_id,
router=llm_router,
**request_data,
)
return response
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"RAG Ingest failed: {e}")
raise HTTPException(
status_code=500,
detail={"error": str(e)},
)
litellm/proxy/route_llm_request.py
+3
View File
@@ -40,6 +40,7 @@ ROUTE_ENDPOINT_MAPPING = {
"alist_skills": "/skills",
"aget_skill": "/skills/{skill_id}",
"adelete_skill": "/skills/{skill_id}",
"aingest": "/rag/ingest",
}
@@ -134,6 +135,7 @@ async def route_request(
"alist_skills",
"aget_skill",
"adelete_skill",
"aingest",
],
):
"""
@@ -190,6 +192,7 @@ async def route_request(
"alist_skills",
"aget_skill",
"adelete_skill",
"aingest",
] and (data.get("model") is None or data.get("model") == ""):
# These endpoints don't need a model, use custom_llm_provider directly
return getattr(litellm, f"{route_type}")(**data)
litellm/proxy/spend_tracking/spend_management_endpoints.py
+2 -2
View File
@@ -1427,7 +1427,7 @@ async def _get_spend_report_for_time_range(
LEFT JOIN
"LiteLLM_TeamTable" t ON s.team_id = t.team_id
WHERE
s."startTime"::DATE >= $1::date AND s."startTime"::DATE <= $2::date
s."startTime" >= $1::date AND s."startTime" < ($2::date + INTERVAL '1 day')
GROUP BY
t.team_alias
ORDER BY
@@ -1441,7 +1441,7 @@ async def _get_spend_report_for_time_range(
jsonb_array_elements_text(request_tags) AS individual_request_tag,
SUM(spend) AS total_spend
FROM "LiteLLM_SpendLogs"
WHERE "startTime"::DATE >= $1::date AND "startTime"::DATE <= $2::date
WHERE "startTime" >= $1::date AND "startTime" < ($2::date + INTERVAL '1 day')
GROUP BY individual_request_tag
ORDER BY total_spend DESC;
"""
litellm/rag/__init__.py
+22
View File
@@ -0,0 +1,22 @@
"""
LiteLLM RAG (Retrieval Augmented Generation) Module.
Provides an all-in-one API for document ingestion:
Upload -> (OCR) -> Chunk -> Embed -> Vector Store
"""
from litellm.rag.main import aingest, ingest
__all__ = ["ingest", "aingest"]
# Expose at litellm.rag level for convenience
async def arag_ingest(*args, **kwargs):
"""Alias for aingest."""
return await aingest(*args, **kwargs)
def rag_ingest(*args, **kwargs):
"""Alias for ingest."""
return ingest(*args, **kwargs)
litellm/rag/ingestion/__init__.py
+14
View File
@@ -0,0 +1,14 @@
"""
RAG Ingestion classes for different providers.
"""
from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion
from litellm.rag.ingestion.bedrock_ingestion import BedrockRAGIngestion
from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion
__all__ = [
"BaseRAGIngestion",
"BedrockRAGIngestion",
"OpenAIRAGIngestion",
]
litellm/rag/ingestion/base_ingestion.py
+314
View File
@@ -0,0 +1,314 @@
"""
Base RAG Ingestion class.
Provides abstract methods for:
- OCR
- Chunking
- Embedding
- Vector Store operations
Providers can inherit and override methods as needed.
"""
from __future__ import annotations
import base64
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm._uuid import uuid4
from litellm.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE
from litellm.rag.text_splitters import RecursiveCharacterTextSplitter
from litellm.types.rag import RAGIngestOptions, RAGIngestResponse
if TYPE_CHECKING:
from litellm import Router
class BaseRAGIngestion(ABC):
"""
Base class for RAG ingestion.
Providers should inherit from this class and override methods as needed.
For example, OpenAI handles embedding internally when attaching files to
vector stores, so it overrides the embedding step to be a no-op.
"""
def __init__(
self,
ingest_options: RAGIngestOptions,
router: Optional["Router"] = None,
):
self.ingest_options = ingest_options
self.router = router
self.ingest_id = f"ingest_{uuid4()}"
# Extract configs from options
self.ocr_config = ingest_options.get("ocr")
self.chunking_strategy = ingest_options.get("chunking_strategy", {"type": "auto"})
self.embedding_config = ingest_options.get("embedding")
self.vector_store_config = ingest_options.get("vector_store") or {}
self.ingest_name = ingest_options.get("name")
@property
def custom_llm_provider(self) -> str:
"""Get the vector store provider."""
return self.vector_store_config.get("custom_llm_provider", "openai")
async def upload(
self,
file_data: Optional[Tuple[str, bytes, str]] = None,
file_url: Optional[str] = None,
file_id: Optional[str] = None,
) -> Tuple[Optional[str], Optional[bytes], Optional[str], Optional[str]]:
"""
Upload / prepare file for ingestion.
Args:
file_data: Tuple of (filename, content_bytes, content_type)
file_url: URL to fetch file from
file_id: Existing file ID to use
Returns:
Tuple of (filename, file_content, content_type, existing_file_id)
"""
if file_data:
filename, file_content, content_type = file_data
return filename, file_content, content_type, None
if file_url:
async with httpx.AsyncClient() as http_client:
response = await http_client.get(file_url)
response.raise_for_status()
file_content = response.content
filename = file_url.split("/")[-1] or "document"
content_type = response.headers.get("content-type", "application/octet-stream")
return filename, file_content, content_type, None
if file_id:
return None, None, None, file_id
raise ValueError("Must provide file_data, file_url, or file_id")
async def ocr(
self,
file_content: Optional[bytes],
content_type: Optional[str],
) -> Optional[str]:
"""
Perform OCR on file content to extract text.
Args:
file_content: Raw file bytes
content_type: MIME type of the file
Returns:
Extracted text or None if OCR not configured/needed
"""
if not self.ocr_config or not file_content:
return None
ocr_model = self.ocr_config.get("model", "mistral/mistral-ocr-latest")
# Determine document type
if content_type and "image" in content_type:
doc_type, url_key = "image_url", "image_url"
else:
doc_type, url_key = "document_url", "document_url"
# Encode as base64 data URL
b64_content = base64.b64encode(file_content).decode("utf-8")
data_url = f"data:{content_type};base64,{b64_content}"
# Use router if available
if self.router is not None:
ocr_response = await self.router.aocr(
model=ocr_model,
document={"type": doc_type, url_key: data_url},
)
else:
ocr_response = await litellm.aocr(
model=ocr_model,
document={"type": doc_type, url_key: data_url},
)
# Extract text from pages
if hasattr(ocr_response, "pages") and ocr_response.pages: # type: ignore
return "\n\n".join(
page.markdown for page in ocr_response.pages if hasattr(page, "markdown") # type: ignore
)
return None
def chunk(
self,
text: Optional[str],
file_content: Optional[bytes],
ocr_was_used: bool,
) -> List[str]:
"""
Split text into chunks using RecursiveCharacterTextSplitter.
Args:
text: Text from OCR (if used)
file_content: Raw file content bytes
ocr_was_used: Whether OCR was performed
Returns:
List of text chunks
"""
# Get text to chunk
text_to_chunk: Optional[str] = None
if text:
text_to_chunk = text
elif file_content and not ocr_was_used:
try:
text_to_chunk = file_content.decode("utf-8")
except UnicodeDecodeError:
verbose_logger.debug("Binary file detected, skipping text chunking")
return []
if not text_to_chunk:
return []
# Extract RecursiveCharacterTextSplitter args
splitter_args = self.chunking_strategy or {}
chunk_size = splitter_args.get("chunk_size", DEFAULT_CHUNK_SIZE)
chunk_overlap = splitter_args.get("chunk_overlap", DEFAULT_CHUNK_OVERLAP)
separators = splitter_args.get("separators", None)
# Build splitter kwargs
splitter_kwargs: Dict[str, Any] = {
"chunk_size": chunk_size,
"chunk_overlap": chunk_overlap,
}
if separators:
splitter_kwargs["separators"] = separators
text_splitter = RecursiveCharacterTextSplitter(**splitter_kwargs)
return text_splitter.split_text(text_to_chunk)
async def embed(
self,
chunks: List[str],
) -> Optional[List[List[float]]]:
"""
Generate embeddings for text chunks.
Args:
chunks: List of text chunks
Returns:
List of embeddings or None
"""
if not self.embedding_config or not chunks:
return None
embedding_model = self.embedding_config.get("model", "text-embedding-3-small")
if self.router is not None:
response = await self.router.aembedding(model=embedding_model, input=chunks)
else:
response = await litellm.aembedding(model=embedding_model, input=chunks)
return [item["embedding"] for item in response.data]
@abstractmethod
async def store(
self,
file_content: Optional[bytes],
filename: Optional[str],
content_type: Optional[str],
chunks: List[str],
embeddings: Optional[List[List[float]]],
) -> Tuple[Optional[str], Optional[str]]:
"""
Store content in vector store.
This method must be implemented by provider-specific subclasses.
Args:
file_content: Raw file bytes
filename: Name of the file
content_type: MIME type
chunks: Text chunks (if chunking was done locally)
embeddings: Embeddings (if embedding was done locally)
Returns:
Tuple of (vector_store_id, file_id)
"""
pass
async def ingest(
self,
file_data: Optional[Tuple[str, bytes, str]] = None,
file_url: Optional[str] = None,
file_id: Optional[str] = None,
) -> RAGIngestResponse:
"""
Execute the full ingestion pipeline.
Args:
file_data: Tuple of (filename, content_bytes, content_type)
file_url: URL to fetch file from
file_id: Existing file ID to use
Returns:
RAGIngestResponse with status and IDs
Raises:
ValueError: If no input source is provided
"""
# Step 1: Upload (raises ValueError if no input provided)
filename, file_content, content_type, existing_file_id = await self.upload(
file_data=file_data,
file_url=file_url,
file_id=file_id,
)
try:
# Step 2: OCR (optional)
extracted_text = await self.ocr(
file_content=file_content,
content_type=content_type,
)
# Step 3: Chunking
chunks = self.chunk(
text=extracted_text,
file_content=file_content,
ocr_was_used=self.ocr_config is not None,
)
# Step 4: Embedding (optional - some providers handle this internally)
embeddings = await self.embed(chunks=chunks)
# Step 5: Store in vector store
vector_store_id, result_file_id = await self.store(
file_content=file_content,
filename=filename,
content_type=content_type,
chunks=chunks,
embeddings=embeddings,
)
return RAGIngestResponse(
id=self.ingest_id,
status="completed",
vector_store_id=vector_store_id or "",
file_id=result_file_id or existing_file_id,
)
except Exception as e:
verbose_logger.exception(f"RAG Pipeline failed: {e}")
return RAGIngestResponse(
id=self.ingest_id,
status="failed",
vector_store_id="",
file_id=None,
)
litellm/rag/ingestion/bedrock_ingestion.py
+619
View File
@@ -0,0 +1,619 @@
"""
Bedrock-specific RAG Ingestion implementation.
Bedrock Knowledge Bases handle embedding internally when files are ingested,
so this implementation uploads files to S3 and triggers ingestion jobs.
Supports two modes:
1. Use existing KB: Provide vector_store_id (KB ID)
2. Auto-create KB: Don't provide vector_store_id - creates all AWS resources automatically
"""
from __future__ import annotations
import json
import time
import uuid
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
from litellm._logging import verbose_logger
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion
if TYPE_CHECKING:
from litellm import Router
from litellm.types.rag import RAGIngestOptions
def _get_str_or_none(value: Any) -> Optional[str]:
"""Cast config value to Optional[str]."""
return str(value) if value is not None else None
def _get_int(value: Any, default: int) -> int:
"""Cast config value to int with default."""
if value is None:
return default
return int(value)
class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM):
"""
Bedrock Knowledge Base RAG ingestion.
Supports two modes:
1. **Use existing KB**: Provide vector_store_id
2. **Auto-create KB**: Don't provide vector_store_id - creates S3 bucket,
OpenSearch Serverless collection, IAM role, KB, and data source automatically
Optional config:
- vector_store_id: Existing KB ID (if not provided, auto-creates)
- s3_bucket: S3 bucket (auto-created if not provided)
- embedding_model: Bedrock embedding model (default: amazon.titan-embed-text-v2:0)
- wait_for_ingestion: Wait for completion (default: True)
- ingestion_timeout: Max seconds to wait (default: 300)
AWS Auth (uses BaseAWSLLM):
- aws_access_key_id, aws_secret_access_key, aws_session_token
- aws_region_name (default: us-west-2)
- aws_role_name, aws_session_name, aws_profile_name
- aws_web_identity_token, aws_sts_endpoint, aws_external_id
"""
def __init__(
self,
ingest_options: "RAGIngestOptions",
router: Optional["Router"] = None,
):
BaseRAGIngestion.__init__(self, ingest_options=ingest_options, router=router)
BaseAWSLLM.__init__(self)
# Use vector_store_id as unified param (maps to knowledge_base_id)
self.knowledge_base_id = self.vector_store_config.get(
"vector_store_id"
) or self.vector_store_config.get("knowledge_base_id")
# Optional config
self._data_source_id = self.vector_store_config.get("data_source_id")
self._s3_bucket = self.vector_store_config.get("s3_bucket")
self._s3_prefix: Optional[str] = str(self.vector_store_config.get("s3_prefix")) if self.vector_store_config.get("s3_prefix") else None
self.embedding_model = self.vector_store_config.get(
"embedding_model"
) or "amazon.titan-embed-text-v2:0"
self.wait_for_ingestion = self.vector_store_config.get("wait_for_ingestion", False)
self.ingestion_timeout: int = _get_int(self.vector_store_config.get("ingestion_timeout"), 300)
# Get AWS region using BaseAWSLLM method
_aws_region = self.vector_store_config.get("aws_region_name")
self.aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
aws_region_name=str(_aws_region) if _aws_region else None
)
# Will be set during initialization
self.data_source_id: Optional[str] = None
self.s3_bucket: Optional[str] = None
self.s3_prefix: str = self._s3_prefix or "data/"
self._config_initialized = False
# Track resources we create (for cleanup if needed)
self._created_resources: Dict[str, Any] = {}
def _ensure_config_initialized(self):
"""Lazily initialize KB config - either detect from existing or create new."""
if self._config_initialized:
return
if self.knowledge_base_id:
# Use existing KB - auto-detect data source and S3 bucket
self._auto_detect_config()
else:
# No KB provided - create everything from scratch
self._create_knowledge_base_infrastructure()
self._config_initialized = True
def _auto_detect_config(self):
"""Auto-detect data source ID and S3 bucket from existing Knowledge Base."""
verbose_logger.debug(
f"Auto-detecting data source and S3 bucket for KB={self.knowledge_base_id}"
)
bedrock_agent = self._get_boto3_client("bedrock-agent")
# List data sources for this KB
ds_response = bedrock_agent.list_data_sources(
knowledgeBaseId=self.knowledge_base_id
)
data_sources = ds_response.get("dataSourceSummaries", [])
if not data_sources:
raise ValueError(
f"No data sources found for Knowledge Base {self.knowledge_base_id}. "
"Please create a data source first or provide data_source_id and s3_bucket."
)
# Use first data source (or user-provided override)
if self._data_source_id:
self.data_source_id = self._data_source_id
else:
self.data_source_id = data_sources[0]["dataSourceId"]
verbose_logger.info(f"Auto-detected data source: {self.data_source_id}")
# Get data source details for S3 bucket
ds_details = bedrock_agent.get_data_source(
knowledgeBaseId=self.knowledge_base_id,
dataSourceId=self.data_source_id,
)
s3_config = (
ds_details.get("dataSource", {})
.get("dataSourceConfiguration", {})
.get("s3Configuration", {})
)
bucket_arn = s3_config.get("bucketArn", "")
if bucket_arn:
# Extract bucket name from ARN: arn:aws:s3:::bucket-name
self.s3_bucket = self._s3_bucket or bucket_arn.split(":")[-1]
verbose_logger.info(f"Auto-detected S3 bucket: {self.s3_bucket}")
# Use inclusion prefix if available
prefixes = s3_config.get("inclusionPrefixes", [])
if prefixes and not self._s3_prefix:
self.s3_prefix = prefixes[0]
else:
if not self._s3_bucket:
raise ValueError(
f"Could not auto-detect S3 bucket for data source {self.data_source_id}. "
"Please provide s3_bucket in config."
)
self.s3_bucket = self._s3_bucket
def _create_knowledge_base_infrastructure(self):
"""Create all AWS resources needed for a new Knowledge Base."""
verbose_logger.info("Creating new Bedrock Knowledge Base infrastructure...")
# Generate unique names
unique_id = uuid.uuid4().hex[:8]
kb_name = self.ingest_name or f"litellm-kb-{unique_id}"
# Get AWS account ID
sts = self._get_boto3_client("sts")
account_id = sts.get_caller_identity()["Account"]
# Step 1: Create S3 bucket (if not provided)
self.s3_bucket = self._s3_bucket or self._create_s3_bucket(unique_id)
# Step 2: Create OpenSearch Serverless collection
collection_name, collection_arn = self._create_opensearch_collection(
unique_id, account_id
)
# Step 3: Create OpenSearch index
self._create_opensearch_index(collection_name)
# Step 4: Create IAM role for Bedrock
role_arn = self._create_bedrock_role(unique_id, account_id, collection_arn)
# Step 5: Create Knowledge Base
self.knowledge_base_id = self._create_knowledge_base(
kb_name, role_arn, collection_arn
)
# Step 6: Create Data Source
self.data_source_id = self._create_data_source(kb_name)
verbose_logger.info(
f"Created KB infrastructure: kb_id={self.knowledge_base_id}, "
f"ds_id={self.data_source_id}, bucket={self.s3_bucket}"
)
def _create_s3_bucket(self, unique_id: str) -> str:
"""Create S3 bucket for KB data source."""
s3 = self._get_boto3_client("s3")
bucket_name = f"litellm-kb-{unique_id}"
verbose_logger.debug(f"Creating S3 bucket: {bucket_name}")
create_params: Dict[str, Any] = {"Bucket": bucket_name}
if self.aws_region_name != "us-east-1":
create_params["CreateBucketConfiguration"] = {
"LocationConstraint": self.aws_region_name
}
s3.create_bucket(**create_params)
self._created_resources["s3_bucket"] = bucket_name
verbose_logger.info(f"Created S3 bucket: {bucket_name}")
return bucket_name
def _create_opensearch_collection(
self, unique_id: str, account_id: str
) -> Tuple[str, str]:
"""Create OpenSearch Serverless collection for vector storage."""
oss = self._get_boto3_client("opensearchserverless")
collection_name = f"litellm-kb-{unique_id}"
verbose_logger.debug(f"Creating OpenSearch Serverless collection: {collection_name}")
# Create encryption policy
oss.create_security_policy(
name=f"{collection_name}-enc",
type="encryption",
policy=json.dumps({
"Rules": [{"ResourceType": "collection", "Resource": [f"collection/{collection_name}"]}],
"AWSOwnedKey": True,
}),
)
# Create network policy (public access for simplicity)
oss.create_security_policy(
name=f"{collection_name}-net",
type="network",
policy=json.dumps([{
"Rules": [{"ResourceType": "collection", "Resource": [f"collection/{collection_name}"]},
{"ResourceType": "dashboard", "Resource": [f"collection/{collection_name}"]}],
"AllowFromPublic": True,
}]),
)
# Create data access policy
oss.create_access_policy(
name=f"{collection_name}-access",
type="data",
policy=json.dumps([{
"Rules": [
{"ResourceType": "index", "Resource": [f"index/{collection_name}/*"], "Permission": ["aoss:*"]},
{"ResourceType": "collection", "Resource": [f"collection/{collection_name}"], "Permission": ["aoss:*"]},
],
"Principal": [f"arn:aws:iam::{account_id}:root"],
}]),
)
# Create collection
response = oss.create_collection(
name=collection_name,
type="VECTORSEARCH",
)
collection_id = response["createCollectionDetail"]["id"]
self._created_resources["opensearch_collection"] = collection_name
# Wait for collection to be active
verbose_logger.debug("Waiting for OpenSearch collection to be active...")
for _ in range(60): # 5 min timeout
status_response = oss.batch_get_collection(ids=[collection_id])
status = status_response["collectionDetails"][0]["status"]
if status == "ACTIVE":
break
time.sleep(5)
else:
raise TimeoutError("OpenSearch collection did not become active in time")
collection_arn = status_response["collectionDetails"][0]["arn"]
verbose_logger.info(f"Created OpenSearch collection: {collection_name}")
return collection_name, collection_arn
def _create_opensearch_index(self, collection_name: str):
"""Create vector index in OpenSearch collection."""
from opensearchpy import OpenSearch, RequestsHttpConnection
from requests_aws4auth import AWS4Auth
# Get credentials for signing
credentials = self.get_credentials(
aws_access_key_id=_get_str_or_none(self.vector_store_config.get("aws_access_key_id")),
aws_secret_access_key=_get_str_or_none(self.vector_store_config.get("aws_secret_access_key")),
aws_session_token=_get_str_or_none(self.vector_store_config.get("aws_session_token")),
aws_region_name=self.aws_region_name,
)
# Get collection endpoint
oss = self._get_boto3_client("opensearchserverless")
collections = oss.batch_get_collection(names=[collection_name])
endpoint = collections["collectionDetails"][0]["collectionEndpoint"]
host = endpoint.replace("https://", "")
auth = AWS4Auth(
credentials.access_key,
credentials.secret_key,
self.aws_region_name,
"aoss",
session_token=credentials.token,
)
client = OpenSearch(
hosts=[{"host": host, "port": 443}],
http_auth=auth,
use_ssl=True,
verify_certs=True,
connection_class=RequestsHttpConnection,
)
index_name = "bedrock-kb-index"
index_body = {
"settings": {
"index": {"knn": True, "knn.algo_param.ef_search": 512}
},
"mappings": {
"properties": {
"bedrock-knowledge-base-default-vector": {
"type": "knn_vector",
"dimension": 1024,
"method": {"engine": "faiss", "name": "hnsw", "space_type": "l2"},
},
"AMAZON_BEDROCK_METADATA": {"type": "text", "index": False},
"AMAZON_BEDROCK_TEXT_CHUNK": {"type": "text"},
}
},
}
client.indices.create(index=index_name, body=index_body)
verbose_logger.info(f"Created OpenSearch index: {index_name}")
def _create_bedrock_role(
self, unique_id: str, account_id: str, collection_arn: str
) -> str:
"""Create IAM role for Bedrock KB."""
iam = self._get_boto3_client("iam")
role_name = f"litellm-bedrock-kb-{unique_id}"
verbose_logger.debug(f"Creating IAM role: {role_name}")
trust_policy = {
"Version": "2012-10-17",
"Statement": [{
"Effect": "Allow",
"Principal": {"Service": "bedrock.amazonaws.com"},
"Action": "sts:AssumeRole",
"Condition": {
"StringEquals": {"aws:SourceAccount": account_id},
"ArnLike": {"aws:SourceArn": f"arn:aws:bedrock:{self.aws_region_name}:{account_id}:knowledge-base/*"},
},
}],
}
response = iam.create_role(
RoleName=role_name,
AssumeRolePolicyDocument=json.dumps(trust_policy),
)
role_arn = response["Role"]["Arn"]
self._created_resources["iam_role"] = role_name
# Attach permissions policy
permissions_policy = {
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Action": ["bedrock:InvokeModel"],
"Resource": [f"arn:aws:bedrock:{self.aws_region_name}::foundation-model/{self.embedding_model}"],
},
{
"Effect": "Allow",
"Action": ["aoss:APIAccessAll"],
"Resource": [collection_arn],
},
{
"Effect": "Allow",
"Action": ["s3:GetObject", "s3:ListBucket"],
"Resource": [f"arn:aws:s3:::{self.s3_bucket}", f"arn:aws:s3:::{self.s3_bucket}/*"],
},
],
}
iam.put_role_policy(
RoleName=role_name,
PolicyName=f"{role_name}-policy",
PolicyDocument=json.dumps(permissions_policy),
)
# Wait for role to propagate
time.sleep(10)
verbose_logger.info(f"Created IAM role: {role_arn}")
return role_arn
def _create_knowledge_base(
self, kb_name: str, role_arn: str, collection_arn: str
) -> str:
"""Create Bedrock Knowledge Base."""
bedrock_agent = self._get_boto3_client("bedrock-agent")
verbose_logger.debug(f"Creating Knowledge Base: {kb_name}")
response = bedrock_agent.create_knowledge_base(
name=kb_name,
roleArn=role_arn,
knowledgeBaseConfiguration={
"type": "VECTOR",
"vectorKnowledgeBaseConfiguration": {
"embeddingModelArn": f"arn:aws:bedrock:{self.aws_region_name}::foundation-model/{self.embedding_model}",
},
},
storageConfiguration={
"type": "OPENSEARCH_SERVERLESS",
"opensearchServerlessConfiguration": {
"collectionArn": collection_arn,
"fieldMapping": {
"metadataField": "AMAZON_BEDROCK_METADATA",
"textField": "AMAZON_BEDROCK_TEXT_CHUNK",
"vectorField": "bedrock-knowledge-base-default-vector",
},
"vectorIndexName": "bedrock-kb-index",
},
},
)
kb_id = response["knowledgeBase"]["knowledgeBaseId"]
self._created_resources["knowledge_base"] = kb_id
# Wait for KB to be active
verbose_logger.debug("Waiting for Knowledge Base to be active...")
for _ in range(30):
kb_status = bedrock_agent.get_knowledge_base(knowledgeBaseId=kb_id)
status = kb_status["knowledgeBase"]["status"]
if status == "ACTIVE":
break
time.sleep(2)
else:
raise TimeoutError("Knowledge Base did not become active in time")
verbose_logger.info(f"Created Knowledge Base: {kb_id}")
return kb_id
def _create_data_source(self, kb_name: str) -> str:
"""Create Data Source for the Knowledge Base."""
bedrock_agent = self._get_boto3_client("bedrock-agent")
verbose_logger.debug(f"Creating Data Source for KB: {self.knowledge_base_id}")
response = bedrock_agent.create_data_source(
knowledgeBaseId=self.knowledge_base_id,
name=f"{kb_name}-s3-source",
dataSourceConfiguration={
"type": "S3",
"s3Configuration": {
"bucketArn": f"arn:aws:s3:::{self.s3_bucket}",
"inclusionPrefixes": [self.s3_prefix],
},
},
)
ds_id = response["dataSource"]["dataSourceId"]
self._created_resources["data_source"] = ds_id
verbose_logger.info(f"Created Data Source: {ds_id}")
return ds_id
def _get_boto3_client(self, service_name: str):
"""Get a boto3 client for the specified service using BaseAWSLLM auth."""
try:
import boto3
except ImportError:
raise ImportError("boto3 is required for Bedrock ingestion. Install with: pip install boto3")
# Get credentials using BaseAWSLLM's get_credentials method
credentials = self.get_credentials(
aws_access_key_id=_get_str_or_none(self.vector_store_config.get("aws_access_key_id")),
aws_secret_access_key=_get_str_or_none(self.vector_store_config.get("aws_secret_access_key")),
aws_session_token=_get_str_or_none(self.vector_store_config.get("aws_session_token")),
aws_region_name=self.aws_region_name,
aws_session_name=_get_str_or_none(self.vector_store_config.get("aws_session_name")),
aws_profile_name=_get_str_or_none(self.vector_store_config.get("aws_profile_name")),
aws_role_name=_get_str_or_none(self.vector_store_config.get("aws_role_name")),
aws_web_identity_token=_get_str_or_none(self.vector_store_config.get("aws_web_identity_token")),
aws_sts_endpoint=_get_str_or_none(self.vector_store_config.get("aws_sts_endpoint")),
aws_external_id=_get_str_or_none(self.vector_store_config.get("aws_external_id")),
)
# Create session with credentials
session = boto3.Session(
aws_access_key_id=credentials.access_key,
aws_secret_access_key=credentials.secret_key,
aws_session_token=credentials.token,
region_name=self.aws_region_name,
)
return session.client(service_name)
async def embed(
self,
chunks: List[str],
) -> Optional[List[List[float]]]:
"""
Bedrock handles embedding internally - skip this step.
Returns:
None (Bedrock embeds when files are ingested)
"""
return None
async def store(
self,
file_content: Optional[bytes],
filename: Optional[str],
content_type: Optional[str],
chunks: List[str],
embeddings: Optional[List[List[float]]],
) -> Tuple[Optional[str], Optional[str]]:
"""
Store content in Bedrock Knowledge Base.
Bedrock workflow:
1. Auto-detect data source and S3 bucket (if not provided)
2. Upload file to S3 bucket
3. Start ingestion job
4. (Optional) Wait for ingestion to complete
Args:
file_content: Raw file bytes
filename: Name of the file
content_type: MIME type
chunks: Ignored - Bedrock handles chunking
embeddings: Ignored - Bedrock handles embedding
Returns:
Tuple of (knowledge_base_id, file_key)
"""
# Auto-detect data source and S3 bucket if needed
self._ensure_config_initialized()
if not file_content or not filename:
verbose_logger.warning("No file content or filename provided for Bedrock ingestion")
return _get_str_or_none(self.knowledge_base_id), None
# Step 1: Upload file to S3
s3_client = self._get_boto3_client("s3")
s3_key = f"{self.s3_prefix.rstrip('/')}/{filename}"
verbose_logger.debug(f"Uploading file to s3://{self.s3_bucket}/{s3_key}")
s3_client.put_object(
Bucket=self.s3_bucket,
Key=s3_key,
Body=file_content,
ContentType=content_type or "application/octet-stream",
)
verbose_logger.info(f"Uploaded file to s3://{self.s3_bucket}/{s3_key}")
# Step 2: Start ingestion job
bedrock_agent = self._get_boto3_client("bedrock-agent")
verbose_logger.debug(
f"Starting ingestion job for KB={self.knowledge_base_id}, DS={self.data_source_id}"
)
ingestion_response = bedrock_agent.start_ingestion_job(
knowledgeBaseId=self.knowledge_base_id,
dataSourceId=self.data_source_id,
)
job_id = ingestion_response["ingestionJob"]["ingestionJobId"]
verbose_logger.info(f"Started ingestion job: {job_id}")
# Step 3: Wait for ingestion (optional)
if self.wait_for_ingestion:
start_time = time.time()
while time.time() - start_time < self.ingestion_timeout:
job_status = bedrock_agent.get_ingestion_job(
knowledgeBaseId=self.knowledge_base_id,
dataSourceId=self.data_source_id,
ingestionJobId=job_id,
)
status = job_status["ingestionJob"]["status"]
verbose_logger.debug(f"Ingestion job {job_id} status: {status}")
if status == "COMPLETE":
stats = job_status["ingestionJob"].get("statistics", {})
verbose_logger.info(
f"Ingestion complete: {stats.get('numberOfNewDocumentsIndexed', 0)} docs indexed"
)
break
elif status == "FAILED":
failure_reasons = job_status["ingestionJob"].get("failureReasons", [])
verbose_logger.error(f"Ingestion failed: {failure_reasons}")
break
elif status in ("STARTING", "IN_PROGRESS"):
time.sleep(2)
else:
verbose_logger.warning(f"Unknown ingestion status: {status}")
break
return str(self.knowledge_base_id) if self.knowledge_base_id else None, s3_key
litellm/rag/ingestion/openai_ingestion.py
+111
View File
@@ -0,0 +1,111 @@
"""
OpenAI-specific RAG Ingestion implementation.
OpenAI handles embedding internally when files are attached to vector stores,
so this implementation skips the embedding step and directly uploads files.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast
import litellm
from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion
from litellm.vector_store_files.main import acreate as vector_store_file_acreate
from litellm.vector_stores.main import acreate as vector_store_acreate
if TYPE_CHECKING:
from litellm import Router
from litellm.types.rag import RAGIngestOptions
class OpenAIRAGIngestion(BaseRAGIngestion):
"""
OpenAI-specific RAG ingestion.
Key differences from base:
- Embedding is handled by OpenAI when attaching files to vector stores
- Files are uploaded and attached to vector stores directly
- Chunking is done by OpenAI's vector store (uses 'auto' strategy)
"""
def __init__(
self,
ingest_options: "RAGIngestOptions",
router: Optional["Router"] = None,
):
super().__init__(ingest_options=ingest_options, router=router)
async def embed(
self,
chunks: List[str],
) -> Optional[List[List[float]]]:
"""
OpenAI handles embedding internally - skip this step.
Returns:
None (OpenAI embeds when files are attached to vector store)
"""
# OpenAI handles embedding when files are attached to vector stores
return None
async def store(
self,
file_content: Optional[bytes],
filename: Optional[str],
content_type: Optional[str],
chunks: List[str],
embeddings: Optional[List[List[float]]],
) -> Tuple[Optional[str], Optional[str]]:
"""
Store content in OpenAI vector store.
OpenAI workflow:
1. Create vector store (if not provided)
2. Upload file to OpenAI
3. Attach file to vector store (OpenAI handles chunking/embedding)
Args:
file_content: Raw file bytes
filename: Name of the file
content_type: MIME type
chunks: Ignored - OpenAI handles chunking
embeddings: Ignored - OpenAI handles embedding
Returns:
Tuple of (vector_store_id, file_id)
"""
vector_store_id = self.vector_store_config.get("vector_store_id")
ttl_days = self.vector_store_config.get("ttl_days")
# Create vector store if not provided
if not vector_store_id:
expires_after = {"anchor": "last_active_at", "days": ttl_days} if ttl_days else None
create_response = await vector_store_acreate(
name=self.ingest_name or "litellm-rag-ingest",
custom_llm_provider="openai",
expires_after=expires_after,
)
vector_store_id = create_response.get("id")
# Upload file and attach to vector store
result_file_id = None
if file_content and filename and vector_store_id:
# Upload file to OpenAI
file_response = await litellm.acreate_file(
file=(filename, file_content, content_type or "application/octet-stream"),
purpose="assistants",
custom_llm_provider="openai",
)
result_file_id = file_response.id
# Attach file to vector store (OpenAI handles chunking/embedding)
await vector_store_file_acreate(
vector_store_id=vector_store_id,
file_id=result_file_id,
custom_llm_provider="openai",
chunking_strategy=cast(Optional[Dict[str, Any]], self.chunking_strategy),
)
return vector_store_id, result_file_id
litellm/rag/main.py
+240
View File
@@ -0,0 +1,240 @@
"""
RAG Ingest API for LiteLLM.
Provides an all-in-one API for document ingestion:
Upload -> (OCR) -> Chunk -> Embed -> Vector Store
"""
from __future__ import annotations
__all__ = ["ingest", "aingest"]
import asyncio
import contextvars
from functools import partial
from typing import TYPE_CHECKING, Any, Coroutine, Dict, Optional, Tuple, Type, Union
import httpx
import litellm
from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion
from litellm.rag.ingestion.bedrock_ingestion import BedrockRAGIngestion
from litellm.rag.ingestion.openai_ingestion import OpenAIRAGIngestion
from litellm.types.rag import RAGIngestOptions, RAGIngestResponse
from litellm.utils import client
if TYPE_CHECKING:
from litellm import Router
# Registry of provider-specific ingestion classes
INGESTION_REGISTRY: Dict[str, Type[BaseRAGIngestion]] = {
"openai": OpenAIRAGIngestion,
"bedrock": BedrockRAGIngestion,
}
def get_ingestion_class(provider: str) -> Type[BaseRAGIngestion]:
"""
Get the ingestion class for a given provider.
Args:
provider: The vector store provider name (e.g., 'openai')
Returns:
The ingestion class for the provider
Raises:
ValueError: If provider is not supported
"""
ingestion_class = INGESTION_REGISTRY.get(provider)
if ingestion_class is None:
supported = ", ".join(INGESTION_REGISTRY.keys())
raise ValueError(
f"Provider '{provider}' is not supported for RAG ingestion. "
f"Supported providers: {supported}"
)
return ingestion_class
async def _execute_ingest_pipeline(
ingest_options: RAGIngestOptions,
file_data: Optional[Tuple[str, bytes, str]] = None,
file_url: Optional[str] = None,
file_id: Optional[str] = None,
router: Optional["Router"] = None,
) -> RAGIngestResponse:
"""
Execute the RAG ingest pipeline using provider-specific implementation.
Args:
ingest_options: Configuration for the ingest pipeline
file_data: Tuple of (filename, content_bytes, content_type)
file_url: URL to fetch file from
file_id: Existing file ID to use
router: Optional LiteLLM router for load balancing
Returns:
RAGIngestResponse with status and IDs
"""
# Get provider from vector store config
vector_store_config = ingest_options.get("vector_store") or {}
provider = vector_store_config.get("custom_llm_provider", "openai")
# Get provider-specific ingestion class
ingestion_class = get_ingestion_class(provider)
# Create ingestion instance
ingestion = ingestion_class(
ingest_options=ingest_options,
router=router,
)
# Execute ingestion pipeline
return await ingestion.ingest(
file_data=file_data,
file_url=file_url,
file_id=file_id,
)
####### PUBLIC API ###################
@client
async def aingest(
ingest_options: Dict[str, Any],
file_data: Optional[Tuple[str, bytes, str]] = None,
file: Optional[Dict[str, str]] = None,
file_url: Optional[str] = None,
file_id: Optional[str] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
**kwargs,
) -> RAGIngestResponse:
"""
Async: Ingest a document into a vector store.
Args:
ingest_options: Configuration for the ingest pipeline
file_data: Tuple of (filename, content_bytes, content_type)
file: Dict with {filename, content (base64), content_type} - for JSON API
file_url: URL to fetch file from
file_id: Existing file ID to use
Example:
```python
response = await litellm.aingest(
ingest_options={
"vector_store": {"custom_llm_provider": "openai"}
},
file_url="https://example.com/doc.pdf",
)
```
"""
local_vars = locals()
try:
loop = asyncio.get_event_loop()
kwargs["aingest"] = True
func = partial(
ingest,
ingest_options=ingest_options,
file_data=file_data,
file=file,
file_url=file_url,
file_id=file_id,
timeout=timeout,
**kwargs,
)
ctx = contextvars.copy_context()
func_with_context = partial(ctx.run, func)
init_response = await loop.run_in_executor(None, func_with_context)
if asyncio.iscoroutine(init_response):
response = await init_response
else:
response = init_response
return response
except Exception as e:
raise litellm.exception_type(
model=None,
custom_llm_provider=ingest_options.get("vector_store", {}).get("custom_llm_provider"),
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
@client
def ingest(
ingest_options: Dict[str, Any],
file_data: Optional[Tuple[str, bytes, str]] = None,
file: Optional[Dict[str, str]] = None,
file_url: Optional[str] = None,
file_id: Optional[str] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
**kwargs,
) -> Union[RAGIngestResponse, Coroutine[Any, Any, RAGIngestResponse]]:
"""
Ingest a document into a vector store.
Args:
ingest_options: Configuration for the ingest pipeline
file_data: Tuple of (filename, content_bytes, content_type)
file: Dict with {filename, content (base64), content_type} - for JSON API
file_url: URL to fetch file from
file_id: Existing file ID to use
Example:
```python
response = litellm.ingest(
ingest_options={
"vector_store": {"custom_llm_provider": "openai"}
},
file_data=("doc.txt", b"Hello world", "text/plain"),
)
```
"""
import base64
local_vars = locals()
try:
_is_async = kwargs.pop("aingest", False) is True
router: Optional["Router"] = kwargs.get("router")
# Convert file dict to file_data tuple if provided
if file is not None and file_data is None:
filename = file.get("filename", "document")
content_b64 = file.get("content", "")
content_type = file.get("content_type", "application/octet-stream")
content_bytes = base64.b64decode(content_b64)
file_data = (filename, content_bytes, content_type)
if _is_async:
return _execute_ingest_pipeline(
ingest_options=ingest_options, # type: ignore
file_data=file_data,
file_url=file_url,
file_id=file_id,
router=router,
)
else:
return asyncio.get_event_loop().run_until_complete(
_execute_ingest_pipeline(
ingest_options=ingest_options, # type: ignore
file_data=file_data,
file_url=file_url,
file_id=file_id,
router=router,
)
)
except Exception as e:
raise litellm.exception_type(
model=None,
custom_llm_provider=ingest_options.get("vector_store", {}).get("custom_llm_provider"),
original_exception=e,
completion_kwargs=local_vars,
extra_kwargs=kwargs,
)
litellm/rag/text_splitters/__init__.py
+10
View File
@@ -0,0 +1,10 @@
"""
Text splitting utilities for RAG ingestion.
"""
from litellm.rag.text_splitters.recursive_character_text_splitter import (
RecursiveCharacterTextSplitter,
)
__all__ = ["RecursiveCharacterTextSplitter"]
litellm/rag/text_splitters/recursive_character_text_splitter.py
+135
View File
@@ -0,0 +1,135 @@
"""
RecursiveCharacterTextSplitter for RAG ingestion.
A simple implementation that splits text recursively by different separators.
"""
from typing import List, Optional
from litellm.constants import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE
class RecursiveCharacterTextSplitter:
"""
Split text recursively by different separators.
Tries to split by the first separator, then recursively splits
by subsequent separators if chunks are still too large.
"""
def __init__(
self,
chunk_size: int = DEFAULT_CHUNK_SIZE,
chunk_overlap: int = DEFAULT_CHUNK_OVERLAP,
separators: Optional[List[str]] = None,
):
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
self.separators = separators or ["\n\n", "\n", " ", ""]
def split_text(self, text: str) -> List[str]:
"""Split text into chunks."""
return self._split_text(text, self.separators)
def _split_text(self, text: str, separators: List[str], depth: int = 0) -> List[str]:
"""Recursively split text using separators."""
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
if depth > DEFAULT_MAX_RECURSE_DEPTH:
# Max depth reached, return text as-is split into chunk_size pieces
return [text[i:i + self.chunk_size] for i in range(0, len(text), self.chunk_size)]
final_chunks: List[str] = []
# Get the appropriate separator
separator = separators[-1]
new_separators: List[str] = []
for i, sep in enumerate(separators):
if sep == "":
separator = sep
break
if sep in text:
separator = sep
new_separators = separators[i + 1 :]
break
# Split by the chosen separator
if separator:
splits = text.split(separator)
else:
splits = list(text)
# Merge splits into chunks
good_splits: List[str] = []
for split in splits:
if len(split) < self.chunk_size:
good_splits.append(split)
else:
# Chunk is too big, merge what we have and recurse
if good_splits:
merged = self._merge_splits(good_splits, separator)
final_chunks.extend(merged)
good_splits = []
if new_separators:
# Recursively split with finer separators
other_chunks = self._split_text(split, new_separators, depth + 1)
final_chunks.extend(other_chunks)
else:
# No more separators, force split
final_chunks.extend(self._force_split(split))
# Merge remaining good splits
if good_splits:
merged = self._merge_splits(good_splits, separator)
final_chunks.extend(merged)
return final_chunks
def _merge_splits(self, splits: List[str], separator: str) -> List[str]:
"""Merge splits into chunks respecting chunk_size and chunk_overlap."""
chunks: List[str] = []
current_chunk: List[str] = []
current_length = 0
for split in splits:
split_len = len(split)
sep_len = len(separator) if current_chunk else 0
if current_length + split_len + sep_len > self.chunk_size:
if current_chunk:
chunk_text = separator.join(current_chunk).strip()
if chunk_text:
chunks.append(chunk_text)
# Handle overlap
while current_length > self.chunk_overlap and len(current_chunk) > 1:
removed = current_chunk.pop(0)
current_length -= len(removed) + len(separator)
current_chunk.append(split)
current_length += split_len + sep_len
# Add remaining
if current_chunk:
chunk_text = separator.join(current_chunk).strip()
if chunk_text:
chunks.append(chunk_text)
return chunks
def _force_split(self, text: str) -> List[str]:
"""Force split text by chunk_size when no separator works."""
chunks: List[str] = []
start = 0
while start < len(text):
end = start + self.chunk_size
chunk = text[start:end].strip()
if chunk:
chunks.append(chunk)
start = end - self.chunk_overlap if end < len(text) else len(text)
return chunks
litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py
+2 -1
View File
@@ -2,7 +2,7 @@ from typing import TYPE_CHECKING, Optional
from typing_extensions import TypedDict
from ..utils import CompletionTokensDetails, PromptTokensDetailsWrapper
from ..utils import CompletionTokensDetails, PromptTokensDetailsWrapper, ServerToolUse
class UsagePerChunk(TypedDict):
@@ -10,6 +10,7 @@ class UsagePerChunk(TypedDict):
completion_tokens: int
cache_creation_input_tokens: Optional[int]
cache_read_input_tokens: Optional[int]
server_tool_use: Optional[ServerToolUse]
web_search_requests: Optional[int]
completion_tokens_details: Optional[CompletionTokensDetails]
prompt_tokens_details: Optional[PromptTokensDetailsWrapper]
litellm/types/llms/anthropic.py
+74
View File
@@ -36,12 +36,20 @@ class AnthropicOutputSchema(TypedDict, total=False):
schema: Required[dict]
class AnthropicOutputConfig(TypedDict, total=False):
"""Configuration for controlling Claude's output behavior."""
effort: Literal["high", "medium", "low"]
class AnthropicMessagesTool(TypedDict, total=False):
name: Required[str]
description: str
input_schema: Optional[AnthropicInputSchema]
type: Literal["custom"]
cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
defer_loading: bool
allowed_callers: Optional[List[str]]
input_examples: Optional[List[Dict[str, Any]]]
class AnthropicComputerTool(TypedDict, total=False):
@@ -67,24 +75,78 @@ class AnthropicWebSearchTool(TypedDict, total=False):
cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
max_uses: Optional[int]
user_location: Optional[AnthropicWebSearchUserLocation]
defer_loading: Optional[bool]
allowed_callers: Optional[List[str]]
input_examples: Optional[List[Dict[str, Any]]]
class AnthropicHostedTools(TypedDict, total=False): # for bash_tool and text_editor
type: Required[str]
name: Required[str]
cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
defer_loading: Optional[bool]
allowed_callers: Optional[List[str]]
input_examples: Optional[List[Dict[str, Any]]]
class AnthropicCodeExecutionTool(TypedDict, total=False):
type: Required[str]
name: Required[Literal["code_execution"]]
cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
defer_loading: Optional[bool]
allowed_callers: Optional[List[str]]
input_examples: Optional[List[Dict[str, Any]]]
class AnthropicMemoryTool(TypedDict, total=False):
type: Required[str]
name: Required[Literal["memory"]]
cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
defer_loading: Optional[bool]
allowed_callers: Optional[List[str]]
input_examples: Optional[List[Dict[str, Any]]]
class AnthropicToolSearchToolRegex(TypedDict, total=False):
"""Tool search tool using regex patterns for tool discovery."""
type: Required[Literal["tool_search_tool_regex_20251119"]]
name: Required[str]
class AnthropicToolSearchToolBM25(TypedDict, total=False):
"""Tool search tool using BM25 algorithm for tool discovery."""
type: Required[Literal["tool_search_tool_bm25_20251119"]]
name: Required[str]
cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
defer_loading: Optional[bool]
allowed_callers: Optional[List[str]]
input_examples: Optional[List[Dict[str, Any]]]
class ToolReference(TypedDict, total=False):
"""Reference to a tool that should be expanded from deferred tools."""
type: Required[Literal["tool_reference"]]
tool_name: Required[str]
class DirectToolCaller(TypedDict, total=False):
"""Indicates a tool was called directly by Claude."""
type: Required[Literal["direct"]]
class CodeExecutionToolCaller(TypedDict, total=False):
"""Indicates a tool was called programmatically from code execution."""
type: Required[Literal["code_execution_20250825"]]
tool_id: Required[str] # ID of the code execution tool that made the call
ToolCaller = Union[DirectToolCaller, CodeExecutionToolCaller]
class AnthropicContainer(TypedDict, total=False):
"""Container metadata for code execution."""
id: Required[str]
expires_at: Optional[str] # ISO 8601 timestamp
AllAnthropicToolsValues = Union[
@@ -94,6 +156,8 @@ AllAnthropicToolsValues = Union[
AnthropicWebSearchTool,
AnthropicCodeExecutionTool,
AnthropicMemoryTool,
AnthropicToolSearchToolRegex,
AnthropicToolSearchToolBM25,
]
@@ -121,6 +185,7 @@ class AnthropicMessagesToolUseParam(TypedDict, total=False):
name: str
input: dict
cache_control: Optional[Union[dict, ChatCompletionCachedContent]]
caller: Optional[ToolCaller]
AnthropicMessagesAssistantMessageValues = Union[
@@ -372,6 +437,7 @@ class ToolUseBlock(TypedDict):
name: str
type: Literal["tool_use"]
caller: Optional[ToolCaller]
class TextBlock(TypedDict):
@@ -565,3 +631,11 @@ class ANTHROPIC_BETA_HEADER_VALUES(str, Enum):
WEB_FETCH_2025_09_10 = "web-fetch-2025-09-10"
CONTEXT_MANAGEMENT_2025_06_27 = "context-management-2025-06-27"
STRUCTURED_OUTPUT_2025_09_25 = "structured-outputs-2025-11-13"
ADVANCED_TOOL_USE_2025_11_20 = "advanced-tool-use-2025-11-20"
# Tool search beta header constant
ANTHROPIC_TOOL_SEARCH_BETA_HEADER = "advanced-tool-use-2025-11-20"
# Effort beta header constant
ANTHROPIC_EFFORT_BETA_HEADER = "effort-2025-11-24"
litellm/types/llms/openai.py
+12 -1
View File
@@ -1,6 +1,17 @@
from enum import Enum
from os import PathLike
from typing import IO, Any, Dict, Iterable, List, Literal, Mapping, Optional, Tuple, Union
from typing import (
IO,
Any,
Dict,
Iterable,
List,
Literal,
Mapping,
Optional,
Tuple,
Union,
)
import httpx
from openai._legacy_response import (
litellm/types/llms/vertex_ai.py
+5
View File
@@ -183,8 +183,13 @@ GeminiResponseModalities = Literal["TEXT", "IMAGE", "AUDIO", "VIDEO"]
GeminiImageAspectRatio = Literal["1:1", "2:3", "3:2", "3:4", "4:3", "9:16", "16:9", "21:9"]
GeminiImageSize = Literal["1K", "2K", "4K"]
class GeminiImageConfig(TypedDict, total=False):
aspectRatio: GeminiImageAspectRatio
imageSize: GeminiImageSize
class PrebuiltVoiceConfig(TypedDict):
voiceName: str
litellm/types/rag.py
+146
View File
@@ -0,0 +1,146 @@
"""
Type definitions for RAG (Retrieval Augmented Generation) Ingest API.
"""
from typing import Any, Dict, List, Literal, Optional, Union
from pydantic import BaseModel
from typing_extensions import TypedDict
class RAGChunkingStrategy(TypedDict, total=False):
"""
Chunking strategy config for RAG ingest using RecursiveCharacterTextSplitter.
See: https://docs.langchain.com/oss/python/langchain/rag
"""
chunk_size: int # Maximum size of chunks (default: 1000)
chunk_overlap: int # Overlap between chunks (default: 200)
separators: Optional[List[str]] # Custom separators for splitting
class RAGIngestOCROptions(TypedDict, total=False):
"""OCR configuration for RAG ingest pipeline."""
model: str # e.g., "mistral/mistral-ocr-latest"
class RAGIngestEmbeddingOptions(TypedDict, total=False):
"""Embedding configuration for RAG ingest pipeline."""
model: str # e.g., "text-embedding-3-small"
class OpenAIVectorStoreOptions(TypedDict, total=False):
"""
OpenAI vector store configuration.
Example (auto-create):
{"custom_llm_provider": "openai"}
Example (use existing):
{"custom_llm_provider": "openai", "vector_store_id": "vs_xxx"}
"""
custom_llm_provider: Literal["openai"]
vector_store_id: Optional[str] # Existing VS ID (auto-creates if not provided)
ttl_days: Optional[int] # Time-to-live in days for indexed content
class BedrockVectorStoreOptions(TypedDict, total=False):
"""
Bedrock Knowledge Base configuration.
Example (auto-create KB and all resources):
{"custom_llm_provider": "bedrock"}
Example (use existing KB):
{"custom_llm_provider": "bedrock", "vector_store_id": "KB_ID"}
Auto-creation creates: S3 bucket, OpenSearch Serverless collection,
IAM role, Knowledge Base, and Data Source.
"""
custom_llm_provider: Literal["bedrock"]
vector_store_id: Optional[str] # Existing KB ID (auto-creates if not provided)
# Bedrock-specific options
s3_bucket: Optional[str] # S3 bucket (auto-created if not provided)
s3_prefix: Optional[str] # S3 key prefix (default: "data/")
embedding_model: Optional[str] # Embedding model (default: amazon.titan-embed-text-v2:0)
data_source_id: Optional[str] # For existing KB: override auto-detected DS
wait_for_ingestion: Optional[bool] # Wait for completion (default: False - returns immediately)
ingestion_timeout: Optional[int] # Timeout in seconds if wait_for_ingestion=True (default: 300)
# AWS auth (uses BaseAWSLLM)
aws_access_key_id: Optional[str]
aws_secret_access_key: Optional[str]
aws_session_token: Optional[str]
aws_region_name: Optional[str] # default: us-west-2
aws_role_name: Optional[str]
aws_session_name: Optional[str]
aws_profile_name: Optional[str]
aws_web_identity_token: Optional[str]
aws_sts_endpoint: Optional[str]
aws_external_id: Optional[str]
# Union type for vector store options
RAGIngestVectorStoreOptions = Union[OpenAIVectorStoreOptions, BedrockVectorStoreOptions]
class RAGIngestOptions(TypedDict, total=False):
"""
Combined options for RAG ingest pipeline.
Unified interface - just specify custom_llm_provider:
Example (OpenAI):
from litellm.types.rag import RAGIngestOptions, OpenAIVectorStoreOptions
options: RAGIngestOptions = {
"vector_store": OpenAIVectorStoreOptions(
custom_llm_provider="openai",
vector_store_id="vs_xxx", # optional
)
}
Example (Bedrock):
from litellm.types.rag import RAGIngestOptions, BedrockVectorStoreOptions
options: RAGIngestOptions = {
"vector_store": BedrockVectorStoreOptions(
custom_llm_provider="bedrock",
vector_store_id="KB_ID", # optional - auto-creates if not provided
wait_for_ingestion=True,
)
}
"""
name: Optional[str] # Optional pipeline name for logging
ocr: Optional[RAGIngestOCROptions] # Optional OCR step
chunking_strategy: Optional[RAGChunkingStrategy] # RecursiveCharacterTextSplitter args
embedding: Optional[RAGIngestEmbeddingOptions] # Embedding model config
vector_store: RAGIngestVectorStoreOptions # OpenAI or Bedrock config
class RAGIngestResponse(TypedDict, total=False):
"""Response from RAG ingest API."""
id: str # Unique ingest job ID
status: Literal["completed", "in_progress", "failed"]
vector_store_id: str # The vector store ID (created or existing)
file_id: Optional[str] # The file ID in the vector store
class RAGIngestRequest(BaseModel):
"""Request body for RAG ingest API (for validation)."""
file_url: Optional[str] = None # URL to fetch file from
file_id: Optional[str] = None # Existing file ID
ingest_options: Dict[str, Any] # RAGIngestOptions as dict for flexibility
class Config:
extra = "allow" # Allow additional fields
litellm/types/utils.py
+2 -1
View File
@@ -999,7 +999,8 @@ class PromptTokensDetailsWrapper(
class ServerToolUse(BaseModel):
web_search_requests: Optional[int]
web_search_requests: Optional[int] = None
tool_search_requests: Optional[int] = None
class Usage(CompletionUsage):
litellm/utils.py
+11 -1
View File
@@ -3719,7 +3719,17 @@ def get_optional_params( # noqa: PLR0915
else False
),
)
elif bedrock_route == "openai":
optional_params = litellm.AmazonBedrockOpenAIConfig().map_openai_params(
model=model,
non_default_params=non_default_params,
optional_params=optional_params,
drop_params=(
drop_params
if drop_params is not None and isinstance(drop_params, bool)
else False
),
)
elif "anthropic" in bedrock_base_model and bedrock_route == "invoke":
if bedrock_base_model.startswith("anthropic.claude-3"):
optional_params = (
model_prices_and_context_window.json
+52
View File
@@ -24605,6 +24605,58 @@
"supports_tool_choice": true,
"supports_vision": true
},
"vertex_ai/claude-opus-4-5": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "vertex_ai-anthropic_models",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": true,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
},
"vertex_ai/claude-opus-4-5@20251101": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_read_input_token_cost": 5e-07,
"input_cost_per_token": 5e-06,
"litellm_provider": "vertex_ai-anthropic_models",
"max_input_tokens": 200000,
"max_output_tokens": 64000,
"max_tokens": 64000,
"mode": "chat",
"output_cost_per_token": 2.5e-05,
"search_context_cost_per_query": {
"search_context_size_high": 0.01,
"search_context_size_low": 0.01,
"search_context_size_medium": 0.01
},
"supports_assistant_prefill": true,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
},
"vertex_ai/claude-sonnet-4-5": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
pyproject.toml
+1
View File
@@ -67,6 +67,7 @@ polars = {version = "^1.31.0", optional = true, python = ">=3.10"}
semantic-router = {version = ">=0.1.12", optional = true, python = ">=3.9,<3.14"}
mlflow = {version = ">3.1.4", optional = true, python = ">=3.10"}
soundfile = {version = "^0.12.1", optional = true}
grpcio = ">=1.62.3,<1.68.0" # Constrain to < 1.68.0 to avoid resource exhausted bug (https://github.com/grpc/grpc/issues/38290). Minimum 1.62.3 required by grpcio-status.
[tool.poetry.extras]
proxy = [
requirements.txt
+1
View File
@@ -39,6 +39,7 @@ azure-storage-file-datalake==12.20.0 # for azure buck storage logging
opentelemetry-api==1.25.0
opentelemetry-sdk==1.25.0
opentelemetry-exporter-otlp==1.25.0
grpcio>=1.62.3,<1.68.0 # Constraint for opentelemetry-exporter-otlp-proto-grpc to avoid resource exhausted bug (https://github.com/grpc/grpc/issues/38290)
sentry_sdk==2.21.0 # for sentry error handling
detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
cryptography==44.0.1
tests/code_coverage_tests/recursive_detector.py
+3
View File
@@ -32,6 +32,9 @@ IGNORE_FUNCTIONS = [
"_redact_base64", # max depth set.
"_contains_vision_content", # max depth set.
"_read_all_bytes", # max depth set.
"_fix_enum_types", # max depth set.
"_collect_argument_paths", # max depth set.
"_split_text", # max depth set.
]
tests/litellm/llms/vertex_ai/gemini/test_transformation.py
+35 -1
View File
@@ -191,4 +191,38 @@ async def test__transform_request_body_image_config_snake_case():
assert "generationConfig" in rb
assert "image_config" in rb["generationConfig"]
assert rb["generationConfig"]["image_config"] == {"aspect_ratio": "16:9"}
assert rb["generationConfig"]["image_config"] == {"aspect_ratio": "16:9"}
@pytest.mark.asyncio
async def test__transform_request_body_image_config_with_image_size():
"""Test imageSize parameter support in imageConfig"""
model = "gemini-3-pro-image-preview"
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Generate a 4K image of Tokyo skyline"}
]
}
]
optional_params = {
"imageConfig": {"aspectRatio": "16:9", "imageSize": "4K"},
"responseModalities": ["Image"]
}
litellm_params = {}
transform_request_params = {
"messages": messages,
"model": model,
"optional_params": optional_params,
"custom_llm_provider": "gemini",
"litellm_params": litellm_params,
"cached_content": None,
}
rb: RequestBody = transformation._transform_request_body(**transform_request_params)
assert "generationConfig" in rb
assert "imageConfig" in rb["generationConfig"]
assert rb["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9"
assert rb["generationConfig"]["imageConfig"]["imageSize"] == "4K"
tests/llm_translation/test_bedrock_completion.py
+97
View File
@@ -3434,3 +3434,100 @@ async def test_bedrock_streaming_passthrough_test1(monkeypatch):
print(mock_callback.call_args.kwargs.keys())
assert "standard_logging_object" in mock_callback.call_args.kwargs["kwargs"]
assert "response_cost" in mock_callback.call_args.kwargs["kwargs"]
def test_bedrock_openai_imported_model():
"""
Test that Bedrock imported models using OpenAI format work correctly.
This test validates:
1. The request body follows OpenAI Chat Completions format
2. The URL is correctly constructed for Bedrock invoke endpoint
3. Messages with system, user roles and image_url content are preserved
"""
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
# Sample base64 image data (truncated for test)
sample_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
messages = [
{
"role": "system",
"content": "You are a helpful assistant that can analyze images.",
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "Spot the difference between the two images?",
},
{
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{sample_base64}"},
},
{
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{sample_base64}"},
},
],
},
]
with patch.object(client, "post") as mock_post:
try:
response = completion(
model="bedrock/openai/arn:aws:bedrock:us-east-1:117159858402:imported-model/m4gc1mrfuddy",
messages=messages,
max_tokens=300,
temperature=0.5,
client=client,
)
except Exception as e:
print(f"Exception (expected during mock): {e}")
mock_post.assert_called_once()
# Validate URL
url = mock_post.call_args.kwargs["url"]
print(f"URL: {url}")
assert "bedrock-runtime.us-east-1.amazonaws.com" in url
assert "arn:aws:bedrock:us-east-1:117159858402:imported-model/m4gc1mrfuddy" in url
assert "/invoke" in url
# Validate request body follows OpenAI format
request_body = json.loads(mock_post.call_args.kwargs["data"])
print(f"Request body: {json.dumps(request_body, indent=2)}")
# Check messages structure
assert "messages" in request_body
assert len(request_body["messages"]) == 2
# Check system message
system_msg = request_body["messages"][0]
assert system_msg["role"] == "system"
assert "helpful assistant" in system_msg["content"]
# Check user message with image content
user_msg = request_body["messages"][1]
assert user_msg["role"] == "user"
assert isinstance(user_msg["content"], list)
assert len(user_msg["content"]) == 3
# Check text content
assert user_msg["content"][0]["type"] == "text"
assert "Spot the difference" in user_msg["content"][0]["text"]
# Check image_url content
assert user_msg["content"][1]["type"] == "image_url"
assert "image_url" in user_msg["content"][1]
assert user_msg["content"][1]["image_url"]["url"].startswith("data:image/jpeg;base64,")
assert user_msg["content"][2]["type"] == "image_url"
assert "image_url" in user_msg["content"][2]
# Check max_tokens and temperature
assert request_body["max_tokens"] == 300
assert request_body["temperature"] == 0.5
tests/llm_translation/test_gemini.py
+1
View File
@@ -295,6 +295,7 @@ def test_gemini_image_generation():
[
"gemini/gemini-2.5-flash-image-preview",
"gemini/gemini-2.0-flash-preview-image-generation",
"gemini/gemini-3-pro-image-preview",
],
)
def test_gemini_flash_image_preview_models(model_name: str):
tests/proxy_unit_tests/test_proxy_token_counter.py
+17
View File
@@ -819,3 +819,20 @@ async def test_vertex_ai_anthropic_token_counting():
assert response.original_response is not None
assert "input_tokens" in response.original_response
assert response.original_response["input_tokens"] == 15
@pytest.mark.parametrize("vertex_location", ["global", "us-central1"])
def test_vertex_ai_gemini_token_counting_endpoint(vertex_location):
from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import (
VertexAIPartnerModelsTokenCounter,
)
endpoint = VertexAIPartnerModelsTokenCounter()._build_count_tokens_endpoint(
model="gemini-2.5-pro",
project_id="test-project",
vertex_location=vertex_location,
api_base=None,
)
if vertex_location == "global":
assert endpoint == "https://aiplatform.googleapis.com"
else:
assert endpoint == f"https://{vertex_location}-aiplatform.googleapis.com"
tests/test_litellm/completion_extras/test_litellm_responses_transformation_transformation.py
+136
View File
@@ -0,0 +1,136 @@
"""
Test for response_format to text.format conversion in completion -> responses bridge
"""
import pytest
from litellm.completion_extras.litellm_responses_transformation.transformation import (
LiteLLMResponsesTransformationHandler,
)
def test_transform_response_format_to_text_format_json_schema():
"""Test conversion of response_format with json_schema to text.format"""
handler = LiteLLMResponsesTransformationHandler()
# Chat Completion format
response_format = {
"type": "json_schema",
"json_schema": {
"name": "person_schema",
"schema": {
"type": "object",
"properties": {
"name": {"type": "string"},
"age": {"type": "integer"}
},
"required": ["name", "age"],
"additionalProperties": False
},
"strict": True
}
}
# Convert to Responses API format
result = handler._transform_response_format_to_text_format(response_format)
# Verify conversion
assert result is not None
assert "format" in result
assert result["format"]["type"] == "json_schema"
assert result["format"]["name"] == "person_schema"
assert result["format"]["strict"] is True
assert "schema" in result["format"]
assert result["format"]["schema"]["type"] == "object"
assert "properties" in result["format"]["schema"]
def test_transform_response_format_to_text_format_json_object():
"""Test conversion of response_format with json_object to text.format"""
handler = LiteLLMResponsesTransformationHandler()
response_format = {
"type": "json_object"
}
result = handler._transform_response_format_to_text_format(response_format)
assert result is not None
assert "format" in result
assert result["format"]["type"] == "json_object"
def test_transform_response_format_to_text_format_text():
"""Test conversion of response_format with text to text.format"""
handler = LiteLLMResponsesTransformationHandler()
response_format = {
"type": "text"
}
result = handler._transform_response_format_to_text_format(response_format)
assert result is not None
assert "format" in result
assert result["format"]["type"] == "text"
def test_transform_response_format_to_text_format_none():
"""Test that None input returns None"""
handler = LiteLLMResponsesTransformationHandler()
result = handler._transform_response_format_to_text_format(None)
assert result is None
def test_transform_request_with_response_format():
"""Test that transform_request correctly handles response_format parameter"""
handler = LiteLLMResponsesTransformationHandler()
messages = [
{"role": "user", "content": "Extract person info: John Doe, 30 years old"}
]
optional_params = {
"response_format": {
"type": "json_schema",
"json_schema": {
"name": "person_schema",
"schema": {
"type": "object",
"properties": {
"name": {"type": "string"},
"age": {"type": "integer"}
},
"required": ["name", "age"],
"additionalProperties": False
},
"strict": True
}
}
}
litellm_params = {}
headers = {}
# Mock logging object
class MockLoggingObj:
pass
litellm_logging_obj = MockLoggingObj()
result = handler.transform_request(
model="o3-pro",
messages=messages,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers,
litellm_logging_obj=litellm_logging_obj,
)
# Verify that text parameter was set with converted format
assert "text" in result
assert result["text"] is not None
assert "format" in result["text"]
assert result["text"]["format"]["type"] == "json_schema"
assert result["text"]["format"]["name"] == "person_schema"
assert "schema" in result["text"]["format"]
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
+81
View File
@@ -16,6 +16,7 @@ from litellm.types.utils import (
Function,
ModelResponseStream,
PromptTokensDetails,
ServerToolUse,
StreamingChoices,
Usage,
)
@@ -325,3 +326,83 @@ def test_stream_chunk_builder_litellm_usage_chunks():
assert usage.prompt_tokens == 50
assert usage.completion_tokens == 27
assert usage.total_tokens == 77
def test_stream_chunk_builder_anthropic_web_search():
# Prepare two mocked streaming chunks with usage split across them
chunk1 = ModelResponseStream(
id="chatcmpl-mocked-usage-1",
created=1745513206,
model="claude-sonnet-4-5-20250929",
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(
provider_specific_fields=None,
content="",
role=None,
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields=None,
stream_options={"include_usage": True},
usage=Usage(
completion_tokens=0,
prompt_tokens=50,
total_tokens=50,
completion_tokens_details=None,
server_tool_use=ServerToolUse(web_search_requests=2),
prompt_tokens_details=None,
),
)
chunk2 = ModelResponseStream(
id="chatcmpl-mocked-usage-1",
created=1745513207,
model="claude-sonnet-4-5-20250929",
object="chat.completion.chunk",
system_fingerprint=None,
choices=[
StreamingChoices(
finish_reason="stop",
index=0,
delta=Delta(
provider_specific_fields=None,
content=None,
role=None,
function_call=None,
tool_calls=None,
audio=None,
),
logprobs=None,
)
],
provider_specific_fields=None,
stream_options={"include_usage": True},
usage=Usage(
completion_tokens=27,
prompt_tokens=0,
total_tokens=27,
completion_tokens_details=None,
prompt_tokens_details=None,
),
)
chunks = [chunk1, chunk2]
processor = ChunkProcessor(chunks=chunks)
usage = processor.calculate_usage(
chunks=chunks, model="claude-sonnet-4-5-20250929", completion_output=""
)
assert usage.prompt_tokens == 50
assert usage.completion_tokens == 27
assert usage.total_tokens == 77
assert usage.server_tool_use['web_search_requests'] == 2
tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py
+741
View File
@@ -556,3 +556,744 @@ def test_anthropic_structured_output_beta_header():
"structured-outputs-2025-11-13"
in response["raw_request_headers"]["anthropic-beta"]
)
# ============ Tool Search Tests ============
def test_tool_search_regex_detection():
"""Test that tool search regex tools are properly detected"""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
config = AnthropicModelInfo()
# Test with tool search regex tool
tools = [
{
"type": "tool_search_tool_regex_20251119",
"name": "tool_search_tool_regex"
}
]
assert config.is_tool_search_used(tools) is True
# Test without tool search
tools = [
{
"type": "function",
"function": {"name": "get_weather"}
}
]
assert config.is_tool_search_used(tools) is False
def test_tool_search_bm25_detection():
"""Test that tool search BM25 tools are properly detected"""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
config = AnthropicModelInfo()
# Test with tool search BM25 tool
tools = [
{
"type": "tool_search_tool_bm25_20251119",
"name": "tool_search_tool_bm25"
}
]
assert config.is_tool_search_used(tools) is True
def test_tool_search_beta_header():
"""Test that tool search beta header is automatically added"""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
config = AnthropicModelInfo()
headers = config.get_anthropic_headers(
api_key="test-key",
tool_search_used=True,
)
assert "anthropic-beta" in headers
assert "advanced-tool-use-2025-11-20" in headers["anthropic-beta"]
def test_tool_search_regex_mapping():
"""Test that tool search regex tools are properly mapped"""
config = AnthropicConfig()
tool = {
"type": "tool_search_tool_regex_20251119",
"name": "tool_search_tool_regex"
}
mapped_tool, mcp_server = config._map_tool_helper(tool)
assert mapped_tool is not None
assert mapped_tool["type"] == "tool_search_tool_regex_20251119"
assert mapped_tool["name"] == "tool_search_tool_regex"
assert mcp_server is None
def test_tool_search_bm25_mapping():
"""Test that tool search BM25 tools are properly mapped"""
config = AnthropicConfig()
tool = {
"type": "tool_search_tool_bm25_20251119",
"name": "tool_search_tool_bm25"
}
mapped_tool, mcp_server = config._map_tool_helper(tool)
assert mapped_tool is not None
assert mapped_tool["type"] == "tool_search_tool_bm25_20251119"
assert mapped_tool["name"] == "tool_search_tool_bm25"
assert mcp_server is None
def test_deferred_tools_separation():
"""Test that deferred and non-deferred tools are properly separated"""
config = AnthropicConfig()
tools = [
{
"type": "tool_search_tool_regex_20251119",
"name": "tool_search_tool_regex"
},
{
"type": "function",
"function": {"name": "get_weather"},
"defer_loading": True
},
{
"type": "function",
"function": {"name": "search_files"},
"defer_loading": False
}
]
non_deferred, deferred = config._separate_deferred_tools(tools)
assert len(non_deferred) == 2 # tool_search and search_files
assert len(deferred) == 1 # get_weather
def test_server_tool_use_in_response():
"""Test that server_tool_use blocks are parsed correctly"""
config = AnthropicConfig()
completion_response = {
"content": [
{
"type": "server_tool_use",
"id": "srvtoolu_01ABC123",
"name": "tool_search_tool_regex",
"input": {"query": "weather"}
}
]
}
text, citations, thinking_blocks, reasoning_content, tool_calls = config.extract_response_content(
completion_response
)
assert len(tool_calls) == 1
assert tool_calls[0]["id"] == "srvtoolu_01ABC123"
assert tool_calls[0]["function"]["name"] == "tool_search_tool_regex"
def test_tool_search_usage_tracking():
"""Test that tool_search_requests are tracked in usage"""
config = AnthropicConfig()
usage_object = {
"input_tokens": 100,
"output_tokens": 50,
"server_tool_use": {
"tool_search_requests": 2
}
}
usage = config.calculate_usage(usage_object=usage_object, reasoning_content=None)
assert usage.server_tool_use is not None
assert usage.server_tool_use.tool_search_requests == 2
def test_tool_reference_expansion():
"""Test that tool_reference blocks are expanded correctly"""
config = AnthropicConfig()
deferred_tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather"
}
}
]
content = [
{"type": "text", "text": "I'll search for tools"},
{"type": "tool_reference", "tool_name": "get_weather"}
]
expanded = config._expand_tool_references(content, deferred_tools)
assert len(expanded) == 2
assert expanded[0]["type"] == "text"
assert expanded[1]["type"] == "function"
assert expanded[1]["function"]["name"] == "get_weather"
def test_defer_loading_preserved_in_transformation():
"""Test that defer_loading parameter is preserved when transforming tools"""
config = AnthropicConfig()
tool = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather information",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
},
"required": ["location"]
}
},
"defer_loading": True
}
mapped_tool, mcp_server = config._map_tool_helper(tool)
assert mapped_tool is not None
assert mapped_tool.get("defer_loading") is True
assert mapped_tool["name"] == "get_weather"
assert mcp_server is None
def test_tool_search_complete_response_parsing():
"""Test parsing a complete tool search response with server_tool_use and tool_search_tool_result blocks"""
config = AnthropicConfig()
# Simulating actual Anthropic API response with tool search
completion_response = {
"content": [
{
"type": "text",
"text": "I'll search for weather-related tools that can help you."
},
{
"type": "server_tool_use",
"id": "srvtoolu_015i6aVA2niwzv4RG4DtnxDJ",
"name": "tool_search_tool_regex",
"input": {"pattern": "weather", "limit": 5},
"caller": {"type": "direct"}
},
{
"type": "tool_search_tool_result",
"tool_use_id": "srvtoolu_015i6aVA2niwzv4RG4DtnxDJ",
"content": {
"type": "tool_search_tool_search_result",
"tool_references": [{"type": "tool_reference", "tool_name": "get_weather"}]
}
},
{
"type": "text",
"text": "Great! I found a weather tool."
},
{
"type": "tool_use",
"id": "toolu_01CrCNx4ntSaeeV9iArT4JfQ",
"name": "get_weather",
"input": {"location": "San Francisco"}
}
],
"usage": {
"input_tokens": 1639,
"output_tokens": 170,
"server_tool_use": {"web_search_requests": 0}
}
}
# Extract content
text, citations, thinking_blocks, reasoning_content, tool_calls = config.extract_response_content(
completion_response
)
# Verify text extraction (should concatenate both text blocks)
assert "I'll search for weather-related tools" in text
assert "Great! I found a weather tool" in text
# Verify tool calls (should have both server_tool_use and tool_use)
assert len(tool_calls) == 2
assert tool_calls[0]["function"]["name"] == "tool_search_tool_regex"
assert tool_calls[1]["function"]["name"] == "get_weather"
# Verify usage calculation counts tool_search_requests from content
usage = config.calculate_usage(
usage_object=completion_response["usage"],
reasoning_content=None,
completion_response=completion_response
)
assert usage.server_tool_use is not None
assert usage.server_tool_use.web_search_requests == 0
assert usage.server_tool_use.tool_search_requests == 1 # Counted from server_tool_use blocks
def test_allowed_callers_field_preservation():
"""Test that allowed_callers field is preserved during tool transformation."""
config = AnthropicConfig()
# Test with top-level allowed_callers
tool_with_allowed_callers = {
"type": "function",
"function": {
"name": "query_database",
"description": "Execute a SQL query",
"parameters": {
"type": "object",
"properties": {
"sql": {"type": "string"}
},
"required": ["sql"]
}
},
"allowed_callers": ["code_execution_20250825"]
}
transformed_tool, _ = config._map_tool_helper(tool_with_allowed_callers)
assert transformed_tool is not None
assert "allowed_callers" in transformed_tool
assert transformed_tool["allowed_callers"] == ["code_execution_20250825"]
def test_programmatic_tool_calling_beta_header():
"""Test that beta header is automatically added when programmatic tool calling is detected."""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
model_info = AnthropicModelInfo()
# Test detection with allowed_callers
tools = [
{
"type": "code_execution_20250825",
"name": "code_execution"
},
{
"type": "function",
"function": {
"name": "query_database",
"description": "Execute a SQL query",
"parameters": {"type": "object", "properties": {}}
},
"allowed_callers": ["code_execution_20250825"]
}
]
is_programmatic = model_info.is_programmatic_tool_calling_used(tools)
assert is_programmatic is True
# Test header generation
headers = model_info.get_anthropic_headers(
api_key="test-key",
programmatic_tool_calling_used=True
)
assert "anthropic-beta" in headers
assert "advanced-tool-use-2025-11-20" in headers["anthropic-beta"]
def test_caller_field_in_response():
"""Test that caller field is correctly parsed from tool_use blocks."""
config = AnthropicConfig()
# Mock response with programmatic tool call
completion_response = {
"id": "msg_test",
"type": "message",
"role": "assistant",
"content": [
{
"type": "text",
"text": "I'll query the database."
},
{
"type": "tool_use",
"id": "toolu_123",
"name": "query_database",
"input": {"sql": "SELECT * FROM users"},
"caller": {
"type": "code_execution_20250825",
"tool_id": "srvtoolu_abc"
}
}
],
"stop_reason": "tool_use",
"usage": {"input_tokens": 100, "output_tokens": 50}
}
text, citations, thinking, reasoning, tool_calls = config.extract_response_content(completion_response)
assert len(tool_calls) == 1
assert tool_calls[0]["id"] == "toolu_123"
assert tool_calls[0]["function"]["name"] == "query_database"
assert "caller" in tool_calls[0]
assert tool_calls[0]["caller"]["type"] == "code_execution_20250825"
assert tool_calls[0]["caller"]["tool_id"] == "srvtoolu_abc"
def test_code_execution_20250825_tool_type():
"""Test that code_execution_20250825 tool type is handled correctly."""
config = AnthropicConfig()
tool = {
"type": "code_execution_20250825",
"name": "code_execution"
}
transformed_tool, _ = config._map_tool_helper(tool)
assert transformed_tool is not None
assert transformed_tool["type"] == "code_execution_20250825"
assert transformed_tool["name"] == "code_execution"
def test_allowed_callers_in_function_field():
"""Test that allowed_callers in function field is also preserved."""
config = AnthropicConfig()
# Test with function.allowed_callers
tool = {
"type": "function",
"function": {
"name": "query_database",
"description": "Execute a SQL query",
"parameters": {
"type": "object",
"properties": {
"sql": {"type": "string"}
},
"required": ["sql"]
},
"allowed_callers": ["code_execution_20250825"]
}
}
transformed_tool, _ = config._map_tool_helper(tool)
assert transformed_tool is not None
assert "allowed_callers" in transformed_tool
assert transformed_tool["allowed_callers"] == ["code_execution_20250825"]
def test_input_examples_field_preservation():
"""Test that input_examples field is preserved during tool transformation."""
config = AnthropicConfig()
# Test with top-level input_examples
tool_with_examples = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get the current weather in a given location",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"},
"unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}
},
"required": ["location"]
}
},
"input_examples": [
{"location": "San Francisco, CA", "unit": "fahrenheit"},
{"location": "Tokyo, Japan", "unit": "celsius"}
]
}
transformed_tool, _ = config._map_tool_helper(tool_with_examples)
assert transformed_tool is not None
assert "input_examples" in transformed_tool
assert len(transformed_tool["input_examples"]) == 2
assert transformed_tool["input_examples"][0]["location"] == "San Francisco, CA"
def test_input_examples_beta_header():
"""Test that beta header is automatically added when input_examples is detected."""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
model_info = AnthropicModelInfo()
# Test detection with input_examples
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather information",
"parameters": {"type": "object", "properties": {}}
},
"input_examples": [
{"location": "San Francisco, CA"}
]
}
]
is_examples_used = model_info.is_input_examples_used(tools)
assert is_examples_used is True
# Test header generation
headers = model_info.get_anthropic_headers(
api_key="test-key",
input_examples_used=True
)
assert "anthropic-beta" in headers
assert "advanced-tool-use-2025-11-20" in headers["anthropic-beta"]
def test_input_examples_in_function_field():
"""Test that input_examples in function field is also preserved."""
config = AnthropicConfig()
# Test with function.input_examples
tool = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather information",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
},
"required": ["location"]
},
"input_examples": [
{"location": "Paris, France"},
{"location": "London, UK"}
]
}
}
transformed_tool, _ = config._map_tool_helper(tool)
assert transformed_tool is not None
assert "input_examples" in transformed_tool
assert len(transformed_tool["input_examples"]) == 2
def test_input_examples_with_other_features():
"""Test that input_examples works alongside other tool features."""
config = AnthropicConfig()
# Tool with input_examples, defer_loading, and allowed_callers
tool = {
"type": "function",
"function": {
"name": "query_database",
"description": "Execute a SQL query",
"parameters": {
"type": "object",
"properties": {
"sql": {"type": "string"}
},
"required": ["sql"]
}
},
"input_examples": [
{"sql": "SELECT * FROM users WHERE id = 1"}
],
"defer_loading": True,
"allowed_callers": ["code_execution_20250825"]
}
transformed_tool, _ = config._map_tool_helper(tool)
assert transformed_tool is not None
assert "input_examples" in transformed_tool
assert "defer_loading" in transformed_tool
assert "allowed_callers" in transformed_tool
assert transformed_tool["defer_loading"] is True
assert transformed_tool["allowed_callers"] == ["code_execution_20250825"]
def test_input_examples_empty_list_not_added():
"""Test that empty input_examples list is not added to transformed tool."""
config = AnthropicConfig()
# Tool with empty input_examples
tool = {
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather information",
"parameters": {
"type": "object",
"properties": {
"location": {"type": "string"}
},
"required": ["location"]
}
},
"input_examples": []
}
transformed_tool, _ = config._map_tool_helper(tool)
assert transformed_tool is not None
# Empty list should not be added
assert "input_examples" not in transformed_tool or len(transformed_tool.get("input_examples", [])) == 0
# ============ Effort Parameter Tests ============
def test_effort_output_config_preservation():
"""Test that output_config with effort is preserved in transformation."""
config = AnthropicConfig()
messages = [{"role": "user", "content": "Analyze this code"}]
optional_params = {
"output_config": {
"effort": "medium"
}
}
result = config.transform_request(
model="claude-opus-4-5-20251101",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
)
assert "output_config" in result
assert result["output_config"]["effort"] == "medium"
def test_effort_beta_header_injection():
"""Test that effort beta header is automatically added when output_config is detected."""
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
model_info = AnthropicModelInfo()
# Test with effort parameter
optional_params = {
"output_config": {
"effort": "low"
}
}
effort_used = model_info.is_effort_used(optional_params=optional_params)
assert effort_used is True
headers = model_info.get_anthropic_headers(
api_key="test-key",
effort_used=effort_used
)
assert "anthropic-beta" in headers
assert "effort-2025-11-24" in headers["anthropic-beta"]
def test_effort_validation():
"""Test that only valid effort values are accepted."""
config = AnthropicConfig()
messages = [{"role": "user", "content": "Test"}]
# Valid values should work
for effort in ["high", "medium", "low"]:
optional_params = {"output_config": {"effort": effort}}
result = config.transform_request(
model="claude-opus-4-5-20251101",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
)
assert result["output_config"]["effort"] == effort
# Invalid value should raise error
with pytest.raises(ValueError, match="Invalid effort value"):
optional_params = {"output_config": {"effort": "invalid"}}
config.transform_request(
model="claude-opus-4-5-20251101",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
)
def test_effort_with_claude_opus_45():
"""Test effort parameter works with Claude Opus 4.5 model."""
config = AnthropicConfig()
messages = [{"role": "user", "content": "Complex analysis task"}]
optional_params = {
"output_config": {
"effort": "high"
}
}
result = config.transform_request(
model="claude-opus-4-5-20251101",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
)
assert "output_config" in result
assert result["output_config"]["effort"] == "high"
assert result["model"] == "claude-opus-4-5-20251101"
def test_effort_with_other_features():
"""Test effort works alongside other features (thinking, tools)."""
config = AnthropicConfig()
messages = [{"role": "user", "content": "Use tools efficiently"}]
tools = [
{
"type": "function",
"function": {
"name": "get_data",
"description": "Get data",
"parameters": {
"type": "object",
"properties": {
"query": {"type": "string"}
},
"required": ["query"]
}
}
}
]
optional_params = {
"output_config": {
"effort": "low"
},
"tools": tools,
"thinking": {
"type": "enabled",
"budget_tokens": 1000
}
}
result = config.transform_request(
model="claude-opus-4-5-20251101",
messages=messages,
optional_params=optional_params,
litellm_params={},
headers={}
)
# Verify all features are present
assert "output_config" in result
assert result["output_config"]["effort"] == "low"
assert "tools" in result
assert len(result["tools"]) > 0
assert "thinking" in result
tests/test_litellm/proxy/guardrails/test_pillar_guardrails.py
+227
View File
@@ -234,6 +234,22 @@ def pillar_async_response():
)
@pytest.fixture
def user_api_key_dict_with_context():
"""Fixture providing UserAPIKeyAuth with complete context."""
return UserAPIKeyAuth(
token="hashed-test-token",
key_name="production-api-key",
key_alias="prod-key",
user_id="user-123",
user_email="test@example.com",
team_id="team-456",
team_alias="engineering-team",
org_id="org-789",
metadata={"environment": "production", "region": "us-east-1"},
)
@pytest.fixture
def mock_llm_response_with_tools():
"""Fixture providing a mock LLM response with tool calls."""
@@ -502,6 +518,217 @@ async def test_pre_call_hook_custom_header_overrides(
assert captured_headers.get("plr_evidence") == "false"
# =========================================================================
# LITELLM KEY CONTEXT HEADER TESTS
# =========================================================================
@pytest.mark.asyncio
async def test_litellm_context_headers_automatically_added(
sample_request_data,
user_api_key_dict_with_context,
dual_cache,
pillar_clean_response,
):
"""Test that LiteLLM context headers are automatically added (always enabled)."""
guardrail = PillarGuardrail(
guardrail_name="pillar-context-enabled",
api_key="test-pillar-key",
api_base="https://api.pillar.security",
)
captured_headers: Dict[str, str] = {}
async def _mock_post(*args, **kwargs):
captured_headers.update(kwargs.get("headers", {}))
return pillar_clean_response
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new=_mock_post,
):
await guardrail.async_pre_call_hook(
data=sample_request_data,
cache=dual_cache,
user_api_key_dict=user_api_key_dict_with_context,
call_type="completion",
)
# Verify LiteLLM context headers are present
assert "X-LiteLLM-Key-Name" in captured_headers
assert captured_headers["X-LiteLLM-Key-Name"] == "production-api-key"
assert "X-LiteLLM-Key-Alias" in captured_headers
assert captured_headers["X-LiteLLM-Key-Alias"] == "prod-key"
assert "X-LiteLLM-User-Id" in captured_headers
assert captured_headers["X-LiteLLM-User-Id"] == "user-123"
assert "X-LiteLLM-User-Email" in captured_headers
assert captured_headers["X-LiteLLM-User-Email"] == "test@example.com"
assert "X-LiteLLM-Team-Id" in captured_headers
assert captured_headers["X-LiteLLM-Team-Id"] == "team-456"
assert "X-LiteLLM-Team-Name" in captured_headers
assert captured_headers["X-LiteLLM-Team-Name"] == "engineering-team"
assert "X-LiteLLM-Org-Id" in captured_headers
assert captured_headers["X-LiteLLM-Org-Id"] == "org-789"
# Metadata is NOT sent (may contain sensitive information)
assert "X-LiteLLM-Metadata" not in captured_headers
@pytest.mark.asyncio
async def test_litellm_context_with_partial_fields(
sample_request_data,
dual_cache,
pillar_clean_response,
):
"""Test that partial LiteLLM context (only some fields present) is handled correctly."""
# Create UserAPIKeyAuth with only some fields populated
partial_context = UserAPIKeyAuth(
user_id="user-only",
team_id="team-only",
)
guardrail = PillarGuardrail(
guardrail_name="pillar-partial-context",
api_key="test-pillar-key",
api_base="https://api.pillar.security",
pass_litellm_key_header=True,
)
captured_headers: Dict[str, str] = {}
async def _mock_post(*args, **kwargs):
captured_headers.update(kwargs.get("headers", {}))
return pillar_clean_response
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new=_mock_post,
):
await guardrail.async_pre_call_hook(
data=sample_request_data,
cache=dual_cache,
user_api_key_dict=partial_context,
call_type="completion",
)
# Verify only populated fields are present
assert "X-LiteLLM-User-Id" in captured_headers
assert captured_headers["X-LiteLLM-User-Id"] == "user-only"
assert "X-LiteLLM-Team-Id" in captured_headers
assert captured_headers["X-LiteLLM-Team-Id"] == "team-only"
# Verify empty fields are not present
assert "X-LiteLLM-Key-Name" not in captured_headers
assert "X-LiteLLM-User-Email" not in captured_headers
# =========================================================================
# MULTI-MODAL CONTENT TESTS
# =========================================================================
@pytest.mark.asyncio
async def test_multimodal_image_url_support(
user_api_key_dict,
dual_cache,
pillar_clean_response,
):
"""Test that messages with image URLs are properly handled."""
multimodal_data = {
"model": "gpt-4-vision-preview",
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "What's in this image?"},
{
"type": "image_url",
"image_url": {
"url": "https://example.com/image.jpg",
"detail": "high",
},
},
],
}
],
}
guardrail = PillarGuardrail(
guardrail_name="pillar-multimodal",
api_key="test-pillar-key",
api_base="https://api.pillar.security",
)
captured_payload: Dict[str, Any] = {}
async def _mock_post(*args, **kwargs):
captured_payload.update(kwargs.get("json", {}))
return pillar_clean_response
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new=_mock_post,
):
result = await guardrail.async_pre_call_hook(
data=multimodal_data,
cache=dual_cache,
user_api_key_dict=user_api_key_dict,
call_type="completion",
)
# Verify multimodal message structure is preserved
assert result == multimodal_data
assert "messages" in captured_payload
assert len(captured_payload["messages"]) == 1
assert isinstance(captured_payload["messages"][0]["content"], list)
assert captured_payload["messages"][0]["content"][1]["type"] == "image_url"
@pytest.mark.asyncio
async def test_multimodal_with_attachments(
user_api_key_dict,
dual_cache,
pillar_clean_response,
):
"""Test that messages with file attachments are properly handled."""
multimodal_data = {
"model": "gpt-4",
"messages": [
{
"role": "user",
"content": "Analyze this document",
"attachments": [
{
"file_id": "file-abc123",
"tools": [{"type": "code_interpreter"}],
}
],
}
],
}
guardrail = PillarGuardrail(
guardrail_name="pillar-attachments",
api_key="test-pillar-key",
api_base="https://api.pillar.security",
)
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=pillar_clean_response,
):
result = await guardrail.async_pre_call_hook(
data=multimodal_data,
cache=dual_cache,
user_api_key_dict=user_api_key_dict,
call_type="completion",
)
# Verify attachment structure is preserved
assert result == multimodal_data
assert result["messages"][0]["attachments"] is not None
# ============================================================================
# EDGE CASE TESTS
# ============================================================================
tests/vector_store_tests/rag/base_rag_tests.py
+173
View File
@@ -0,0 +1,173 @@
"""
Base RAG test class that enforces common tests across all providers.
Providers should inherit from BaseRAGTest and implement the abstract methods.
"""
import os
import sys
import uuid
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
import litellm
from litellm.types.rag import (
RAGIngestOptions,
OpenAIVectorStoreOptions,
BedrockVectorStoreOptions,
)
class BaseRAGTest(ABC):
"""
Abstract base test class for RAG ingestion tests.
Providers should inherit from this class and implement:
- get_base_ingest_options(): Returns provider-specific ingest options
- query_vector_store(): Queries the vector store after ingestion
"""
@abstractmethod
def get_base_ingest_options(self) -> RAGIngestOptions:
"""
Must return the base ingest options for the provider.
Example for OpenAI:
return {
"vector_store": OpenAIVectorStoreOptions(
custom_llm_provider="openai",
)
}
Example for Bedrock:
return {
"vector_store": BedrockVectorStoreOptions(
custom_llm_provider="bedrock",
)
}
"""
pass
@abstractmethod
async def query_vector_store(
self,
vector_store_id: str,
query: str,
) -> Optional[Dict[str, Any]]:
"""
Query the vector store to verify ingestion.
Args:
vector_store_id: The ID of the vector store to query
query: The search query
Returns:
Search results dict or None if no results found
"""
pass
def get_unique_filename(self, prefix: str = "test") -> str:
"""Generate a unique filename for test documents."""
unique_id = uuid.uuid4().hex[:8]
return f"{prefix}_{unique_id}.txt", unique_id
@pytest.mark.asyncio
async def test_basic_ingest(self):
"""
Test basic text file ingestion to vector store.
"""
litellm._turn_on_debug()
filename, unique_id = self.get_unique_filename("basic_ingest")
text_content = f"Test document {unique_id} for RAG ingestion.".encode("utf-8")
file_data = (filename, text_content, "text/plain")
ingest_options = self.get_base_ingest_options()
ingest_options["name"] = f"test-basic-ingest-{unique_id}"
try:
response = await litellm.rag.aingest(
ingest_options=ingest_options,
file_data=file_data,
)
print(f"RAG Ingest Response: {response}")
assert "id" in response
assert response["id"].startswith("ingest_")
assert "status" in response
assert response["status"] in ["completed", "failed"]
assert "vector_store_id" in response
if response["status"] == "completed":
assert response["vector_store_id"]
print(f"Vector store ID: {response['vector_store_id']}")
except litellm.InternalServerError:
pytest.skip("Skipping test due to litellm.InternalServerError")
@pytest.mark.asyncio
async def test_ingest_and_query(self):
"""
Test full RAG flow: ingest a document and then query it.
"""
import asyncio
litellm._turn_on_debug()
filename, unique_id = self.get_unique_filename("ingest_query")
text_content = f"""
Test document {unique_id} for RAG ingestion and query.
LiteLLM provides a unified interface for 100+ LLMs.
This content should be retrievable via semantic search.
""".encode("utf-8")
file_data = (filename, text_content, "text/plain")
ingest_options = self.get_base_ingest_options()
ingest_options["name"] = f"test-ingest-query-{unique_id}"
try:
# Step 1: Ingest
ingest_response = await litellm.rag.aingest(
ingest_options=ingest_options,
file_data=file_data,
)
print(f"Ingest Response: {ingest_response}")
assert ingest_response["status"] == "completed"
vector_store_id = ingest_response["vector_store_id"]
assert vector_store_id
# Step 2: Query with retry (indexing may take time)
search_results = None
max_retries = 10
for attempt in range(max_retries):
await asyncio.sleep(3)
search_results = await self.query_vector_store(
vector_store_id=vector_store_id,
query=f"Test document {unique_id}",
)
if search_results:
break
print(
f"Attempt {attempt + 1}/{max_retries}: "
"Waiting for document to be indexed..."
)
print(f"Search Results: {search_results}")
# Validate search results
assert search_results is not None, "Document not found after retries"
print("Query successful!")
except litellm.InternalServerError:
pytest.skip("Skipping test due to litellm.InternalServerError")
tests/vector_store_tests/rag/test_document.txt
+4
View File
@@ -0,0 +1,4 @@
Test document abc123 for RAG ingestion.
This is a sample document to test the RAG ingest API.
LiteLLM provides a unified interface for vector stores.
tests/vector_store_tests/rag/test_rag_bedrock.py
+91
View File
@@ -0,0 +1,91 @@
"""
Bedrock Knowledge Base RAG ingestion tests.
Requires environment variables:
- AWS_ACCESS_KEY_ID
- AWS_SECRET_ACCESS_KEY
- AWS_REGION_NAME (optional, defaults to us-west-2)
Optional (for using existing KB instead of auto-creating):
- BEDROCK_KNOWLEDGE_BASE_ID
"""
import os
import sys
from typing import Any, Dict, Optional
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
import litellm
from litellm.types.rag import RAGIngestOptions, BedrockVectorStoreOptions
from tests.vector_store_tests.rag.base_rag_tests import BaseRAGTest
class TestRAGBedrock(BaseRAGTest):
"""Test RAG Ingest with Bedrock Knowledge Base."""
@pytest.fixture(autouse=True)
def check_env_vars(self):
"""Check required environment variables before each test."""
aws_key = os.environ.get("AWS_ACCESS_KEY_ID")
aws_secret = os.environ.get("AWS_SECRET_ACCESS_KEY")
if not aws_key or not aws_secret:
pytest.skip("Skipping Bedrock test: AWS credentials required")
def get_base_ingest_options(self) -> RAGIngestOptions:
"""
Return Bedrock-specific ingest options.
Uses unified interface - no vector_store_id means auto-create KB.
If BEDROCK_KNOWLEDGE_BASE_ID is set, uses existing KB.
"""
# Use existing KB if provided, otherwise auto-create
existing_kb_id = os.environ.get("BEDROCK_KNOWLEDGE_BASE_ID")
return {
"vector_store": BedrockVectorStoreOptions(
custom_llm_provider="bedrock",
vector_store_id=existing_kb_id, # None = auto-create
# wait_for_ingestion defaults to False - returns immediately
),
}
async def query_vector_store(
self,
vector_store_id: str,
query: str,
) -> Optional[Dict[str, Any]]:
"""Query Bedrock Knowledge Base."""
try:
import boto3
except ImportError:
pytest.skip("boto3 required for Bedrock tests")
session = boto3.Session(
aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
region_name=os.environ.get("AWS_REGION_NAME", "us-west-2"),
)
bedrock_agent_runtime = session.client("bedrock-agent-runtime")
response = bedrock_agent_runtime.retrieve(
knowledgeBaseId=vector_store_id,
retrievalQuery={"text": query},
retrievalConfiguration={
"vectorSearchConfiguration": {"numberOfResults": 5}
},
)
if response.get("retrievalResults") and len(response["retrievalResults"]) > 0:
# Check if query terms appear in results
for result in response["retrievalResults"]:
# Extract unique_id from query if present
if query in result["content"]["text"]:
return response
# Return results even if exact match not found
return response
return None
tests/vector_store_tests/rag/test_rag_openai.py
+45
View File
@@ -0,0 +1,45 @@
"""
OpenAI RAG ingestion tests.
"""
import os
import sys
from typing import Any, Dict, Optional
import pytest
sys.path.insert(0, os.path.abspath("../../.."))
import litellm
from litellm.types.rag import RAGIngestOptions, OpenAIVectorStoreOptions
from tests.vector_store_tests.rag.base_rag_tests import BaseRAGTest
class TestRAGOpenAI(BaseRAGTest):
"""Test RAG Ingest with OpenAI provider."""
def get_base_ingest_options(self) -> RAGIngestOptions:
"""Return OpenAI-specific ingest options."""
return {
"vector_store": OpenAIVectorStoreOptions(
custom_llm_provider="openai",
),
}
async def query_vector_store(
self,
vector_store_id: str,
query: str,
) -> Optional[Dict[str, Any]]:
"""Query OpenAI vector store."""
search_response = await litellm.vector_stores.asearch(
vector_store_id=vector_store_id,
query=query,
custom_llm_provider="openai",
)
if search_response.get("data") and len(search_response["data"]) > 0:
return search_response
return None
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx
+47 -56
View File
@@ -1,26 +1,9 @@
import * as useAuthorizedModule from "@/app/(dashboard)/hooks/useAuthorized";
import * as useTeamsModule from "@/app/(dashboard)/hooks/useTeams";
import { render, screen, waitFor } from "@testing-library/react";
import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
import { beforeEach, describe, expect, it, vi } from "vitest";
import AllModelsTab from "./AllModelsTab";
// Mock window.matchMedia for Ant Design components
beforeAll(() => {
Object.defineProperty(window, "matchMedia", {
writable: true,
value: (query: string) => ({
matches: false,
media: query,
onchange: null,
addListener: () => {},
removeListener: () => {},
addEventListener: () => {},
removeEventListener: () => {},
dispatchEvent: () => false,
}),
});
});
describe("AllModelsTab", () => {
const mockSetSelectedModelGroup = vi.fn();
const mockSetSelectedModelId = vi.fn();
@@ -51,24 +34,18 @@ describe("AllModelsTab", () => {
showSSOBanner: false,
};
beforeAll(() => {
// Mock useAuthorized hook
beforeEach(() => {
vi.clearAllMocks();
vi.spyOn(useAuthorizedModule, "default").mockReturnValue(mockUseAuthorized);
});
beforeEach(() => {
vi.clearAllMocks();
});
it("should render with empty data", () => {
// Mock useTeams hook
vi.spyOn(useTeamsModule, "default").mockReturnValue({
teams: [],
setTeams: vi.fn(),
});
const { container } = render(<AllModelsTab {...defaultProps} />);
expect(container).toBeTruthy();
render(<AllModelsTab {...defaultProps} />);
expect(screen.getByText("Current Team:")).toBeInTheDocument();
});
@@ -89,7 +66,6 @@ describe("AllModelsTab", () => {
},
];
// Mock useTeams hook with team data
vi.spyOn(useTeamsModule, "default").mockReturnValue({
teams: mockTeams,
setTeams: vi.fn(),
@@ -101,7 +77,7 @@ describe("AllModelsTab", () => {
model_name: "gpt-4-accessible",
model_info: {
id: "model-1",
access_via_team_ids: ["team-456"], // Direct team access
access_via_team_ids: ["team-456"],
access_groups: [],
},
},
@@ -109,7 +85,7 @@ describe("AllModelsTab", () => {
model_name: "gpt-3.5-turbo-blocked",
model_info: {
id: "model-2",
access_via_team_ids: ["team-789"], // Different team
access_via_team_ids: ["team-789"],
access_groups: [],
},
},
@@ -118,7 +94,6 @@ describe("AllModelsTab", () => {
render(<AllModelsTab {...defaultProps} modelData={modelData} />);
// Initially on "personal" team, should show 0 results (no models have direct_access)
await waitFor(() => {
expect(screen.getByText("Showing 0 results")).toBeInTheDocument();
});
@@ -129,7 +104,7 @@ describe("AllModelsTab", () => {
{
team_id: "team-sales",
team_alias: "Sales Team",
models: ["sales-model-group"], // Team has this model group
models: ["sales-model-group"],
max_budget: null,
budget_duration: null,
tpm_limit: null,
@@ -141,7 +116,6 @@ describe("AllModelsTab", () => {
},
];
// Mock useTeams hook
vi.spyOn(useTeamsModule, "default").mockReturnValue({
teams: mockTeams,
setTeams: vi.fn(),
@@ -153,8 +127,8 @@ describe("AllModelsTab", () => {
model_name: "gpt-4-sales",
model_info: {
id: "model-sales-1",
access_via_team_ids: [], // No direct team access
access_groups: ["sales-model-group"], // But has access group that matches team's models
access_via_team_ids: [],
access_groups: ["sales-model-group"],
},
},
{
@@ -162,7 +136,7 @@ describe("AllModelsTab", () => {
model_info: {
id: "model-eng-1",
access_via_team_ids: [],
access_groups: ["engineering-model-group"], // Different access group
access_groups: ["engineering-model-group"],
},
},
],
@@ -170,14 +144,12 @@ describe("AllModelsTab", () => {
render(<AllModelsTab {...defaultProps} modelData={modelData} />);
// Initially on "personal" team, should show 0 results
await waitFor(() => {
expect(screen.getByText("Showing 0 results")).toBeInTheDocument();
});
});
it("should filter models by direct_access for personal team", async () => {
// Mock useTeams hook
vi.spyOn(useTeamsModule, "default").mockReturnValue({
teams: [],
setTeams: vi.fn(),
@@ -189,7 +161,7 @@ describe("AllModelsTab", () => {
model_name: "gpt-4-personal",
model_info: {
id: "model-personal-1",
direct_access: true, // Available for personal use
direct_access: true,
access_via_team_ids: [],
access_groups: [],
},
@@ -198,7 +170,7 @@ describe("AllModelsTab", () => {
model_name: "gpt-4-team-only",
model_info: {
id: "model-team-1",
direct_access: false, // Not available for personal use
direct_access: false,
access_via_team_ids: ["team-123"],
access_groups: [],
},
@@ -208,16 +180,12 @@ describe("AllModelsTab", () => {
render(<AllModelsTab {...defaultProps} modelData={modelData} />);
// When currentTeam is "personal" (default), it should filter by direct_access === true
// This tests the personal access logic in lines 72-73
// Should show 1 result (only gpt-4-personal with direct_access=true)
await waitFor(() => {
expect(screen.getByText("Showing 1 - 1 of 1 results")).toBeInTheDocument();
});
});
it("should show disabled delete icon for config models", async () => {
// Mock useTeams hook
it("should show config model status for models defined in configs", async () => {
vi.spyOn(useTeamsModule, "default").mockReturnValue({
teams: [],
setTeams: vi.fn(),
@@ -231,7 +199,7 @@ describe("AllModelsTab", () => {
provider: "openai",
model_info: {
id: "model-config-1",
db_model: false, // Config model (no db_model)
db_model: false,
direct_access: true,
access_via_team_ids: [],
access_groups: [],
@@ -246,7 +214,7 @@ describe("AllModelsTab", () => {
provider: "openai",
model_info: {
id: "model-db-1",
db_model: true, // DB model
db_model: true,
direct_access: true,
access_via_team_ids: [],
access_groups: [],
@@ -258,19 +226,42 @@ describe("AllModelsTab", () => {
],
};
const { container } = render(<AllModelsTab {...defaultProps} modelData={modelData} />);
render(<AllModelsTab {...defaultProps} modelData={modelData} />);
await waitFor(() => {
expect(screen.getByText(/Showing \d+ - \d+ of 2 results/)).toBeInTheDocument();
expect(screen.getByText("Config Model")).toBeInTheDocument();
expect(screen.getByText("DB Model")).toBeInTheDocument();
});
});
it("should show 'Defined in config' for models defined in configs", async () => {
vi.spyOn(useTeamsModule, "default").mockReturnValue({
teams: [],
setTeams: vi.fn(),
});
const disabledIcons = container.querySelectorAll(".opacity-50.cursor-not-allowed");
expect(disabledIcons.length).toBeGreaterThan(0);
const modelData = {
data: [
{
model_name: "gpt-4-config-model",
litellm_model_name: "gpt-4-config-model",
provider: "openai",
model_info: {
id: "model-config-defined",
db_model: false,
direct_access: true,
access_via_team_ids: [],
access_groups: [],
created_by: "user-123",
created_at: "2024-01-01",
updated_at: "2024-01-01",
},
},
],
};
const configModelIcon = Array.from(disabledIcons).find((icon) => {
const parent = icon.closest('[class*="actions"], [class*="flex items-center justify-end"]');
return parent !== null;
});
expect(configModelIcon).toBeTruthy();
render(<AllModelsTab {...defaultProps} modelData={modelData} />);
expect(screen.getByText("Defined in config")).toBeInTheDocument();
});
});
ui/litellm-dashboard/src/components/OldTeams.test.tsx
+82 -21
View File
@@ -1,5 +1,6 @@
import { act, fireEvent, render, screen } from "@testing-library/react";
import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key";
import { teamCreateCall } from "./networking";
import OldTeams from "./OldTeams";
@@ -23,6 +24,28 @@ vi.mock("./molecules/notifications_manager", () => ({
},
}));
vi.mock("./key_team_helpers/fetch_available_models_team_key", () => ({
fetchAvailableModelsForTeamOrKey: vi.fn(),
getModelDisplayName: vi.fn((model: string) => model),
unfurlWildcardModelsInList: vi.fn((teamModels: string[], allModels: string[]) => {
const wildcardDisplayNames: string[] = [];
const expandedModels: string[] = [];
teamModels.forEach((teamModel) => {
if (teamModel.endsWith("/*")) {
const provider = teamModel.replace("/*", "");
const matchingModels = allModels.filter((model) => model.startsWith(provider + "/"));
expandedModels.push(...matchingModels);
wildcardDisplayNames.push(teamModel);
} else {
expandedModels.push(teamModel);
}
});
return [...wildcardDisplayNames, ...expandedModels].filter((item, index, array) => array.indexOf(item) === index);
}),
}));
describe("OldTeams - handleCreate organization handling", () => {
beforeEach(() => {
vi.clearAllMocks();
@@ -236,7 +259,7 @@ describe("OldTeams - handleCreate organization handling", () => {
});
it("should clear the delete modal when the cancel button is clicked", async () => {
const { getByRole, getByTestId } = render(
render(
<OldTeams
teams={[
{
@@ -261,7 +284,7 @@ describe("OldTeams - handleCreate organization handling", () => {
organizations={[]}
/>,
);
const deleteTeamButton = getByTestId("delete-team-button");
const deleteTeamButton = screen.getByTestId("delete-team-button");
act(() => {
fireEvent.click(deleteTeamButton);
});
@@ -275,7 +298,7 @@ describe("OldTeams - empty state", () => {
});
it("should display empty state message when teams array is empty", () => {
const { getByText } = render(
render(
<OldTeams
teams={[]}
searchParams={{}}
@@ -287,12 +310,12 @@ describe("OldTeams - empty state", () => {
/>,
);
expect(getByText("No teams found")).toBeInTheDocument();
expect(getByText("Adjust your filters or create a new team")).toBeInTheDocument();
expect(screen.getByText("No teams found")).toBeInTheDocument();
expect(screen.getByText("Adjust your filters or create a new team")).toBeInTheDocument();
});
it("should display empty state message when teams is null", () => {
const { getByText } = render(
render(
<OldTeams
teams={null}
searchParams={{}}
@@ -304,12 +327,12 @@ describe("OldTeams - empty state", () => {
/>,
);
expect(getByText("No teams found")).toBeInTheDocument();
expect(getByText("Adjust your filters or create a new team")).toBeInTheDocument();
expect(screen.getByText("No teams found")).toBeInTheDocument();
expect(screen.getByText("Adjust your filters or create a new team")).toBeInTheDocument();
});
it("should not display empty state when teams array has items", () => {
const { queryByText, getByText } = render(
render(
<OldTeams
teams={[
{
@@ -335,9 +358,9 @@ describe("OldTeams - empty state", () => {
/>,
);
expect(queryByText("No teams found")).not.toBeInTheDocument();
expect(queryByText("Adjust your filters or create a new team")).not.toBeInTheDocument();
expect(getByText("Test Team")).toBeInTheDocument();
expect(screen.queryByText("No teams found")).not.toBeInTheDocument();
expect(screen.queryByText("Adjust your filters or create a new team")).not.toBeInTheDocument();
expect(screen.getByText("Test Team")).toBeInTheDocument();
});
});
@@ -473,7 +496,7 @@ describe("OldTeams - Default Team Settings tab visibility", () => {
});
it("should show Default Team Settings tab for Admin role", () => {
const { getByRole } = render(
render(
<OldTeams
teams={[
{
@@ -499,11 +522,11 @@ describe("OldTeams - Default Team Settings tab visibility", () => {
/>,
);
expect(getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument();
expect(screen.getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument();
});
it("should show Default Team Settings tab for proxy_admin role", () => {
const { getByRole } = render(
render(
<OldTeams
teams={[
{
@@ -529,11 +552,11 @@ describe("OldTeams - Default Team Settings tab visibility", () => {
/>,
);
expect(getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument();
expect(screen.getByRole("tab", { name: "Default Team Settings" })).toBeInTheDocument();
});
it("should not show Default Team Settings tab for proxy_admin_viewer role", () => {
const { queryByRole } = render(
render(
<OldTeams
teams={[
{
@@ -559,11 +582,11 @@ describe("OldTeams - Default Team Settings tab visibility", () => {
/>,
);
expect(queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument();
expect(screen.queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument();
});
it("should not show Default Team Settings tab for Admin Viewer role", () => {
const { queryByRole } = render(
render(
<OldTeams
teams={[
{
@@ -589,6 +612,44 @@ describe("OldTeams - Default Team Settings tab visibility", () => {
/>,
);
expect(queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument();
expect(screen.queryByRole("tab", { name: "Default Team Settings" })).not.toBeInTheDocument();
});
});
describe("OldTeams - all-proxy-models dropdown visibility", () => {
beforeEach(() => {
vi.clearAllMocks();
vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]);
});
it("should not show all-proxy-models option when user has no access to it", async () => {
vi.mocked(fetchAvailableModelsForTeamOrKey).mockResolvedValue(["gpt-4", "gpt-3.5-turbo"]);
render(
<OldTeams
teams={[]}
searchParams={{}}
accessToken="test-token"
setTeams={vi.fn()}
userID="user-123"
userRole="Admin"
organizations={[]}
/>,
);
await waitFor(() => {
expect(fetchAvailableModelsForTeamOrKey).toHaveBeenCalled();
});
const createButton = screen.getByRole("button", { name: /create new team/i });
act(() => {
fireEvent.click(createButton);
});
await waitFor(() => {
expect(screen.getByLabelText(/models/i)).toBeInTheDocument();
});
const allProxyModelsOption = screen.queryByText("All Proxy Models");
expect(allProxyModelsOption).not.toBeInTheDocument();
});
});
ui/litellm-dashboard/src/components/OldTeams.tsx
+11 -3
View File
@@ -1139,12 +1139,20 @@ const Teams: React.FC<TeamProps> = ({
</Tooltip>
</span>
}
rules={[
{
required: true,
message: "Please select at least one model",
},
]}
name="models"
>
<Select2 mode="multiple" placeholder="Select models" style={{ width: "100%" }}>
<Select2.Option key="all-proxy-models" value="all-proxy-models">
All Proxy Models
</Select2.Option>
{(isProxyAdminRole(userRole || "") || userModels.includes("all-proxy-models")) && (
<Select2.Option key="all-proxy-models" value="all-proxy-models">
All Proxy Models
</Select2.Option>
)}
<Select2.Option key="no-default-models" value="no-default-models">
No Default Models
</Select2.Option>
ui/litellm-dashboard/src/components/molecules/models/columns.tsx
+14 -7
View File
@@ -1,9 +1,9 @@
import { KeyIcon, TrashIcon } from "@heroicons/react/outline";
import { ColumnDef } from "@tanstack/react-table";
import { Button, Badge, Icon } from "@tremor/react";
import { Badge, Button, Icon } from "@tremor/react";
import { Tooltip } from "antd";
import { getProviderLogoAndName } from "../../provider_info_helpers";
import { ModelData } from "../../model_dashboard/types";
import { TrashIcon, KeyIcon } from "@heroicons/react/outline";
import { getProviderLogoAndName } from "../../provider_info_helpers";
export const columns = (
userRole: string,
@@ -135,18 +135,25 @@ export const columns = (
size: 160, // Fixed column width
cell: ({ row }) => {
const model = row.original;
const isConfigModel = !model.model_info?.db_model;
const createdBy = model.model_info.created_by;
const createdAt = model.model_info.created_at ? new Date(model.model_info.created_at).toLocaleDateString() : null;
return (
<div className="flex flex-col min-w-0 max-w-[160px]">
{/* Created By - Primary */}
<div className="text-xs font-medium text-gray-900 truncate" title={createdBy || "Unknown"}>
{createdBy || "Unknown"}
<div
className="text-xs font-medium text-gray-900 truncate"
title={isConfigModel ? "Defined in config" : createdBy || "Unknown"}
>
{isConfigModel ? "Defined in config" : createdBy || "Unknown"}
</div>
{/* Created At - Secondary */}
<div className="text-xs text-gray-500 truncate mt-0.5" title={createdAt || "Unknown date"}>
{createdAt || "Unknown date"}
<div
className="text-xs text-gray-500 truncate mt-0.5"
title={isConfigModel ? "Config file" : createdAt || "Unknown date"}
>
{isConfigModel ? "-" : createdAt || "Unknown date"}
</div>
</div>
);
ui/litellm-dashboard/src/components/tag_management/TagTable.test.tsx
+101
View File
@@ -0,0 +1,101 @@
import { render, screen } from "@testing-library/react";
import { beforeEach, describe, expect, it, vi } from "vitest";
import TagTable from "./TagTable";
import { Tag } from "./types";
describe("TagTable", () => {
const mockOnEdit = vi.fn();
const mockOnDelete = vi.fn();
const mockOnSelectTag = vi.fn();
const mockTag: Tag = {
name: "test-tag",
description: "Test description",
models: ["model-1", "model-2"],
model_info: {
"model-1": "GPT-4",
"model-2": "Claude-3",
},
created_at: "2024-01-01T00:00:00Z",
updated_at: "2024-01-01T00:00:00Z",
};
const mockDynamicSpendTag: Tag = {
name: "dynamic-spend-tag",
description:
"This is just a spend tag that was passed dynamically in a request. It does not control any LLM models.",
models: [],
created_at: "2024-01-01T00:00:00Z",
updated_at: "2024-01-01T00:00:00Z",
};
const defaultProps = {
data: [],
onEdit: mockOnEdit,
onDelete: mockOnDelete,
onSelectTag: mockOnSelectTag,
};
beforeEach(() => {
vi.clearAllMocks();
});
it("should render", () => {
render(<TagTable {...defaultProps} />);
expect(screen.getByText("Tag Name")).toBeInTheDocument();
expect(screen.getByText("Description")).toBeInTheDocument();
expect(screen.getByText("Allowed Models")).toBeInTheDocument();
expect(screen.getByText("Created")).toBeInTheDocument();
expect(screen.getByText("Actions")).toBeInTheDocument();
});
it("should display no tags found message when data is empty", () => {
render(<TagTable {...defaultProps} />);
expect(screen.getByText("No tags found")).toBeInTheDocument();
});
it("should display tag name", () => {
render(<TagTable {...defaultProps} data={[mockTag]} />);
expect(screen.getByText("test-tag")).toBeInTheDocument();
});
it("should display tag description", () => {
render(<TagTable {...defaultProps} data={[mockTag]} />);
expect(screen.getByText("Test description")).toBeInTheDocument();
});
it("should display All Models badge when models array is empty", () => {
const tagWithNoModels: Tag = {
...mockTag,
models: [],
};
render(<TagTable {...defaultProps} data={[tagWithNoModels]} />);
expect(screen.getByText("All Models")).toBeInTheDocument();
});
it("should display formatted created date", () => {
render(<TagTable {...defaultProps} data={[mockTag]} />);
const formattedDate = new Date(mockTag.created_at).toLocaleDateString();
expect(screen.getByText(formattedDate)).toBeInTheDocument();
});
it("should disable tag name button for dynamic spend tags", () => {
render(<TagTable {...defaultProps} data={[mockDynamicSpendTag]} />);
const tagButton = screen.getByRole("button", { name: "dynamic-spend-tag" });
expect(tagButton).toBeDisabled();
});
it("should disable edit icon for dynamic spend tags", () => {
render(<TagTable {...defaultProps} data={[mockDynamicSpendTag]} />);
const editIcon = screen.getByLabelText("Edit tag (disabled)");
expect(editIcon).toBeInTheDocument();
expect(editIcon).toHaveClass("cursor-not-allowed");
});
it("should disable delete icon for dynamic spend tags", () => {
render(<TagTable {...defaultProps} data={[mockDynamicSpendTag]} />);
const deleteIcon = screen.getByLabelText("Delete tag (disabled)");
expect(deleteIcon).toBeInTheDocument();
expect(deleteIcon).toHaveClass("cursor-not-allowed");
});
});
ui/litellm-dashboard/src/components/tag_management/TagTable.tsx
+66 -20
View File
@@ -1,18 +1,4 @@
import React from "react";
import {
Table,
TableBody,
TableCell,
TableHead,
TableHeaderCell,
TableRow,
Icon,
Button,
Badge,
Text,
} from "@tremor/react";
import { PencilAltIcon, TrashIcon, SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon } from "@heroicons/react/outline";
import { Tooltip } from "antd";
import { ChevronDownIcon, ChevronUpIcon, PencilAltIcon, SwitchVerticalIcon, TrashIcon } from "@heroicons/react/outline";
import {
ColumnDef,
flexRender,
@@ -21,6 +7,20 @@ import {
SortingState,
useReactTable,
} from "@tanstack/react-table";
import {
Badge,
Button,
Icon,
Table,
TableBody,
TableCell,
TableHead,
TableHeaderCell,
TableRow,
Text,
} from "@tremor/react";
import { Tooltip } from "antd";
import React from "react";
import { Tag } from "./types";
interface TagTableProps {
@@ -30,6 +30,9 @@ interface TagTableProps {
onSelectTag: (tagName: string) => void;
}
const DYNAMIC_SPEND_TAG_DESCRIPTION =
"This is just a spend tag that was passed dynamically in a request. It does not control any LLM models.";
const TagTable: React.FC<TagTableProps> = ({ data, onEdit, onDelete, onSelectTag }) => {
const [sorting, setSorting] = React.useState<SortingState>([{ id: "created_at", desc: true }]);
@@ -39,14 +42,20 @@ const TagTable: React.FC<TagTableProps> = ({ data, onEdit, onDelete, onSelectTag
accessorKey: "name",
cell: ({ row }) => {
const tag = row.original;
const isDynamicSpendTag = tag.description === DYNAMIC_SPEND_TAG_DESCRIPTION;
return (
<div className="overflow-hidden">
<Tooltip title={tag.name}>
<Tooltip
title={
isDynamicSpendTag ? "You cannot view the information of a dynamically generated spend tag" : tag.name
}
>
<Button
size="xs"
variant="light"
className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5"
onClick={() => onSelectTag(tag.name)}
disabled={isDynamicSpendTag}
>
{tag.name}
</Button>
@@ -68,7 +77,7 @@ const TagTable: React.FC<TagTableProps> = ({ data, onEdit, onDelete, onSelectTag
},
},
{
header: "Allowed LLMs",
header: "Allowed Models",
accessorKey: "models",
cell: ({ row }) => {
const tag = row.original;
@@ -102,13 +111,50 @@ const TagTable: React.FC<TagTableProps> = ({ data, onEdit, onDelete, onSelectTag
},
{
id: "actions",
header: "",
header: "Actions",
cell: ({ row }) => {
const tag = row.original;
const isDynamicSpendTag = tag.description === DYNAMIC_SPEND_TAG_DESCRIPTION;
return (
<div className="flex space-x-2">
<Icon icon={PencilAltIcon} size="sm" onClick={() => onEdit(tag)} className="cursor-pointer" />
<Icon icon={TrashIcon} size="sm" onClick={() => onDelete(tag.name)} className="cursor-pointer" />
{isDynamicSpendTag ? (
<Tooltip title="Dynamically generated spend tags cannot be edited">
<Icon
icon={PencilAltIcon}
size="sm"
className="opacity-50 cursor-not-allowed"
aria-label="Edit tag (disabled)"
/>
</Tooltip>
) : (
<Tooltip title="Edit tag">
<Icon
icon={PencilAltIcon}
size="sm"
onClick={() => onEdit(tag)}
className="cursor-pointer hover:text-blue-500"
/>
</Tooltip>
)}
{isDynamicSpendTag ? (
<Tooltip title="Dynamically generated spend tags cannot be deleted">
<Icon
icon={TrashIcon}
size="sm"
className="opacity-50 cursor-not-allowed"
aria-label="Delete tag (disabled)"
/>
</Tooltip>
) : (
<Tooltip title="Delete tag">
<Icon
icon={TrashIcon}
size="sm"
onClick={() => onDelete(tag.name)}
className="cursor-pointer hover:text-red-500"
/>
</Tooltip>
)}
</div>
);
},
ui/litellm-dashboard/src/components/tag_management/components/CreateTagModal.test.tsx
+64
View File
@@ -0,0 +1,64 @@
import { render, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import CreateTagModal from "./CreateTagModal";
describe("CreateTagModal", () => {
const mockOnCancel = vi.fn();
const mockOnSubmit = vi.fn();
const mockAvailableModels = [
{
model_name: "GPT-4",
litellm_params: { model: "gpt-4" },
model_info: { id: "model-1" },
},
{
model_name: "Claude-3",
litellm_params: { model: "claude-3" },
model_info: { id: "model-2" },
},
];
const defaultProps = {
visible: true,
onCancel: mockOnCancel,
onSubmit: mockOnSubmit,
availableModels: mockAvailableModels,
};
beforeEach(() => {
vi.clearAllMocks();
});
it("should render the modal", () => {
render(<CreateTagModal {...defaultProps} />);
expect(screen.getByRole("dialog")).toBeInTheDocument();
expect(screen.getByText("Create New Tag")).toBeInTheDocument();
});
it("should submit form with required tag name", async () => {
const user = userEvent.setup();
render(<CreateTagModal {...defaultProps} />);
const tagNameInput = screen.getByLabelText("Tag Name");
await user.type(tagNameInput, "test-tag");
const submitButton = screen.getByRole("button", { name: /Create Tag/i });
await user.click(submitButton);
expect(mockOnSubmit).toHaveBeenCalledWith({
tag_name: "test-tag",
});
});
it("should not submit form when tag name is missing", async () => {
const user = userEvent.setup();
render(<CreateTagModal {...defaultProps} />);
const submitButton = screen.getByRole("button", { name: /Create Tag/i });
await user.click(submitButton);
// Form validation should prevent submission
expect(mockOnSubmit).not.toHaveBeenCalled();
});
});
ui/litellm-dashboard/src/components/tag_management/components/CreateTagModal.tsx
+11 -33
View File
@@ -1,9 +1,9 @@
import React from "react";
import { Button, TextInput, Accordion, AccordionHeader, AccordionBody, Title } from "@tremor/react";
import { Modal, Form, Select as Select2, Tooltip, Input } from "antd";
import { InfoCircleOutlined } from "@ant-design/icons";
import NumericalInput from "../../shared/numerical_input";
import { Accordion, AccordionBody, AccordionHeader, Button, TextInput, Title } from "@tremor/react";
import { Form, Input, Modal, Select as Select2, Tooltip } from "antd";
import React from "react";
import BudgetDurationDropdown from "../../common_components/budget_duration_dropdown";
import NumericalInput from "../../shared/numerical_input";
interface ModelInfo {
model_name: string;
@@ -22,12 +22,7 @@ interface CreateTagModalProps {
availableModels: ModelInfo[];
}
const CreateTagModal: React.FC<CreateTagModalProps> = ({
visible,
onCancel,
onSubmit,
availableModels,
}) => {
const CreateTagModal: React.FC<CreateTagModalProps> = ({ visible, onCancel, onSubmit, availableModels }) => {
const [form] = Form.useForm();
const handleFinish = (values: any) => {
@@ -41,25 +36,9 @@ const CreateTagModal: React.FC<CreateTagModalProps> = ({
};
return (
<Modal
title="Create New Tag"
visible={visible}
width={800}
footer={null}
onCancel={handleCancel}
>
<Form
form={form}
onFinish={handleFinish}
labelCol={{ span: 8 }}
wrapperCol={{ span: 16 }}
labelAlign="left"
>
<Form.Item
label="Tag Name"
name="tag_name"
rules={[{ required: true, message: "Please input a tag name" }]}
>
<Modal title="Create New Tag" visible={visible} width={800} footer={null} onCancel={handleCancel}>
<Form form={form} onFinish={handleFinish} labelCol={{ span: 8 }} wrapperCol={{ span: 16 }} labelAlign="left">
<Form.Item label="Tag Name" name="tag_name" rules={[{ required: true, message: "Please input a tag name" }]}>
<TextInput />
</Form.Item>
@@ -70,15 +49,15 @@ const CreateTagModal: React.FC<CreateTagModalProps> = ({
<Form.Item
label={
<span>
Allowed Models{" "}
<Tooltip title="Select which LLMs are allowed to process requests from this tag">
Allowed Models
<Tooltip title="Select which models are allowed to process requests from this tag">
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</Tooltip>
</span>
}
name="allowed_llms"
>
<Select2 mode="multiple" placeholder="Select LLMs">
<Select2 mode="multiple" placeholder="Select Models">
{availableModels.map((model) => (
<Select2.Option key={model.model_info.id} value={model.model_info.id}>
<div>
@@ -150,4 +129,3 @@ const CreateTagModal: React.FC<CreateTagModalProps> = ({
};
export default CreateTagModal;
ui/litellm-dashboard/src/components/tag_management/tag_info.tsx
+37 -24
View File
@@ -1,5 +1,15 @@
import React, { useState, useEffect } from "react";
import { Card, Text, Title, Button, Badge, Accordion, AccordionHeader, AccordionBody, Title as TremorTitle } from "@tremor/react";
import {
Card,
Text,
Title,
Button,
Badge,
Accordion,
AccordionHeader,
AccordionBody,
Title as TremorTitle,
} from "@tremor/react";
import { Form, Input, Select as Select2, Tooltip } from "antd";
import { InfoCircleOutlined } from "@ant-design/icons";
import { fetchUserModels } from "../organisms/create_key_button";
@@ -131,7 +141,7 @@ const TagInfoView: React.FC<TagInfoViewProps> = ({ tagId, onClose, accessToken,
<Card>
<Form form={form} onFinish={handleSave} layout="vertical" initialValues={tagDetails}>
<Form.Item label="Tag Name" name="name" rules={[{ required: true, message: "Please input a tag name" }]}>
<Input />
<Input className="rounded-md border-gray-300" />
</Form.Item>
<Form.Item label="Description" name="description">
@@ -141,15 +151,15 @@ const TagInfoView: React.FC<TagInfoViewProps> = ({ tagId, onClose, accessToken,
<Form.Item
label={
<span>
Allowed LLMs{" "}
<Tooltip title="Select which LLMs are allowed to process this type of data">
Allowed Models
<Tooltip title="Select which models are allowed to process this type of data">
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</Tooltip>
</span>
}
name="models"
>
<Select2 mode="multiple" placeholder="Select LLMs">
<Select2 mode="multiple" placeholder="Select Models">
{userModels.map((modelId) => (
<Select2.Option key={modelId} value={modelId}>
{getModelDisplayName(modelId)}
@@ -228,7 +238,7 @@ const TagInfoView: React.FC<TagInfoViewProps> = ({ tagId, onClose, accessToken,
<Text>{tagDetails.description || "-"}</Text>
</div>
<div>
<Text className="font-medium">Allowed LLMs</Text>
<Text className="font-medium">Allowed Models</Text>
<div className="flex flex-wrap gap-2 mt-2">
{!tagDetails.models || tagDetails.models.length === 0 ? (
<Badge color="red">All Models</Badge>
@@ -256,30 +266,33 @@ const TagInfoView: React.FC<TagInfoViewProps> = ({ tagId, onClose, accessToken,
<Card>
<Title>Budget & Rate Limits</Title>
<div className="space-y-4 mt-4">
{tagDetails.litellm_budget_table.max_budget !== undefined && tagDetails.litellm_budget_table.max_budget !== null && (
<div>
<Text className="font-medium">Max Budget</Text>
<Text>${tagDetails.litellm_budget_table.max_budget}</Text>
</div>
)}
{tagDetails.litellm_budget_table.max_budget !== undefined &&
tagDetails.litellm_budget_table.max_budget !== null && (
<div>
<Text className="font-medium">Max Budget</Text>
<Text>${tagDetails.litellm_budget_table.max_budget}</Text>
</div>
)}
{tagDetails.litellm_budget_table.budget_duration && (
<div>
<Text className="font-medium">Budget Duration</Text>
<Text>{tagDetails.litellm_budget_table.budget_duration}</Text>
</div>
)}
{tagDetails.litellm_budget_table.tpm_limit !== undefined && tagDetails.litellm_budget_table.tpm_limit !== null && (
<div>
<Text className="font-medium">TPM Limit</Text>
<Text>{tagDetails.litellm_budget_table.tpm_limit.toLocaleString()}</Text>
</div>
)}
{tagDetails.litellm_budget_table.rpm_limit !== undefined && tagDetails.litellm_budget_table.rpm_limit !== null && (
<div>
<Text className="font-medium">RPM Limit</Text>
<Text>{tagDetails.litellm_budget_table.rpm_limit.toLocaleString()}</Text>
</div>
)}
{tagDetails.litellm_budget_table.tpm_limit !== undefined &&
tagDetails.litellm_budget_table.tpm_limit !== null && (
<div>
<Text className="font-medium">TPM Limit</Text>
<Text>{tagDetails.litellm_budget_table.tpm_limit.toLocaleString()}</Text>
</div>
)}
{tagDetails.litellm_budget_table.rpm_limit !== undefined &&
tagDetails.litellm_budget_table.rpm_limit !== null && (
<div>
<Text className="font-medium">RPM Limit</Text>
<Text>{tagDetails.litellm_budget_table.rpm_limit.toLocaleString()}</Text>
</div>
)}
</div>
</Card>
)}
ui/litellm-dashboard/src/components/team/team_info.test.tsx
+84 -4
View File
@@ -1,7 +1,7 @@
import * as networking from "@/components/networking";
import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
import { afterEach, describe, expect, it, vi } from "vitest";
import TeamInfoView from "./team_info";
import { render, waitFor } from "@testing-library/react";
import * as networking from "@/components/networking";
// Mock the networking module
vi.mock("@/components/networking", () => ({
@@ -61,7 +61,7 @@ describe("TeamInfoView", () => {
vi.mocked(networking.getGuardrailsList).mockResolvedValue([]);
vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]);
const { getByText } = render(
render(
<TeamInfoView
teamId="123"
onUpdate={() => {}}
@@ -75,7 +75,87 @@ describe("TeamInfoView", () => {
/>,
);
await waitFor(() => {
expect(getByText("User ID")).toBeInTheDocument();
expect(screen.queryByText("User ID")).not.toBeNull();
});
});
it("should not show all-proxy-models option when user has no access to it", async () => {
vi.mocked(networking.teamInfoCall).mockResolvedValue({
team_id: "123",
team_info: {
team_alias: "Test Team",
team_id: "123",
organization_id: null,
admins: ["admin@test.com"],
members: ["user1@test.com", "user2@test.com"],
members_with_roles: [
{
user_id: "user1@test.com",
user_email: "user1@test.com",
role: "member",
spend: 0,
budget_id: "budget1",
},
],
metadata: {},
tpm_limit: null,
rpm_limit: null,
max_budget: null,
budget_duration: null,
models: ["gpt-4"],
blocked: false,
spend: 0,
max_parallel_requests: null,
budget_reset_at: null,
model_id: null,
litellm_model_table: null,
created_at: "2024-01-01T00:00:00Z",
team_member_budget_table: null,
},
keys: [],
team_memberships: [],
});
vi.mocked(networking.getGuardrailsList).mockResolvedValue([]);
vi.mocked(networking.fetchMCPAccessGroups).mockResolvedValue([]);
render(
<TeamInfoView
teamId="123"
onUpdate={() => {}}
onClose={() => {}}
accessToken="123"
is_team_admin={true}
is_proxy_admin={true}
userModels={["gpt-4", "gpt-3.5-turbo"]}
editTeam={false}
premiumUser={false}
/>,
);
await waitFor(() => {
expect(screen.getAllByText("Test Team")).not.toBeNull();
});
const settingsTab = screen.getByRole("tab", { name: "Settings" });
act(() => {
fireEvent.click(settingsTab);
});
await waitFor(() => {
expect(screen.getByText("Team Settings")).toBeInTheDocument();
});
const editButton = screen.getByRole("button", { name: "Edit Settings" });
act(() => {
fireEvent.click(editButton);
});
await waitFor(() => {
expect(screen.getByLabelText("Models")).toBeInTheDocument();
});
const allProxyModelsOption = screen.queryByText("All Proxy Models");
expect(allProxyModelsOption).not.toBeInTheDocument();
});
});
ui/litellm-dashboard/src/components/team/team_info.tsx
+48 -42
View File
@@ -1,50 +1,50 @@
import React, { useState, useEffect } from "react";
import NumericalInput from "../shared/numerical_input";
import UserSearchModal from "@/components/common_components/user_search_modal";
import {
Card,
Title,
Text,
Tab,
TabList,
TabGroup,
TabPanel,
TabPanels,
Grid,
Badge,
Button as TremorButton,
TextInput,
} from "@tremor/react";
import TeamMembersComponent from "./team_member_view";
import MemberPermissions from "./member_permissions";
import {
teamInfoCall,
teamMemberDeleteCall,
teamMemberAddCall,
teamMemberUpdateCall,
Member,
teamUpdateCall,
getGuardrailsList,
Member,
teamInfoCall,
teamMemberAddCall,
teamMemberDeleteCall,
teamMemberUpdateCall,
teamUpdateCall,
} from "@/components/networking";
import { Button, Form, Input, Select, Switch, message, Tooltip } from "antd";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils";
import { InfoCircleOutlined } from "@ant-design/icons";
import { ArrowLeftIcon } from "@heroicons/react/outline";
import MemberModal from "./edit_membership";
import UserSearchModal from "@/components/common_components/user_search_modal";
import {
Badge,
Card,
Grid,
Tab,
TabGroup,
TabList,
TabPanel,
TabPanels,
Text,
TextInput,
Title,
Button as TremorButton,
} from "@tremor/react";
import { Button, Form, Input, message, Select, Switch, Tooltip } from "antd";
import { CheckIcon, CopyIcon } from "lucide-react";
import React, { useEffect, useState } from "react";
import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
import DeleteResourceModal from "../common_components/DeleteResourceModal";
import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector";
import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key";
import ObjectPermissionsView from "../object_permissions_view";
import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
import LoggingSettingsView from "../logging_settings_view";
import MCPServerSelector from "../mcp_server_management/MCPServerSelector";
import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import EditLoggingSettings from "./EditLoggingSettings";
import LoggingSettingsView from "../logging_settings_view";
import { fetchMCPAccessGroups } from "../networking";
import { CheckIcon, CopyIcon } from "lucide-react";
import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
import NotificationsManager from "../molecules/notifications_manager";
import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector";
import { mapEmptyStringToNull } from "@/utils/keyUpdateUtils";
import DeleteResourceModal from "../common_components/DeleteResourceModal";
import { fetchMCPAccessGroups } from "../networking";
import ObjectPermissionsView from "../object_permissions_view";
import NumericalInput from "../shared/numerical_input";
import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
import MemberModal from "./edit_membership";
import EditLoggingSettings from "./EditLoggingSettings";
import MemberPermissions from "./member_permissions";
import TeamMembersComponent from "./team_member_view";
export interface TeamMembership {
user_id: string;
@@ -586,11 +586,17 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<Input type="" />
</Form.Item>
<Form.Item label="Models" name="models">
<Form.Item
label="Models"
name="models"
rules={[{ required: true, message: "Please select at least one model" }]}
>
<Select mode="multiple" placeholder="Select models">
<Select.Option key="all-proxy-models" value="all-proxy-models">
All Proxy Models
</Select.Option>
{(is_proxy_admin || userModels.includes("all-proxy-models")) && (
<Select.Option key="all-proxy-models" value="all-proxy-models">
All Proxy Models
</Select.Option>
)}
<Select.Option key="no-default-models" value="no-default-models">
No Default Models
</Select.Option>
ui/litellm-dashboard/src/components/view_users/columns.tsx
+12 -2
View File
@@ -22,10 +22,12 @@ export const columns = (
handleUserClick: (userId: string, openInEditMode?: boolean) => void,
selectionOptions?: SelectionOptions,
): ColumnDef<UserInfo>[] => {
// Backend sortable columns: user_id, user_email, created_at, spend, user_alias, user_role
const baseColumns: ColumnDef<UserInfo>[] = [
{
header: "User ID",
accessorKey: "user_id",
enableSorting: true,
cell: ({ row }) => (
<Tooltip title={row.original.user_id}>
<span className="text-xs">{row.original.user_id ? `${row.original.user_id.slice(0, 7)}...` : "-"}</span>
@@ -35,16 +37,19 @@ export const columns = (
{
header: "Email",
accessorKey: "user_email",
enableSorting: true,
cell: ({ row }) => <span className="text-xs">{row.original.user_email || "-"}</span>,
},
{
header: "Global Proxy Role",
accessorKey: "user_role",
enableSorting: true,
cell: ({ row }) => <span className="text-xs">{possibleUIRoles?.[row.original.user_role]?.ui_label || "-"}</span>,
},
{
header: "Spend (USD)",
accessorKey: "spend",
enableSorting: true,
cell: ({ row }) => (
<span className="text-xs">{row.original.spend ? formatNumberWithCommas(row.original.spend, 4) : "-"}</span>
),
@@ -52,6 +57,7 @@ export const columns = (
{
header: "Budget (USD)",
accessorKey: "max_budget",
enableSorting: false,
cell: ({ row }) => (
<span className="text-xs">{row.original.max_budget !== null ? row.original.max_budget : "Unlimited"}</span>
),
@@ -66,6 +72,7 @@ export const columns = (
</div>
),
accessorKey: "sso_user_id",
enableSorting: false,
cell: ({ row }) => (
<span className="text-xs">{row.original.sso_user_id !== null ? row.original.sso_user_id : "-"}</span>
),
@@ -73,6 +80,7 @@ export const columns = (
{
header: "API Keys",
accessorKey: "key_count",
enableSorting: false,
cell: ({ row }) => (
<Grid numItems={2}>
{row.original.key_count > 0 ? (
@@ -90,7 +98,7 @@ export const columns = (
{
header: "Created At",
accessorKey: "created_at",
sortingFn: "datetime",
enableSorting: true,
cell: ({ row }) => (
<span className="text-xs">
{row.original.created_at ? new Date(row.original.created_at).toLocaleDateString() : "-"}
@@ -100,7 +108,7 @@ export const columns = (
{
header: "Updated At",
accessorKey: "updated_at",
sortingFn: "datetime",
enableSorting: false,
cell: ({ row }) => (
<span className="text-xs">
{row.original.updated_at ? new Date(row.original.updated_at).toLocaleDateString() : "-"}
@@ -110,6 +118,7 @@ export const columns = (
{
id: "actions",
header: "Actions",
enableSorting: false,
cell: ({ row }) => (
<div className="flex gap-2">
<Tooltip title="Edit user details">
@@ -148,6 +157,7 @@ export const columns = (
return [
{
id: "select",
enableSorting: false,
header: () => (
<Checkbox
indeterminate={isIndeterminate}
ui/litellm-dashboard/src/components/view_users/table.test.tsx
+55 -4
View File
@@ -1,6 +1,5 @@
import { render } from "@testing-library/react";
import { act, fireEvent, render, screen } from "@testing-library/react";
import { describe, expect, it, vi } from "vitest";
import React from "react";
import { UserDataTable } from "./table";
@@ -21,7 +20,7 @@ describe("UserDataTable", () => {
const updateFilters = vi.fn();
const { getByText } = render(
render(
<UserDataTable
data={[]}
columns={[]}
@@ -41,6 +40,58 @@ describe("UserDataTable", () => {
/>,
);
expect(getByText("Filters")).toBeInTheDocument();
expect(screen.getByText("Filters")).toBeInTheDocument();
});
it("should call onSortChange when clicking a sortable header", () => {
const filters = {
email: "",
user_id: "",
user_role: "",
sso_user_id: "",
team: "",
model: "",
min_spend: null,
max_spend: null,
sort_by: "created_at",
sort_order: "desc" as const,
};
const updateFilters = vi.fn();
const onSortChange = vi.fn();
const possibleUIRoles = {
admin: { ui_label: "Admin" },
user: { ui_label: "User" },
};
render(
<UserDataTable
data={[]}
columns={[]}
accessToken={null}
userRole={"Admin"}
possibleUIRoles={possibleUIRoles}
filters={filters}
updateFilters={updateFilters}
initialFilters={filters}
teams={[]}
handleEdit={vi.fn()}
handleDelete={vi.fn()}
handleResetPassword={vi.fn()}
userListResponse={{ users: [], total: 0, page: 1, page_size: 25, total_pages: 1 }}
currentPage={1}
handlePageChange={vi.fn()}
onSortChange={onSortChange}
currentSort={{ sortBy: filters.sort_by, sortOrder: filters.sort_order }}
/>,
);
const emailHeader = screen.getByRole("columnheader", { name: /email/i });
act(() => {
fireEvent.click(emailHeader);
});
expect(onSortChange).toHaveBeenCalledWith("user_email", "desc");
});
});
ui/litellm-dashboard/src/components/view_users/table.tsx
+15 -16
View File
@@ -1,11 +1,4 @@
import {
ColumnDef,
flexRender,
getCoreRowModel,
getSortedRowModel,
SortingState,
useReactTable,
} from "@tanstack/react-table";
import { ColumnDef, flexRender, getCoreRowModel, SortingState, useReactTable } from "@tanstack/react-table";
import React from "react";
import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell, Select, SelectItem } from "@tremor/react";
import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon } from "@heroicons/react/outline";
@@ -167,17 +160,23 @@ export function UserDataTable({
state: {
sorting,
},
onSortingChange: (newSorting: any) => {
onSortingChange: (updaterOrValue: any) => {
const newSorting = typeof updaterOrValue === "function" ? updaterOrValue(sorting) : updaterOrValue;
setSorting(newSorting);
if (newSorting.length > 0) {
if (newSorting && Array.isArray(newSorting) && newSorting.length > 0 && newSorting[0]) {
const sortState = newSorting[0];
const sortBy = sortState.id;
const sortOrder = sortState.desc ? "desc" : "asc";
onSortChange?.(sortBy, sortOrder);
if (sortState.id) {
const sortBy = sortState.id;
const sortOrder = sortState.desc ? "desc" : "asc";
onSortChange?.(sortBy, sortOrder);
}
} else {
// Reset to default sort when no sorting is selected
onSortChange?.("created_at", "desc");
}
},
getCoreRowModel: getCoreRowModel(),
getSortedRowModel: getSortedRowModel(),
manualSorting: true,
enableSorting: true,
});
@@ -403,7 +402,7 @@ export function UserDataTable({
header.id === "actions"
? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)]"
: ""
}`}
} ${header.column.getCanSort() ? "cursor-pointer hover:bg-gray-50" : ""}`}
onClick={header.column.getToggleSortingHandler()}
>
<div className="flex items-center justify-between gap-2">
@@ -412,7 +411,7 @@ export function UserDataTable({
? null
: flexRender(header.column.columnDef.header, header.getContext())}
</div>
{header.id !== "actions" && (
{header.id !== "actions" && header.column.getCanSort() && (
<div className="w-4">
{header.column.getIsSorted() ? (
{