mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-06 14:24:50 +00:00
Merge pull request #17124 from BerriAI/litellm_gemini_file_search
Add gemini file search support
This commit is contained in:
Sameer Kankute
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -6,7 +6,7 @@ All-in-one document ingestion pipeline: **Upload → Chunk → Embed → Vector
|
||||
|---------|-----------|
|
||||
| Cost Tracking | ❌ |
|
||||
| Logging | ✅ |
|
||||
| Supported Providers | `openai`, `bedrock` |
|
||||
| Supported Providers | `openai`, `bedrock`, `gemini` |
|
||||
|
||||
## Quick Start
|
||||
|
||||
@@ -50,6 +50,52 @@ curl -X POST "http://localhost:4000/v1/rag/ingest" \
|
||||
}"
|
||||
```
|
||||
|
||||
### 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
|
||||
|
||||
@@ -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>
|
||||
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
"""Gemini File Search Vector Store module."""
|
||||
|
||||
from .transformation import GeminiVectorStoreConfig
|
||||
|
||||
__all__ = ["GeminiVectorStoreConfig"]
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -14,7 +14,7 @@ from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -49,9 +49,14 @@ class BaseRAGIngestion(ABC):
|
||||
|
||||
# Extract configs from options
|
||||
self.ocr_config = ingest_options.get("ocr")
|
||||
self.chunking_strategy = ingest_options.get("chunking_strategy", {"type": "auto"})
|
||||
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 = ingest_options.get("vector_store") or {}
|
||||
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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -19,6 +19,7 @@ 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
|
||||
@@ -31,6 +32,7 @@ if TYPE_CHECKING:
|
||||
INGESTION_REGISTRY: Dict[str, Type[BaseRAGIngestion]] = {
|
||||
"openai": OpenAIRAGIngestion,
|
||||
"bedrock": BedrockRAGIngestion,
|
||||
"gemini": GeminiRAGIngestion,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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",
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user