From fe444f3ed5cae334d013450f96befdad113a09a2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 27 Jan 2026 14:45:26 -0800 Subject: [PATCH] [Feat] RAG API - Add support for using s3 Vectors as Vector Store Provider for /rag/ingest (#19888) * init S3VectorsRAGIngestion as a supported ingestion provider for RAG API * test: TestRAGS3Vectors * init S3VectorsVectorStoreOptions * init s3 vectors * code clean up + QA * fix: get_credentials * S3VectorsRAGIngestion * TestRAGS3Vectors * docs: AWS S3 Vectors * add asyncio QA checks * fix: S3_VECTORS_DEFAULT_DIMENSION --- docs/my-website/docs/rag_ingest.md | 78 ++- litellm/constants.py | 7 + ...odel_prices_and_context_window_backup.json | 42 ++ litellm/rag/ingestion/s3_vectors_ingestion.py | 546 ++++++++++++++++++ litellm/rag/main.py | 2 + litellm/types/rag.py | 46 +- .../rag/test_rag_s3_vectors.py | 107 ++++ 7 files changed, 826 insertions(+), 2 deletions(-) create mode 100644 litellm/rag/ingestion/s3_vectors_ingestion.py create mode 100644 tests/vector_store_tests/rag/test_rag_s3_vectors.py diff --git a/docs/my-website/docs/rag_ingest.md b/docs/my-website/docs/rag_ingest.md index 1133b85f20..7adc2d70b5 100644 --- a/docs/my-website/docs/rag_ingest.md +++ b/docs/my-website/docs/rag_ingest.md @@ -5,7 +5,7 @@ All-in-one document ingestion pipeline: **Upload → Chunk → Embed → Vector | Feature | Supported | |---------|-----------| | Logging | Yes | -| Supported Providers | `openai`, `bedrock`, `vertex_ai`, `gemini` | +| Supported Providers | `openai`, `bedrock`, `vertex_ai`, `gemini`, `s3_vectors` | :::tip After ingesting documents, use [/rag/query](./rag_query.md) to search and generate responses with your ingested content. @@ -75,6 +75,31 @@ curl -X POST "http://localhost:4000/v1/rag/ingest" \ }" ``` +### AWS S3 Vectors + +```bash showLineNumbers title="Ingest to S3 Vectors" +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\": { + \"embedding\": { + \"model\": \"text-embedding-3-small\" + }, + \"vector_store\": { + \"custom_llm_provider\": \"s3_vectors\", + \"vector_bucket_name\": \"my-embeddings\", + \"aws_region_name\": \"us-west-2\" + } + } + }" +``` + ## Response ```json @@ -265,6 +290,57 @@ When `vector_store_id` is omitted, LiteLLM automatically creates: 4. Install: `pip install 'google-cloud-aiplatform>=1.60.0'` ::: +### vector_store (AWS S3 Vectors) + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `custom_llm_provider` | string | - | `"s3_vectors"` | +| `vector_bucket_name` | string | **required** | S3 vector bucket name | +| `index_name` | string | auto-create | Vector index name | +| `dimension` | integer | auto-detect | Vector dimension (auto-detected from embedding model) | +| `distance_metric` | string | `cosine` | Distance metric: `cosine` or `euclidean` | +| `non_filterable_metadata_keys` | array | `["source_text"]` | Metadata keys excluded from filtering | +| `aws_region_name` | string | `us-west-2` | AWS region | +| `aws_access_key_id` | string | env | AWS access key | +| `aws_secret_access_key` | string | env | AWS secret key | + +:::info S3 Vectors Auto-Creation +When `index_name` is omitted, LiteLLM automatically creates: +- S3 vector bucket (if it doesn't exist) +- Vector index with auto-detected dimensions from your embedding model + +**Dimension Auto-Detection**: The vector dimension is automatically detected by making a test embedding request to your specified model. No need to manually specify dimensions! + +**Supported Embedding Models**: Works with any LiteLLM-supported embedding model (OpenAI, Cohere, Bedrock, Azure, etc.) +::: + +**Example with auto-detection:** +```json +{ + "embedding": { + "model": "text-embedding-3-small" // Dimension auto-detected as 1536 + }, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": "my-embeddings" + } +} +``` + +**Example with custom embedding provider:** +```json +{ + "embedding": { + "model": "cohere/embed-english-v3.0" // Dimension auto-detected as 1024 + }, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": "my-embeddings", + "distance_metric": "cosine" + } +} +``` + ## Input Examples ### File (Base64) diff --git a/litellm/constants.py b/litellm/constants.py index 0525bf3843..0c9d1d941b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1327,6 +1327,13 @@ COROUTINE_CHECKER_MAX_SIZE_IN_MEMORY = int( DEFAULT_CHUNK_SIZE = int(os.getenv("DEFAULT_CHUNK_SIZE", 1000)) DEFAULT_CHUNK_OVERLAP = int(os.getenv("DEFAULT_CHUNK_OVERLAP", 200)) +########################### S3 Vectors RAG Constants ########################### +S3_VECTORS_DEFAULT_DIMENSION = int(os.getenv("S3_VECTORS_DEFAULT_DIMENSION", 1024)) +S3_VECTORS_DEFAULT_DISTANCE_METRIC = str( + os.getenv("S3_VECTORS_DEFAULT_DISTANCE_METRIC", "cosine") +) +S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS = ["source_text"] + ########################### Microsoft SSO Constants ########################### MICROSOFT_USER_EMAIL_ATTRIBUTE = str( os.getenv("MICROSOFT_USER_EMAIL_ATTRIBUTE", "userPrincipalName") diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7f33e96dbb..4ac4159558 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10232,6 +10232,48 @@ "mode": "completion", "output_cost_per_token": 5e-07 }, + "deepseek-v3-2-251201": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 98304, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "glm-4-7-251222": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "kimi-k2-thinking-251104": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 229376, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "doubao-embedding": { "input_cost_per_token": 0.0, "litellm_provider": "volcengine", diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py new file mode 100644 index 0000000000..13964cd89a --- /dev/null +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -0,0 +1,546 @@ +""" +S3 Vectors-specific RAG Ingestion implementation. + +S3 Vectors is AWS's native vector storage service that provides: +- Purpose-built vector buckets for storing and querying vectors +- Vector indexes with configurable dimensions and distance metrics +- Metadata filtering for semantic search + +This implementation: +1. Auto-creates vector buckets and indexes if not provided +2. Uses LiteLLM's embedding API (supports any provider) +3. Uses httpx + AWS SigV4 signing (no boto3 dependency for S3 Vectors APIs) +4. Stores vectors with metadata using PutVectors API +""" + +from __future__ import annotations + +import hashlib +import uuid +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import litellm +from litellm._logging import verbose_logger +from litellm.constants import ( + S3_VECTORS_DEFAULT_DIMENSION, + S3_VECTORS_DEFAULT_DISTANCE_METRIC, + S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, +) +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion + +if TYPE_CHECKING: + from litellm import Router + from litellm.types.rag import RAGIngestOptions + + +class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): + """ + S3 Vectors RAG ingestion using httpx + AWS SigV4 signing. + + Workflow: + 1. Auto-create vector bucket if needed (CreateVectorBucket API) + 2. Auto-create vector index if needed (CreateVectorIndex API) + 3. Generate embeddings using LiteLLM (supports any provider) + 4. Store vectors with PutVectors API + + Configuration: + - vector_bucket_name: S3 vector bucket name (required) + - index_name: Vector index name (auto-creates if not provided) + - dimension: Vector dimension (default: S3_VECTORS_DEFAULT_DIMENSION) + - distance_metric: "cosine" or "euclidean" (default: S3_VECTORS_DEFAULT_DISTANCE_METRIC) + - non_filterable_metadata_keys: List of metadata keys to exclude from filtering + """ + + def __init__( + self, + ingest_options: "RAGIngestOptions", + router: Optional["Router"] = None, + ): + BaseRAGIngestion.__init__(self, ingest_options=ingest_options, router=router) + BaseAWSLLM.__init__(self) + + # Extract config + self.vector_bucket_name = self.vector_store_config["vector_bucket_name"] + self.index_name = self.vector_store_config.get("index_name") + self.distance_metric = self.vector_store_config.get( + "distance_metric", S3_VECTORS_DEFAULT_DISTANCE_METRIC + ) + self.non_filterable_metadata_keys = self.vector_store_config.get( + "non_filterable_metadata_keys", + S3_VECTORS_DEFAULT_NON_FILTERABLE_METADATA_KEYS, + ) + + # Get dimension from config (will be auto-detected on first use if not provided) + self.dimension = self._get_dimension_from_config() + + # 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 + ) + + # Create httpx client (similar to s3_v2.py) + ssl_verify = self._get_ssl_verify( + ssl_verify=self.vector_store_config.get("ssl_verify") + ) + self.async_httpx_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.RAG, + params={"ssl_verify": ssl_verify} if ssl_verify is not None else None, + ) + + # Track if infrastructure is initialized + self._config_initialized = False + + async def _get_dimension_from_embedding_request(self) -> int: + """ + Auto-detect dimension by making a test embedding request. + + Makes a single embedding request with a test string to determine + the output dimension of the embedding model. + """ + if not self.embedding_config or "model" not in self.embedding_config: + return S3_VECTORS_DEFAULT_DIMENSION + + try: + model_name = self.embedding_config["model"] + verbose_logger.debug( + f"Auto-detecting dimension by making test embedding request to {model_name}" + ) + + # Make a test embedding request + test_input = "test" + if self.router: + response = await self.router.aembedding(model=model_name, input=[test_input]) + else: + response = await litellm.aembedding(model=model_name, input=[test_input]) + + # Get dimension from the response + if response.data and len(response.data) > 0: + dimension = len(response.data[0]["embedding"]) + verbose_logger.debug( + f"Auto-detected dimension {dimension} for embedding model {model_name}" + ) + return dimension + except Exception as e: + verbose_logger.warning( + f"Could not auto-detect dimension from embedding model: {e}. " + f"Using default dimension of {S3_VECTORS_DEFAULT_DIMENSION}." + ) + + return S3_VECTORS_DEFAULT_DIMENSION + + def _get_dimension_from_config(self) -> Optional[int]: + """ + Get vector dimension from config if explicitly provided. + + Returns None if dimension should be auto-detected. + """ + if "dimension" in self.vector_store_config: + return int(self.vector_store_config["dimension"]) + return None + + async def _ensure_config_initialized(self): + """Lazily initialize S3 Vectors infrastructure.""" + if self._config_initialized: + return + + # Auto-detect dimension if not provided + if self.dimension is None: + self.dimension = await self._get_dimension_from_embedding_request() + + # Ensure vector bucket exists + await self._ensure_vector_bucket_exists() + + # Ensure vector index exists + if not self.index_name: + # Auto-generate index name + unique_id = uuid.uuid4().hex[:8] + self.index_name = f"litellm-index-{unique_id}" + + await self._ensure_vector_index_exists() + + self._config_initialized = True + + async def _sign_and_execute_request( + self, + method: str, + url: str, + data: Optional[str] = None, + headers: Optional[Dict[str, str]] = None, + ) -> Any: + """ + Helper to sign and execute AWS API requests using httpx + SigV4. + + Pattern from litellm/integrations/s3_v2.py + """ + try: + import requests + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError( + "Missing botocore to call S3 Vectors. Run 'pip install boto3'." + ) + + # Get AWS credentials using BaseAWSLLM's get_credentials method + credentials = self.get_credentials( + aws_access_key_id=self.vector_store_config.get("aws_access_key_id"), + aws_secret_access_key=self.vector_store_config.get("aws_secret_access_key"), + aws_session_token=self.vector_store_config.get("aws_session_token"), + aws_region_name=self.aws_region_name, + aws_session_name=self.vector_store_config.get("aws_session_name"), + aws_profile_name=self.vector_store_config.get("aws_profile_name"), + aws_role_name=self.vector_store_config.get("aws_role_name"), + aws_web_identity_token=self.vector_store_config.get("aws_web_identity_token"), + aws_sts_endpoint=self.vector_store_config.get("aws_sts_endpoint"), + aws_external_id=self.vector_store_config.get("aws_external_id"), + ) + + # Prepare headers + if headers is None: + headers = {} + + if data: + headers["Content-Type"] = "application/json" + # Calculate SHA256 hash of the content + content_hash = hashlib.sha256(data.encode("utf-8")).hexdigest() + headers["x-amz-content-sha256"] = content_hash + else: + # For requests without body, use hash of empty string + headers["x-amz-content-sha256"] = hashlib.sha256(b"").hexdigest() + + # Prepare the request + req = requests.Request(method, url, data=data, headers=headers) + prepped = req.prepare() + + # Sign the request + aws_request = AWSRequest( + method=prepped.method, + url=prepped.url, + data=prepped.body, + headers=prepped.headers, + ) + SigV4Auth(credentials, "s3vectors", self.aws_region_name).add_auth(aws_request) + + # Prepare the signed headers + signed_headers = dict(aws_request.headers.items()) + + # Make the request using specific method (pattern from s3_v2.py) + method_upper = method.upper() + if method_upper == "PUT": + response = await self.async_httpx_client.put( + url, data=data, headers=signed_headers + ) + elif method_upper == "POST": + response = await self.async_httpx_client.post( + url, data=data, headers=signed_headers + ) + elif method_upper == "GET": + response = await self.async_httpx_client.get(url, headers=signed_headers) + else: + raise ValueError(f"Unsupported HTTP method: {method}") + + return response + + async def _ensure_vector_bucket_exists(self): + """Create vector bucket if it doesn't exist using GetVectorBucket and CreateVectorBucket APIs.""" + verbose_logger.debug( + f"Ensuring S3 vector bucket exists: {self.vector_bucket_name}" + ) + + # Try to get bucket info using GetVectorBucket API + get_url = f"https://s3vectors.{self.aws_region_name}.api.aws/GetVectorBucket" + get_body = safe_dumps({"vectorBucketName": self.vector_bucket_name}) + + try: + response = await self._sign_and_execute_request("POST", get_url, data=get_body) + if response.status_code == 200: + verbose_logger.debug(f"Vector bucket {self.vector_bucket_name} exists") + return + except Exception as e: + verbose_logger.debug( + f"Bucket check failed (may not exist): {e}, attempting to create" + ) + + # Create vector bucket using CreateVectorBucket API + try: + verbose_logger.debug(f"Creating vector bucket: {self.vector_bucket_name}") + create_url = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateVectorBucket" + create_body = safe_dumps({ + "vectorBucketName": self.vector_bucket_name + }) + + response = await self._sign_and_execute_request("POST", create_url, data=create_body) + + if response.status_code in (200, 201): + verbose_logger.info(f"Created vector bucket: {self.vector_bucket_name}") + elif response.status_code == 409: + # Bucket already exists (ConflictException) + verbose_logger.debug( + f"Vector bucket {self.vector_bucket_name} already exists" + ) + else: + verbose_logger.error(f"CreateVectorBucket failed: {response.status_code} - {response.text}") + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error creating vector bucket: {e}") + raise + + async def _ensure_vector_index_exists(self): + """Create vector index if it doesn't exist using GetIndex and CreateIndex APIs.""" + verbose_logger.debug( + f"Ensuring vector index exists: {self.vector_bucket_name}/{self.index_name}" + ) + + # Try to get index info using GetIndex API + get_url = f"https://s3vectors.{self.aws_region_name}.api.aws/GetIndex" + get_body = safe_dumps({ + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name + }) + + try: + response = await self._sign_and_execute_request("POST", get_url, data=get_body) + if response.status_code == 200: + verbose_logger.debug(f"Vector index {self.index_name} exists") + return + except Exception as e: + verbose_logger.debug( + f"Index check failed (may not exist): {e}, attempting to create" + ) + + # Create vector index using CreateIndex API + try: + verbose_logger.debug( + f"Creating vector index: {self.index_name} with dimension={self.dimension}, metric={self.distance_metric}" + ) + + # Prepare index configuration per AWS API docs + index_config = { + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name, + "dataType": "float32", + "dimension": self.dimension, + "distanceMetric": self.distance_metric, + } + + if self.non_filterable_metadata_keys: + index_config["metadataConfiguration"] = { + "nonFilterableMetadataKeys": self.non_filterable_metadata_keys + } + + create_url = f"https://s3vectors.{self.aws_region_name}.api.aws/CreateIndex" + response = await self._sign_and_execute_request( + "POST", create_url, data=safe_dumps(index_config) + ) + + if response.status_code in (200, 201): + verbose_logger.info(f"Created vector index: {self.index_name}") + elif response.status_code == 409: + verbose_logger.debug(f"Vector index {self.index_name} already exists") + else: + verbose_logger.error(f"CreateIndex failed: {response.status_code} - {response.text}") + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error creating vector index: {e}") + raise + + async def _put_vectors(self, vectors: List[Dict[str, Any]]): + """ + Call PutVectors API to store vectors in S3 Vectors. + + Args: + vectors: List of vector objects with keys: "key", "data", "metadata" + """ + verbose_logger.debug( + f"Storing {len(vectors)} vectors in {self.vector_bucket_name}/{self.index_name}" + ) + + url = f"https://s3vectors.{self.aws_region_name}.api.aws/PutVectors" + + # Prepare request body per AWS API docs + request_body = { + "vectorBucketName": self.vector_bucket_name, + "indexName": self.index_name, + "vectors": vectors + } + + try: + response = await self._sign_and_execute_request( + "POST", url, data=safe_dumps(request_body) + ) + + if response.status_code in (200, 201): + verbose_logger.info( + f"Successfully stored {len(vectors)} vectors in index {self.index_name}" + ) + else: + verbose_logger.error( + f"PutVectors failed with status {response.status_code}: {response.text}" + ) + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error storing vectors: {e}") + raise + + async def embed( + self, + chunks: List[str], + ) -> Optional[List[List[float]]]: + """ + Generate embeddings using LiteLLM's embedding API. + + Supports any embedding provider (OpenAI, Bedrock, Cohere, etc.) + """ + if not chunks: + return None + + # Use embedding config from ingest_options or default + if not self.embedding_config: + verbose_logger.warning( + "No embedding config provided, using default text-embedding-3-small" + ) + self.embedding_config = {"model": "text-embedding-3-small"} + + embedding_model = self.embedding_config.get("model", "text-embedding-3-small") + + verbose_logger.debug( + f"Generating embeddings for {len(chunks)} chunks using {embedding_model}" + ) + + # Convert to list to ensure type compatibility + input_chunks: List[str] = list(chunks) + + if self.router: + response = await self.router.aembedding(model=embedding_model, input=input_chunks) + else: + response = await litellm.aembedding(model=embedding_model, input=input_chunks) + + return [item["embedding"] for item in response.data] + + 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 vectors in S3 Vectors using PutVectors API. + + Steps: + 1. Ensure vector bucket exists (auto-create if needed) + 2. Ensure vector index exists (auto-create if needed) + 3. Prepare vector data with metadata + 4. Call PutVectors API with httpx + SigV4 signing + + Args: + file_content: Raw file bytes (not used for S3 Vectors) + filename: Name of the file + content_type: MIME type (not used for S3 Vectors) + chunks: Text chunks + embeddings: Vector embeddings + + Returns: + Tuple of (index_name, filename) + """ + # Ensure infrastructure exists + await self._ensure_config_initialized() + + if not embeddings or not chunks: + verbose_logger.warning("No embeddings or chunks to store") + return self.index_name, None + + # Prepare vectors for PutVectors API + vectors = [] + for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)): + vector_obj = { + "key": f"{filename}_{i}" if filename else f"chunk_{i}", + "data": {"float32": embedding}, + "metadata": { + "source_text": chunk, # Non-filterable (for reference) + "chunk_index": str(i), # Filterable + }, + } + + if filename: + vector_obj["metadata"]["filename"] = filename # Filterable + + vectors.append(vector_obj) + + # Call PutVectors API + await self._put_vectors(vectors) + + return self.index_name, filename + + async def query_vector_store( + self, vector_store_id: str, query: str, top_k: int = 5 + ) -> Optional[Dict[str, Any]]: + """ + Query S3 Vectors using QueryVectors API. + + Args: + vector_store_id: Index name + query: Query text + top_k: Number of results to return + + Returns: + Query results with vectors and metadata + """ + verbose_logger.debug(f"Querying index {vector_store_id} with query: {query}") + + # Generate query embedding + if not self.embedding_config: + self.embedding_config = {"model": "text-embedding-3-small"} + + embedding_model = self.embedding_config.get("model", "text-embedding-3-small") + + response = await litellm.aembedding(model=embedding_model, input=[query]) + query_embedding = response.data[0]["embedding"] + + # Call QueryVectors API + url = f"https://s3vectors.{self.aws_region_name}.api.aws/QueryVectors" + + request_body = { + "vectorBucketName": self.vector_bucket_name, + "indexName": vector_store_id, + "queryVector": {"float32": query_embedding}, + "topK": top_k, + "returnDistance": True, + "returnMetadata": True, + } + + try: + response = await self._sign_and_execute_request( + "POST", url, data=safe_dumps(request_body) + ) + + if response.status_code == 200: + results = response.json() + verbose_logger.debug(f"Query returned {len(results.get('vectors', []))} results") + + # Check if query terms appear in results + if results.get("vectors"): + for result in results["vectors"]: + metadata = result.get("metadata", {}) + source_text = metadata.get("source_text", "") + if query.lower() in source_text.lower(): + return results + + # Return results even if exact match not found + return results + else: + verbose_logger.error( + f"QueryVectors failed with status {response.status_code}: {response.text}" + ) + return None + except Exception as e: + verbose_logger.exception(f"Error querying vectors: {e}") + return None diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 620097f83d..571f78f2f7 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -31,6 +31,7 @@ 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.rag.ingestion.s3_vectors_ingestion import S3VectorsRAGIngestion from litellm.rag.rag_query import RAGQuery from litellm.types.rag import ( RAGIngestOptions, @@ -48,6 +49,7 @@ INGESTION_REGISTRY: Dict[str, Type[BaseRAGIngestion]] = { "openai": OpenAIRAGIngestion, "bedrock": BedrockRAGIngestion, "gemini": GeminiRAGIngestion, + "s3_vectors": S3VectorsRAGIngestion, } diff --git a/litellm/types/rag.py b/litellm/types/rag.py index fe237a1343..cae0770868 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -129,9 +129,53 @@ class VertexAIVectorStoreOptions(TypedDict, total=False): import_timeout: Optional[int] # Timeout in seconds (default: 600) +class S3VectorsVectorStoreOptions(TypedDict, total=False): + """ + AWS S3 Vectors configuration. + + Example (auto-create): + {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings"} + + Example (use existing): + {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings", + "index_name": "my-index"} + + Example (with credentials): + {"custom_llm_provider": "s3_vectors", "vector_bucket_name": "my-embeddings", + "litellm_credential_name": "my-aws-creds"} + + Auto-creation creates: S3 vector bucket and vector index (if not provided). + Embeddings are generated using LiteLLM's embedding API (supports any provider). + """ + + custom_llm_provider: Literal["s3_vectors"] + vector_bucket_name: str # Required - S3 vector bucket name + index_name: Optional[str] # Vector index name (auto-creates if not provided) + + # Index configuration (for auto-creation) + dimension: Optional[int] # Vector dimension (auto-detected from embedding model, or default: 1024) + distance_metric: Optional[Literal["cosine", "euclidean"]] # Default: cosine + non_filterable_metadata_keys: Optional[List[str]] # Keys excluded from filtering (e.g., ["source_text"]) + + # Credentials (loaded from litellm.credential_list if litellm_credential_name is provided) + litellm_credential_name: Optional[str] # Credential name to load from litellm.credential_list + + # 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, VertexAIVectorStoreOptions + OpenAIVectorStoreOptions, BedrockVectorStoreOptions, VertexAIVectorStoreOptions, S3VectorsVectorStoreOptions ] diff --git a/tests/vector_store_tests/rag/test_rag_s3_vectors.py b/tests/vector_store_tests/rag/test_rag_s3_vectors.py new file mode 100644 index 0000000000..cd8a362a7b --- /dev/null +++ b/tests/vector_store_tests/rag/test_rag_s3_vectors.py @@ -0,0 +1,107 @@ +""" +S3 Vectors RAG ingestion tests. + +Requires environment variables: +- AWS_ACCESS_KEY_ID +- AWS_SECRET_ACCESS_KEY +- AWS_REGION_NAME (optional, defaults to us-west-2) + +Optional: +- S3_VECTOR_BUCKET_NAME (optional, auto-generates if not set) +""" + +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 +from tests.vector_store_tests.rag.base_rag_tests import BaseRAGTest + + +class TestRAGS3Vectors(BaseRAGTest): + """Test RAG Ingest with AWS S3 Vectors.""" + + @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 S3 Vectors test: AWS credentials required") + + def get_base_ingest_options(self) -> RAGIngestOptions: + """ + Return S3 Vectors-specific ingest options. + + Chunking is configured via chunking_strategy (unified interface). + Embeddings are generated using LiteLLM's embedding API. + """ + vector_bucket_name = os.environ.get( + "S3_VECTOR_BUCKET_NAME", "test-litellm-vectors" + ) + aws_region = os.environ.get("AWS_REGION_NAME", "us-west-2") + + return { + "chunking_strategy": { + "chunk_size": 512, + "chunk_overlap": 100, + }, + "embedding": { + "model": "text-embedding-3-small" # Can use any LiteLLM-supported model + }, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": vector_bucket_name, + "index_name": "test-index", + # dimension is auto-detected from embedding model (text-embedding-3-small = 1536) + "distance_metric": "cosine", + "non_filterable_metadata_keys": ["source_text"], + "aws_region_name": aws_region, + }, + } + + async def query_vector_store( + self, + vector_store_id: str, + query: str, + ) -> Optional[Dict[str, Any]]: + """Query S3 Vectors index.""" + try: + # Import the ingestion class to use its query method + from litellm.rag.ingestion.s3_vectors_ingestion import ( + S3VectorsRAGIngestion, + ) + except ImportError: + pytest.skip("S3 Vectors ingestion not available") + + vector_bucket_name = os.environ.get( + "S3_VECTOR_BUCKET_NAME", "test-litellm-vectors" + ) + aws_region = os.environ.get("AWS_REGION_NAME", "us-west-2") + + # Create ingestion instance to use query method + ingest_options = { + "embedding": {"model": "text-embedding-3-small"}, + "vector_store": { + "custom_llm_provider": "s3_vectors", + "vector_bucket_name": vector_bucket_name, + "aws_region_name": aws_region, + }, + } + + ingestion = S3VectorsRAGIngestion(ingest_options=ingest_options) + + # Query the index + results = await ingestion.query_vector_store( + vector_store_id=vector_store_id, + query=query, + top_k=5, + ) + + return results