diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index 258ca7048e7..dace140d8dc 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -64,11 +64,13 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Extract config self.vector_bucket_name = self.vector_store_config["vector_bucket_name"] self.index_name = self.vector_store_config.get("index_name") - self.dimension = int(self.vector_store_config.get("dimension", 1024)) self.distance_metric = self.vector_store_config.get("distance_metric", "cosine") self.non_filterable_metadata_keys = self.vector_store_config.get( "non_filterable_metadata_keys", ["source_text"] ) + + # 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") @@ -88,11 +90,63 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # 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 1024 + + 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}. " + "Using default dimension of 1024." + ) + + return 1024 + + 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() diff --git a/litellm/types/rag.py b/litellm/types/rag.py index e0029ad37be..cae07708686 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -153,7 +153,7 @@ class S3VectorsVectorStoreOptions(TypedDict, total=False): index_name: Optional[str] # Vector index name (auto-creates if not provided) # Index configuration (for auto-creation) - dimension: Optional[int] # Vector dimension (default: 1024) + 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"])