Merge remote-tracking branch 'origin' into litellm_ui_model_page_perf

This commit is contained in:
yuneng-jiang 2025-11-26 09:06:43 -08:00
commit ce65663ad1
69 changed files with 4870 additions and 270 deletions

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

View file

@ -20,6 +20,7 @@ LiteLLM integrates with vector stores, allowing your models to access your organ
- [OpenAI Vector Stores](https://platform.openai.com/docs/api-reference/vector-stores/search)
- [Azure Vector Stores](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/file-search?tabs=python#vector-stores) (Cannot be directly queried. Only available for calling in Assistants messages. We will be adding Azure AI Search Vector Store API support soon.)
- [Vertex AI RAG API](https://cloud.google.com/vertex-ai/generative-ai/docs/rag-overview)
- [Gemini File Search](https://ai.google.dev/gemini-api/docs/file-search)
## Quick Start

View file

@ -0,0 +1,414 @@
import Tabs from '@theme/Tabs';
import TabItem from '@theme/TabItem';
# Gemini File Search
Use Google Gemini's File Search for Retrieval Augmented Generation (RAG) with LiteLLM.
Gemini File Search imports, chunks, and indexes your data to enable fast retrieval of relevant information based on user prompts. This information is then provided as context to the model for more accurate and relevant answers.
[Official Gemini File Search Documentation](https://ai.google.dev/gemini-api/docs/file-search)
## Features
| Feature | Supported | Notes |
|---------|-----------|-------|
| Cost Tracking | ❌ | Cost calculation not yet implemented |
| Logging | ✅ | Full request/response logging |
| RAG Ingest API | ✅ | Upload → Chunk → Embed → Store |
| Vector Store Search | ✅ | Search with metadata filters |
| Custom Chunking | ✅ | Configure chunk size and overlap |
| Metadata Filtering | ✅ | Filter by custom metadata |
| Citations | ✅ | Extract from grounding metadata |
## Quick Start
### Setup
Set your Gemini API key:
```bash
export GEMINI_API_KEY="your-api-key"
# or
export GOOGLE_API_KEY="your-api-key"
```
### Basic RAG Ingest
<Tabs>
<TabItem value="python" label="Python SDK">
```python
import litellm
# Ingest a document
response = await litellm.aingest(
ingest_options={
"name": "my-document-store",
"vector_store": {
"custom_llm_provider": "gemini"
}
},
file_data=("document.txt", b"Your document content", "text/plain")
)
print(f"Vector Store ID: {response['vector_store_id']}")
print(f"File ID: {response['file_id']}")
```
</TabItem>
<TabItem value="proxy" label="LiteLLM Proxy">
```bash
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": {
"name": "my-document-store",
"vector_store": {
"custom_llm_provider": "gemini"
}
}
}'
```
</TabItem>
</Tabs>
### Search Vector Store
<Tabs>
<TabItem value="python" label="Python SDK">
```python
import litellm
# Search the vector store
response = await litellm.vector_stores.asearch(
vector_store_id="fileSearchStores/your-store-id",
query="What is the main topic?",
custom_llm_provider="gemini",
max_num_results=5
)
for result in response["data"]:
print(f"Score: {result.get('score')}")
print(f"Content: {result['content'][0]['text']}")
```
</TabItem>
<TabItem value="proxy" label="LiteLLM Proxy">
```bash
curl -X POST "http://localhost:4000/v1/vector_stores/fileSearchStores/your-store-id/search" \
-H "Authorization: Bearer sk-1234" \
-H "Content-Type: application/json" \
-d '{
"query": "What is the main topic?",
"custom_llm_provider": "gemini",
"max_num_results": 5
}'
```
</TabItem>
</Tabs>
## Advanced Features
### Custom Chunking Configuration
Control how documents are split into chunks:
```python
import litellm
response = await litellm.aingest(
ingest_options={
"name": "custom-chunking-store",
"vector_store": {
"custom_llm_provider": "gemini"
},
"chunking_strategy": {
"white_space_config": {
"max_tokens_per_chunk": 200,
"max_overlap_tokens": 20
}
}
},
file_data=("document.txt", document_content, "text/plain")
)
```
**Chunking Parameters:**
- `max_tokens_per_chunk`: Maximum tokens per chunk (default: 800, min: 100, max: 4096)
- `max_overlap_tokens`: Overlap between chunks (default: 400)
### Metadata Filtering
Attach custom metadata to files and filter searches:
#### Attach Metadata During Ingest
```python
import litellm
response = await litellm.aingest(
ingest_options={
"name": "metadata-store",
"vector_store": {
"custom_llm_provider": "gemini",
"custom_metadata": [
{"key": "author", "string_value": "John Doe"},
{"key": "year", "numeric_value": 2024},
{"key": "category", "string_value": "documentation"}
]
}
},
file_data=("document.txt", document_content, "text/plain")
)
```
#### Search with Metadata Filter
```python
import litellm
response = await litellm.vector_stores.asearch(
vector_store_id="fileSearchStores/your-store-id",
query="What is LiteLLM?",
custom_llm_provider="gemini",
filters={"author": "John Doe", "category": "documentation"}
)
```
**Filter Syntax:**
- Simple equality: `{"key": "value"}`
- Gemini converts to: `key="value"`
- Multiple filters combined with AND
### Using Existing Vector Store
Ingest into an existing File Search store:
```python
import litellm
# First, create a store
create_response = await litellm.vector_stores.acreate(
name="My Persistent Store",
custom_llm_provider="gemini"
)
store_id = create_response["id"]
# Then ingest multiple documents into it
for doc in documents:
await litellm.aingest(
ingest_options={
"vector_store": {
"custom_llm_provider": "gemini",
"vector_store_id": store_id # Reuse existing store
}
},
file_data=(doc["name"], doc["content"], doc["type"])
)
```
### Citation Extraction
Gemini provides grounding metadata with citations:
```python
import litellm
response = await litellm.vector_stores.asearch(
vector_store_id="fileSearchStores/your-store-id",
query="Explain the concept",
custom_llm_provider="gemini"
)
for result in response["data"]:
# Access citation information
if "attributes" in result:
print(f"URI: {result['attributes'].get('uri')}")
print(f"Title: {result['attributes'].get('title')}")
# Content with relevance score
print(f"Score: {result.get('score')}")
print(f"Text: {result['content'][0]['text']}")
```
## Complete Example
End-to-end workflow:
```python
import litellm
# 1. Create a File Search store
store_response = await litellm.vector_stores.acreate(
name="Knowledge Base",
custom_llm_provider="gemini"
)
store_id = store_response["id"]
print(f"Created store: {store_id}")
# 2. Ingest documents with custom chunking and metadata
documents = [
{
"name": "intro.txt",
"content": b"Introduction to LiteLLM...",
"metadata": [
{"key": "section", "string_value": "intro"},
{"key": "priority", "numeric_value": 1}
]
},
{
"name": "advanced.txt",
"content": b"Advanced features...",
"metadata": [
{"key": "section", "string_value": "advanced"},
{"key": "priority", "numeric_value": 2}
]
}
]
for doc in documents:
ingest_response = await litellm.aingest(
ingest_options={
"name": f"ingest-{doc['name']}",
"vector_store": {
"custom_llm_provider": "gemini",
"vector_store_id": store_id,
"custom_metadata": doc["metadata"]
},
"chunking_strategy": {
"white_space_config": {
"max_tokens_per_chunk": 300,
"max_overlap_tokens": 50
}
}
},
file_data=(doc["name"], doc["content"], "text/plain")
)
print(f"Ingested: {doc['name']}")
# 3. Search with filters
search_response = await litellm.vector_stores.asearch(
vector_store_id=store_id,
query="How do I get started?",
custom_llm_provider="gemini",
filters={"section": "intro"},
max_num_results=3
)
# 4. Process results
for i, result in enumerate(search_response["data"]):
print(f"\nResult {i+1}:")
print(f" Score: {result.get('score')}")
print(f" File: {result.get('filename')}")
print(f" Content: {result['content'][0]['text'][:100]}...")
```
## Supported File Types
Gemini File Search supports a wide range of file formats:
### Documents
- PDF (`application/pdf`)
- Microsoft Word (`.docx`, `.doc`)
- Microsoft Excel (`.xlsx`, `.xls`)
- Microsoft PowerPoint (`.pptx`)
- OpenDocument formats (`.odt`, `.ods`, `.odp`)
### Text Files
- Plain text (`text/plain`)
- Markdown (`text/markdown`)
- HTML (`text/html`)
- CSV (`text/csv`)
- JSON (`application/json`)
- XML (`application/xml`)
### Code Files
- Python, JavaScript, TypeScript, Java, C/C++, Go, Rust, etc.
- Most common programming languages supported
See [Gemini's full list of supported file types](https://ai.google.dev/gemini-api/docs/file-search#supported-file-types).
## Pricing
- **Indexing**: $0.15 per 1M tokens (embedding pricing)
- **Storage**: Free
- **Query embeddings**: Free
- **Retrieved tokens**: Charged as regular context tokens
## Supported Models
File Search works with:
- `gemini-3-pro-preview`
- `gemini-2.5-pro`
- `gemini-2.5-flash` (and preview versions)
- `gemini-2.5-flash-lite` (and preview versions)
## Troubleshooting
### Authentication Errors
```python
# Ensure API key is set
import os
os.environ["GEMINI_API_KEY"] = "your-api-key"
# Or pass explicitly
response = await litellm.aingest(
ingest_options={
"vector_store": {
"custom_llm_provider": "gemini",
"api_key": "your-api-key"
}
},
file_data=(...)
)
```
### Store Not Found
Ensure you're using the full store name format:
- ✅ `fileSearchStores/abc123`
- ❌ `abc123`
### Large Files
For files >100MB, split them into smaller chunks before ingestion.
### Slow Indexing
After ingestion, Gemini may need time to index documents. Wait a few seconds before searching:
```python
import time
# After ingest
await litellm.aingest(...)
# Wait for indexing
time.sleep(5)
# Then search
await litellm.vector_stores.asearch(...)
```
## Related Resources
- [Gemini File Search Official Docs](https://ai.google.dev/gemini-api/docs/file-search)
- [LiteLLM RAG Ingest API](/docs/rag_ingest)
- [LiteLLM Vector Store Search](/docs/vector_stores/search)
- [Using Vector Stores with Chat](/docs/completion/knowledgebase)

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

View file

@ -0,0 +1,273 @@
# /rag/ingest
All-in-one document ingestion pipeline: **Upload → Chunk → Embed → Vector Store**
| Feature | Supported |
|---------|-----------|
| Cost Tracking | ❌ |
| Logging | ✅ |
| Supported Providers | `openai`, `bedrock`, `gemini` |
## 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\"
}
}
}"
```
### Gemini
```bash showLineNumbers title="Ingest to Gemini File Search"
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\": \"gemini\"
}
}
}"
```
**With Custom Chunking:**
```bash showLineNumbers title="Ingest with custom chunking"
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": "gemini"
},
"chunking_strategy": {
"white_space_config": {
"max_tokens_per_chunk": 200,
"max_overlap_tokens": 20
}
}
}
}'
```
## 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"}}
}'
```

View file

@ -12,7 +12,7 @@ Search a vector store for relevant chunks based on a query and file attributes f
| Cost Tracking | ✅ | Tracked per search operation |
| Logging | ✅ | Works across all integrations |
| End-user Tracking | ✅ | |
| Support LLM Providers | **OpenAI, Azure OpenAI, Bedrock, Vertex RAG Engine, Azure AI, Milvus** | Full vector stores API support across providers |
| Support LLM Providers | **OpenAI, Azure OpenAI, Bedrock, Vertex RAG Engine, Azure AI, Milvus, Gemini** | Full vector stores API support across providers |
## Usage
@ -164,6 +164,41 @@ print(response)
[See full Milvus vector store documentation](../providers/milvus_vector_stores.md)
</TabItem>
<TabItem value="gemini-provider" label="Gemini Provider">
#### Using Gemini File Search
```python showLineNumbers title="Search Vector Store - Gemini Provider"
import litellm
import os
# Set credentials
os.environ["GEMINI_API_KEY"] = "your-gemini-api-key"
response = await litellm.vector_stores.asearch(
vector_store_id="fileSearchStores/your-store-id",
query="What is the capital of France?",
custom_llm_provider="gemini",
max_num_results=5
)
print(response)
```
**With Metadata Filter:**
```python showLineNumbers title="Search with Metadata Filter"
response = await litellm.vector_stores.asearch(
vector_store_id="fileSearchStores/your-store-id",
query="What is LiteLLM?",
custom_llm_provider="gemini",
filters={"author": "John Doe", "category": "documentation"},
max_num_results=5
)
print(response)
```
[See full Gemini File Search documentation](../providers/gemini_file_search.md)
</TabItem>
</Tabs>

View file

@ -412,6 +412,7 @@ const sidebars = {
"proxy/pass_through"
]
},
"rag_ingest",
"realtime",
"rerank",
"response_api",

View file

@ -1434,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 *
@ -1470,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

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:

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))

View file

@ -580,6 +580,20 @@ class AmazonConverseConfig(BaseConfig):
non_default_params=non_default_params, optional_params=optional_params
)
final_is_thinking_enabled = self.is_thinking_enabled(optional_params)
if (
final_is_thinking_enabled
and "tool_choice" in optional_params
):
tool_choice_block = optional_params["tool_choice"]
if isinstance(tool_choice_block, dict):
if "any" in tool_choice_block or "tool" in tool_choice_block:
verbose_logger.info(
f"{model} does not support forced tool use (tool_choice='required' or specific tool) "
f"when reasoning is enabled. Changing tool_choice to 'auto'."
)
optional_params["tool_choice"] = ToolChoiceValuesBlock(auto={})
return optional_params
def _translate_response_format_param(

View file

@ -1,3 +1,5 @@
from __future__ import annotations
import json
from typing import TYPE_CHECKING, Any, Optional, Union

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")

View file

@ -0,0 +1,6 @@
"""Gemini File Search Vector Store module."""
from .transformation import GeminiVectorStoreConfig
__all__ = ["GeminiVectorStoreConfig"]

View file

@ -0,0 +1,357 @@
"""
Gemini File Search Vector Store Transformation Layer.
Implements the transformation between LiteLLM's unified vector store API
and Google Gemini's File Search API.
"""
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
import httpx
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
from litellm.llms.gemini.common_utils import (
GeminiError,
GeminiModelInfo,
get_api_key_from_env,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.vector_stores import (
VECTOR_STORE_OPENAI_PARAMS,
BaseVectorStoreAuthCredentials,
VectorStoreCreateOptionalRequestParams,
VectorStoreCreateResponse,
VectorStoreFileCounts,
VectorStoreIndexEndpoints,
VectorStoreResultContent,
VectorStoreSearchOptionalRequestParams,
VectorStoreSearchResponse,
VectorStoreSearchResult,
)
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
else:
LiteLLMLoggingObj = Any
class GeminiVectorStoreConfig(BaseVectorStoreConfig):
"""
Vector store configuration for Google Gemini File Search.
"""
def __init__(self) -> None:
super().__init__()
self.model_info = GeminiModelInfo()
self._cached_api_key: Optional[str] = None
def get_auth_credentials(
self, litellm_params: dict
) -> BaseVectorStoreAuthCredentials:
"""Gemini uses API key in query params, not headers."""
return {}
def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints:
"""
Gemini File Search endpoints.
Note: Search is done via generateContent with file_search tool,
not a dedicated search endpoint.
"""
return {
"read": [("POST", "/models/{model}:generateContent")],
"write": [("POST", "/fileSearchStores")],
}
def get_supported_openai_params(
self, model: str
) -> List[VECTOR_STORE_OPENAI_PARAMS]:
"""Supported parameters for Gemini File Search."""
return ["max_num_results", "filters"]
def validate_environment(
self, headers: dict, litellm_params: Optional[GenericLiteLLMParams]
) -> dict:
"""Validate and set up headers for Gemini API."""
headers = headers or {}
headers.setdefault("Content-Type", "application/json")
if litellm_params:
api_key = litellm_params.get("api_key") or get_api_key_from_env()
if api_key:
self._cached_api_key = api_key
return headers
def get_complete_url(self, api_base: Optional[str], litellm_params: dict) -> str:
"""
Get the complete base URL for Gemini API.
Note: This returns the base URL WITHOUT the API key.
The API key will be appended to specific endpoint URLs in the transform methods.
"""
if api_base is None:
api_base = GeminiModelInfo.get_api_base()
if api_base is None:
raise ValueError("GEMINI_API_BASE is not set")
# Ensure we're using the v1beta version for File Search
api_version = "v1beta"
return f"{api_base}/{api_version}"
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> GeminiError:
"""Return Gemini-specific error class."""
return GeminiError(
status_code=status_code,
message=error_message,
headers=headers,
)
def transform_search_vector_store_request(
self,
vector_store_id: str,
query: Union[str, List[str]],
vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams,
api_base: str,
litellm_logging_obj: LiteLLMLoggingObj,
litellm_params: dict,
) -> Tuple[str, Dict]:
"""
Transform search request to Gemini's generateContent format.
Gemini File Search works by calling generateContent with a file_search tool.
"""
# Convert query list to single string if needed
if isinstance(query, list):
query = " ".join(query)
# Get model from litellm_params or use default
# Note: File Search requires gemini-2.5-flash or later
model = litellm_params.get("model") or "gemini-2.5-flash"
if model and model.startswith("gemini/"):
model = model.replace("gemini/", "")
# Get API key - Gemini requires it as a query parameter
api_key = litellm_params.get("api_key") or GeminiModelInfo.get_api_key()
if not api_key:
raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY is required")
# Build the URL for generateContent with API key
url = f"{api_base}/models/{model}:generateContent?key={api_key}"
# Build file_search tool configuration (using snake_case as per Gemini docs)
file_search_config: Dict[str, Any] = {
"file_search_store_names": [vector_store_id]
}
# Add metadata filter if provided
metadata_filter = vector_store_search_optional_params.get("filters")
if metadata_filter:
# Convert to Gemini filter syntax if it's a dict
if isinstance(metadata_filter, dict):
# Simple conversion - may need more sophisticated mapping
filter_parts = []
for key, value in metadata_filter.items():
if isinstance(value, str):
filter_parts.append(f'{key} = "{value}"')
else:
filter_parts.append(f'{key} = {value}')
file_search_config["metadata_filter"] = " AND ".join(filter_parts)
else:
file_search_config["metadata_filter"] = metadata_filter
# Build request body
request_body: Dict[str, Any] = {
"contents": [
{
"parts": [{"text": query}]
}
],
"tools": [
{
"file_search": file_search_config
}
],
}
# Add max_num_results if specified
max_results = vector_store_search_optional_params.get("max_num_results")
if max_results:
# This might need to be added to generationConfig or tool config
# depending on Gemini's API requirements
request_body.setdefault("generationConfig", {})["candidateCount"] = 1
litellm_logging_obj.model_call_details["query"] = query
litellm_logging_obj.model_call_details["vector_store_id"] = vector_store_id
return url, request_body
def transform_search_vector_store_response(
self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj
) -> VectorStoreSearchResponse:
"""
Transform Gemini's generateContent response to standard format.
Extracts grounding metadata and citations from the response.
"""
try:
response_data = response.json()
results: List[VectorStoreSearchResult] = []
# Extract candidates and grounding metadata
candidates = response_data.get("candidates", [])
for candidate in candidates:
grounding_metadata = candidate.get("groundingMetadata", {})
grounding_chunks = grounding_metadata.get("groundingChunks", [])
# Process each grounding chunk
for chunk in grounding_chunks:
retrieved_context = chunk.get("retrievedContext")
if retrieved_context:
# This is from file search
text = retrieved_context.get("text", "")
uri = retrieved_context.get("uri", "")
title = retrieved_context.get("title", "")
# Extract file_id from URI if available
file_id = uri if uri else None
results.append(
VectorStoreSearchResult(
score=None, # Gemini doesn't provide explicit scores
content=[VectorStoreResultContent(text=text, type="text")],
file_id=file_id,
filename=title if title else None,
attributes={
"uri": uri,
"title": title,
},
)
)
# Also extract from grounding supports for more detailed citations
grounding_supports = grounding_metadata.get("groundingSupports", [])
for support in grounding_supports:
segment = support.get("segment", {})
text = segment.get("text", "")
grounding_chunk_indices = support.get("groundingChunkIndices", [])
confidence_scores = support.get("confidenceScores", [])
# Use first confidence score as relevance score
score = confidence_scores[0] if confidence_scores else None
# Only add if we have meaningful text and it's not a duplicate
if text:
already_exists = False
for record in results:
contents = record.get("content") or []
if contents and contents[0].get("text") == text:
already_exists = True
break
if already_exists:
continue
results.append(
VectorStoreSearchResult(
score=score,
content=[VectorStoreResultContent(text=text, type="text")],
attributes={
"grounding_chunk_indices": grounding_chunk_indices,
},
)
)
query = litellm_logging_obj.model_call_details.get("query", "")
return VectorStoreSearchResponse(
object="vector_store.search_results.page",
search_query=query,
data=results,
)
except Exception as e:
raise self.get_error_class(
error_message=f"Failed to parse Gemini response: {str(e)}",
status_code=response.status_code,
headers=response.headers,
)
def transform_create_vector_store_request(
self,
vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams,
api_base: str,
) -> Tuple[str, Dict]:
"""
Transform create request to Gemini's fileSearchStores format.
"""
url = f"{api_base}/fileSearchStores"
# Append API key as query parameter (required by Gemini)
api_key = self._cached_api_key or get_api_key_from_env()
if api_key:
url = f"{url}?key={api_key}"
request_body: Dict[str, Any] = {}
# Add display name if provided
name = vector_store_create_optional_params.get("name")
if name:
request_body["displayName"] = name
return url, request_body
def transform_create_vector_store_response(
self, response: httpx.Response
) -> VectorStoreCreateResponse:
"""
Transform Gemini's fileSearchStore response to standard format.
"""
try:
response_data = response.json()
# Extract store name (format: fileSearchStores/xxxxxxx)
store_name = response_data.get("name", "")
display_name = response_data.get("displayName", "")
create_time = response_data.get("createTime", "")
# Convert ISO timestamp to Unix timestamp
import datetime
created_at = None
if create_time:
try:
dt = datetime.datetime.fromisoformat(create_time.replace("Z", "+00:00"))
created_at = int(dt.timestamp())
except Exception:
created_at = None
return VectorStoreCreateResponse(
id=store_name,
object="vector_store",
created_at=created_at or 0,
name=display_name,
bytes=0, # Gemini doesn't provide size info on creation
file_counts=VectorStoreFileCounts(
in_progress=0,
completed=0,
failed=0,
cancelled=0,
total=0,
),
status="completed",
expires_after=None,
expires_at=None,
last_active_at=None,
metadata=None,
)
except Exception as e:
raise self.get_error_class(
error_message=f"Failed to parse Gemini create response: {str(e)}",
status_code=response.status_code,
headers=response.headers,
)

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"

View file

@ -12,6 +12,8 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
import httpx
from httpx._types import RequestFiles
from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS
from litellm.images.utils import ImageEditRequestUtils
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
from litellm.llms.vertex_ai.common_utils import (
_convert_vertex_datetime_to_openai_datetime,
@ -23,8 +25,6 @@ from litellm.types.videos.utils import (
encode_video_id_with_provider,
extract_original_video_id,
)
from litellm.images.utils import ImageEditRequestUtils
from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@ -172,10 +172,16 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
"""
# Extract Vertex AI parameters using safe helpers from VertexBase
# Use safe_get_* methods that don't mutate litellm_params dict
litellm_params = litellm_params or {}
litellm_params_dict: Dict[str, Any] = (
litellm_params.model_dump() if litellm_params else {}
)
vertex_project = VertexBase.safe_get_vertex_ai_project(litellm_params=litellm_params)
vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(litellm_params=litellm_params)
vertex_project = VertexBase.safe_get_vertex_ai_project(
litellm_params=litellm_params_dict
)
vertex_credentials = VertexBase.safe_get_vertex_ai_credentials(
litellm_params=litellm_params_dict
)
# Get access token from Vertex credentials
access_token, project_id = self.get_access_token(

View file

@ -9522,6 +9522,15 @@
"output_cost_per_token": 0.0,
"supports_embedding_image_input": true
},
"embed-multilingual-light-v3.0": {
"input_cost_per_token": 1e-04,
"litellm_provider": "cohere",
"max_input_tokens": 1024,
"max_tokens": 1024,
"mode": "embedding",
"output_cost_per_token": 0.0,
"supports_embedding_image_input": true
},
"eu.amazon.nova-lite-v1:0": {
"input_cost_per_token": 7.8e-08,
"litellm_provider": "bedrock_converse",

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",

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

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")

View file

@ -148,6 +148,8 @@ class VertexPassthroughLoggingHandler:
logging_obj.model = model
logging_obj.model_call_details["model"] = logging_obj.model
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai"
logging_obj.custom_llm_provider = "vertex_ai"
response_cost = litellm.completion_cost(
completion_response=litellm_prediction_response,
model=model,
@ -156,6 +158,7 @@ class VertexPassthroughLoggingHandler:
kwargs["response_cost"] = response_cost
kwargs["model"] = model
kwargs["custom_llm_provider"] = "vertex_ai"
logging_obj.model_call_details["response_cost"] = response_cost
return {

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
@ -1240,7 +1241,7 @@ def cost_tracking():
global prisma_client
if prisma_client is not None:
litellm.logging_callback_manager.add_litellm_callback(_ProxyDBLogger())
litellm.logging_callback_manager.add_litellm_async_success_callback(_ProxyDBLogger())
async def update_cache( # noqa: PLR0915
token: Optional[str],
@ -10157,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)

View file

@ -0,0 +1,6 @@
"""RAG Endpoints for LiteLLM Proxy."""
from litellm.proxy.rag_endpoints.endpoints import router
__all__ = ["router"]

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)},
)

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)

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;
"""

View file

@ -50,11 +50,11 @@ def _update_request_data_with_litellm_managed_vector_store_registry(
@router.post(
"/v1/vector_stores/{vector_store_id}/search",
"/v1/vector_stores/{vector_store_id:path}/search",
dependencies=[Depends(user_api_key_auth)],
)
@router.post(
"/vector_stores/{vector_store_id}/search", dependencies=[Depends(user_api_key_auth)]
"/vector_stores/{vector_store_id:path}/search", dependencies=[Depends(user_api_key_auth)]
)
async def vector_store_search(
request: Request,

22
litellm/rag/__init__.py Normal file
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)

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",
]

View file

@ -0,0 +1,319 @@
"""
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, cast
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: Dict[str, Any] = cast(
Dict[str, Any],
ingest_options.get("chunking_strategy") or {"type": "auto"},
)
self.embedding_config = ingest_options.get("embedding")
self.vector_store_config: Dict[str, Any] = cast(
Dict[str, Any], 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,
)

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

View file

@ -0,0 +1,319 @@
"""
Gemini-specific RAG Ingestion implementation.
Gemini handles embedding and chunking internally when files are uploaded to File Search 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 httpx
from litellm._logging import verbose_logger
from litellm.llms.gemini.common_utils import GeminiModelInfo
from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion
if TYPE_CHECKING:
from litellm import Router
from litellm.types.rag import RAGIngestOptions
class GeminiRAGIngestion(BaseRAGIngestion):
"""
Gemini-specific RAG ingestion using File Search API.
Key differences from base:
- Embedding is handled by Gemini when files are uploaded to File Search stores
- Files are uploaded using uploadToFileSearchStore API
- Chunking is done by Gemini's File Search (supports custom white_space_config)
- Supports custom metadata attachment
"""
def __init__(
self,
ingest_options: "RAGIngestOptions",
router: Optional["Router"] = None,
):
super().__init__(ingest_options=ingest_options, router=router)
self.model_info = GeminiModelInfo()
async def embed(
self,
chunks: List[str],
) -> Optional[List[List[float]]]:
"""
Gemini handles embedding internally - skip this step.
Returns:
None (Gemini embeds when files are uploaded to File Search store)
"""
# Gemini handles embedding when files are uploaded to File Search 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 Gemini File Search store.
Gemini workflow:
1. Create File Search store (if not provided)
2. Upload file using uploadToFileSearchStore (Gemini handles chunking/embedding)
Args:
file_content: Raw file bytes
filename: Name of the file
content_type: MIME type
chunks: Ignored - Gemini handles chunking
embeddings: Ignored - Gemini handles embedding
Returns:
Tuple of (vector_store_id, file_id)
"""
vector_store_id = self.vector_store_config.get("vector_store_id")
vector_store_config = cast(Dict[str, Any], self.vector_store_config)
# Get API credentials
api_key = cast(Optional[str], vector_store_config.get("api_key")) or GeminiModelInfo.get_api_key()
api_base = cast(Optional[str], vector_store_config.get("api_base")) or GeminiModelInfo.get_api_base()
if not api_key:
raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY is required for Gemini File Search")
if not api_base:
raise ValueError("GEMINI_API_BASE is required")
api_version = "v1beta"
base_url = f"{api_base}/{api_version}"
# Create File Search store if not provided
if not vector_store_id:
vector_store_id = await self._create_file_search_store(
api_key=api_key,
base_url=base_url,
display_name=self.ingest_name or "litellm-rag-ingest",
)
# Upload file to File Search store
result_file_id = None
if file_content and filename and vector_store_id:
result_file_id = await self._upload_to_file_search_store(
api_key=api_key,
base_url=base_url,
vector_store_id=vector_store_id,
filename=filename,
file_content=file_content,
content_type=content_type,
)
return vector_store_id, result_file_id
async def _create_file_search_store(
self,
api_key: str,
base_url: str,
display_name: str,
) -> str:
"""
Create a Gemini File Search store.
Args:
api_key: Gemini API key
base_url: Base URL for Gemini API
display_name: Display name for the store
Returns:
Store name (format: fileSearchStores/xxxxxxx)
"""
url = f"{base_url}/fileSearchStores?key={api_key}"
request_body = {
"displayName": display_name
}
async with httpx.AsyncClient() as client:
response = await client.post(
url,
json=request_body,
headers={"Content-Type": "application/json"},
timeout=60.0,
)
if response.status_code != 200:
error_msg = f"Failed to create File Search store: {response.text}"
verbose_logger.error(error_msg)
raise Exception(error_msg)
response_data = response.json()
store_name = response_data.get("name", "")
verbose_logger.debug(f"Created File Search store: {store_name}")
return store_name
async def _upload_to_file_search_store(
self,
api_key: str,
base_url: str,
vector_store_id: str,
filename: str,
file_content: bytes,
content_type: Optional[str],
) -> str:
"""
Upload a file to Gemini File Search store using resumable upload.
Args:
api_key: Gemini API key
base_url: Base URL for Gemini API
vector_store_id: File Search store name
filename: Name of the file
file_content: File content bytes
content_type: MIME type
Returns:
File ID or document name
"""
# Step 1: Initiate resumable upload
upload_url = await self._initiate_resumable_upload(
api_key=api_key,
base_url=base_url,
vector_store_id=vector_store_id,
filename=filename,
file_size=len(file_content),
content_type=content_type or "application/octet-stream",
)
# Step 2: Upload the file content
file_id = await self._upload_file_content(
upload_url=upload_url,
file_content=file_content,
)
return file_id
async def _initiate_resumable_upload(
self,
api_key: str,
base_url: str,
vector_store_id: str,
filename: str,
file_size: int,
content_type: str,
) -> str:
"""
Initiate a resumable upload session.
Returns:
Upload URL for the resumable session
"""
# Construct the upload URL - need to use the full upload endpoint
# base_url is like: https://generativelanguage.googleapis.com/v1beta
# We need: https://generativelanguage.googleapis.com/upload/v1beta/{store_id}:uploadToFileSearchStore
api_base = base_url.replace("/v1beta", "") # Get base without version
url = f"{api_base}/upload/v1beta/{vector_store_id}:uploadToFileSearchStore?key={api_key}"
# Build request body with chunking config and metadata if provided
request_body: Dict[str, Any] = {
"displayName": filename
}
# Add chunking configuration if provided
chunking_strategy = self.chunking_strategy
if chunking_strategy and isinstance(chunking_strategy, dict):
white_space_config = chunking_strategy.get("white_space_config")
if white_space_config:
request_body["chunkingConfig"] = {
"whiteSpaceConfig": {
"maxTokensPerChunk": white_space_config.get("max_tokens_per_chunk", 800),
"maxOverlapTokens": white_space_config.get("max_overlap_tokens", 400),
}
}
# Add custom metadata if provided in vector_store_config
custom_metadata = cast(Optional[List[Dict[str, Any]]], self.vector_store_config.get("custom_metadata"))
if custom_metadata:
request_body["customMetadata"] = custom_metadata
headers = {
"X-Goog-Upload-Protocol": "resumable",
"X-Goog-Upload-Command": "start",
"X-Goog-Upload-Header-Content-Length": str(file_size),
"X-Goog-Upload-Header-Content-Type": content_type,
"Content-Type": "application/json",
}
verbose_logger.debug(f"Initiating resumable upload: {url}")
async with httpx.AsyncClient() as client:
response = await client.post(
url,
json=request_body,
headers=headers,
timeout=60.0,
)
if response.status_code not in [200, 201]:
error_msg = f"Failed to initiate upload: {response.text}"
verbose_logger.error(error_msg)
raise Exception(error_msg)
verbose_logger.debug(f"Initiate resumable upload response: {response.headers}")
# Extract upload URL from response headers
upload_url = response.headers.get("x-goog-upload-url")
if not upload_url:
raise Exception("No upload URL returned in response headers")
verbose_logger.debug(f"Got upload URL: {upload_url}")
return upload_url
async def _upload_file_content(
self,
upload_url: str,
file_content: bytes,
) -> str:
"""
Upload file content to the resumable upload URL.
Returns:
File ID or document name from the response
"""
headers = {
"Content-Length": str(len(file_content)),
"X-Goog-Upload-Offset": "0",
"X-Goog-Upload-Command": "upload, finalize",
}
verbose_logger.debug(f"Uploading file content ({len(file_content)} bytes)")
async with httpx.AsyncClient() as client:
response = await client.put(
upload_url,
content=file_content,
headers=headers,
timeout=300.0, # Longer timeout for large files
)
if response.status_code not in [200, 201]:
error_msg = f"Failed to upload file: {response.text}"
verbose_logger.error(error_msg)
raise Exception(error_msg)
# Parse response to get file/document ID
try:
response_data = response.json()
# The response should contain the document name or file reference
file_id = response_data.get("name", "") or response_data.get("file", {}).get("name", "")
verbose_logger.debug(f"Upload complete. File ID: {file_id}")
return file_id
except Exception as e:
verbose_logger.warning(f"Could not parse upload response: {e}")
# Return a placeholder if we can't get the ID
return "uploaded"

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

242
litellm/rag/main.py Normal file
View file

@ -0,0 +1,242 @@
"""
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.gemini_ingestion import GeminiRAGIngestion
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,
"gemini": GeminiRAGIngestion,
}
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,
)

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"]

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

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

146
litellm/types/rag.py Normal file
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

View file

@ -7613,6 +7613,12 @@ class ProviderConfigManager:
)
return MilvusVectorStoreConfig()
elif litellm.LlmProviders.GEMINI == provider:
from litellm.llms.gemini.vector_stores.transformation import (
GeminiVectorStoreConfig,
)
return GeminiVectorStoreConfig()
return None
@staticmethod

View file

@ -9522,6 +9522,15 @@
"output_cost_per_token": 0.0,
"supports_embedding_image_input": true
},
"embed-multilingual-light-v3.0": {
"input_cost_per_token": 1e-04,
"litellm_provider": "cohere",
"max_input_tokens": 1024,
"max_tokens": 1024,
"mode": "embedding",
"output_cost_per_token": 0.0,
"supports_embedding_image_input": true
},
"eu.amazon.nova-lite-v1:0": {
"input_cost_per_token": 7.8e-08,
"litellm_provider": "bedrock_converse",

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 = [

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

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.
]

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"

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):

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"

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"]

View file

@ -238,6 +238,30 @@ def test_transform_tool_call_with_cache_control():
assert "cachePoint" in transformed_cache_msg
assert transformed_cache_msg["cachePoint"]["type"] == "default"
def test_reasoning_with_forced_tool_choice_switches_to_auto():
config = AmazonConverseConfig()
non_default_params = {
"tools": [
{
"type": "function",
"function": {"name": "get_current_weather", "parameters": {}},
}
],
"tool_choice": "required",
"reasoning_effort": "low",
}
optional_params = config.map_openai_params(
model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0",
non_default_params=non_default_params,
optional_params={},
drop_params=False,
)
assert optional_params["tool_choice"] == {"auto": {}}
def test_get_supported_openai_params():
config = AmazonConverseConfig()
supported_params = config.get_supported_openai_params(
@ -2592,8 +2616,10 @@ def test_empty_assistant_message_handling():
empty or whitespace-only content with a placeholder to prevent AWS Bedrock
Converse API 400 Bad Request errors.
"""
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
from litellm.litellm_core_utils.prompt_templates.factory import (
_bedrock_converse_messages_pt,
)
# Test case 1: Empty string content - test with modify_params=True to prevent merging
messages = [
{"role": "user", "content": "Hello"},

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
# ============================================================================

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")

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.

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

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

View file

@ -0,0 +1,29 @@
"""
Minimal Gemini File Search vector store tests.
"""
import os
import sys
sys.path.insert(0, os.path.abspath("../.."))
from base_vector_store_test import BaseVectorStoreTest
class TestGeminiVectorStore(BaseVectorStoreTest):
"""Reuses the shared vector store smoke suite with Gemini."""
def get_base_request_args(self) -> dict:
"""Provide arguments for the shared search test."""
return {
"vector_store_id": os.getenv("GEMINI_TEST_STORE_ID", "fileSearchStores/example-test-store"),
"custom_llm_provider": "gemini",
"query": "LiteLLM",
}
def get_base_create_vector_store_args(self) -> dict:
"""Ensure we always call Gemini when creating a vector store."""
return {
"custom_llm_provider": "gemini",
}

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();
});
});

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();
});
});

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>

View file

@ -1,4 +1,5 @@
import { KeyIcon, TrashIcon } from "@heroicons/react/outline";
import { KeyIcon, TrashIcon } from "@heroicons/react/outline";
import { ColumnDef } from "@tanstack/react-table";
import { Badge, Button, Icon } from "@tremor/react";
import { Tooltip } from "antd";
@ -114,18 +115,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>
);

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");
});
});

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>
);
},

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();
});
});

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;

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>
)}

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();
});
});

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>

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}

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");
});
});

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() ? (
{