mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge remote-tracking branch 'origin' into litellm_ui_model_page_perf
This commit is contained in:
commit
ce65663ad1
69 changed files with 4870 additions and 270 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
414
docs/my-website/docs/providers/gemini_file_search.md
Normal file
414
docs/my-website/docs/providers/gemini_file_search.md
Normal 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)
|
||||
|
||||
|
|
@ -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
|
||||
|
|
|
|||
273
docs/my-website/docs/rag_ingest.md
Normal file
273
docs/my-website/docs/rag_ingest.md
Normal 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"}}
|
||||
}'
|
||||
```
|
||||
|
||||
|
|
@ -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>
|
||||
|
||||
|
|
|
|||
|
|
@ -412,6 +412,7 @@ const sidebars = {
|
|||
"proxy/pass_through"
|
||||
]
|
||||
},
|
||||
"rag_ingest",
|
||||
"realtime",
|
||||
"rerank",
|
||||
"response_api",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any, Optional, Union
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
6
litellm/llms/gemini/vector_stores/__init__.py
Normal file
6
litellm/llms/gemini/vector_stores/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
"""Gemini File Search Vector Store module."""
|
||||
|
||||
from .transformation import GeminiVectorStoreConfig
|
||||
|
||||
__all__ = ["GeminiVectorStoreConfig"]
|
||||
|
||||
357
litellm/llms/gemini/vector_stores/transformation.py
Normal file
357
litellm/llms/gemini/vector_stores/transformation.py
Normal 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,
|
||||
)
|
||||
|
||||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
6
litellm/proxy/rag_endpoints/__init__.py
Normal file
6
litellm/proxy/rag_endpoints/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
"""RAG Endpoints for LiteLLM Proxy."""
|
||||
|
||||
from litellm.proxy.rag_endpoints.endpoints import router
|
||||
|
||||
__all__ = ["router"]
|
||||
|
||||
200
litellm/proxy/rag_endpoints/endpoints.py
Normal file
200
litellm/proxy/rag_endpoints/endpoints.py
Normal 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)},
|
||||
)
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
22
litellm/rag/__init__.py
Normal 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)
|
||||
|
||||
14
litellm/rag/ingestion/__init__.py
Normal file
14
litellm/rag/ingestion/__init__.py
Normal 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",
|
||||
]
|
||||
|
||||
319
litellm/rag/ingestion/base_ingestion.py
Normal file
319
litellm/rag/ingestion/base_ingestion.py
Normal 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,
|
||||
)
|
||||
|
||||
619
litellm/rag/ingestion/bedrock_ingestion.py
Normal file
619
litellm/rag/ingestion/bedrock_ingestion.py
Normal 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
|
||||
|
||||
319
litellm/rag/ingestion/gemini_ingestion.py
Normal file
319
litellm/rag/ingestion/gemini_ingestion.py
Normal 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"
|
||||
|
||||
111
litellm/rag/ingestion/openai_ingestion.py
Normal file
111
litellm/rag/ingestion/openai_ingestion.py
Normal 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
242
litellm/rag/main.py
Normal 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,
|
||||
)
|
||||
10
litellm/rag/text_splitters/__init__.py
Normal file
10
litellm/rag/text_splitters/__init__.py
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
"""
|
||||
Text splitting utilities for RAG ingestion.
|
||||
"""
|
||||
|
||||
from litellm.rag.text_splitters.recursive_character_text_splitter import (
|
||||
RecursiveCharacterTextSplitter,
|
||||
)
|
||||
|
||||
__all__ = ["RecursiveCharacterTextSplitter"]
|
||||
|
||||
135
litellm/rag/text_splitters/recursive_character_text_splitter.py
Normal file
135
litellm/rag/text_splitters/recursive_character_text_splitter.py
Normal 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
|
||||
|
||||
|
|
@ -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
146
litellm/types/rag.py
Normal 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
|
||||
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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 = [
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -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"]
|
||||
|
|
@ -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"},
|
||||
|
|
|
|||
|
|
@ -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
|
||||
# ============================================================================
|
||||
|
|
|
|||
173
tests/vector_store_tests/rag/base_rag_tests.py
Normal file
173
tests/vector_store_tests/rag/base_rag_tests.py
Normal 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")
|
||||
|
||||
4
tests/vector_store_tests/rag/test_document.txt
Normal file
4
tests/vector_store_tests/rag/test_document.txt
Normal 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.
|
||||
|
||||
91
tests/vector_store_tests/rag/test_rag_bedrock.py
Normal file
91
tests/vector_store_tests/rag/test_rag_bedrock.py
Normal 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
|
||||
|
||||
45
tests/vector_store_tests/rag/test_rag_openai.py
Normal file
45
tests/vector_store_tests/rag/test_rag_openai.py
Normal 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
|
||||
|
||||
|
||||
29
tests/vector_store_tests/test_gemini_vector_store.py
Normal file
29
tests/vector_store_tests/test_gemini_vector_store.py
Normal 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",
|
||||
}
|
||||
|
||||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
|
|
@ -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>
|
||||
);
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
@ -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;
|
||||
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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() ? (
|
||||
{
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue