From 145b3ba55151f7767422a04527330c56766ed25f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 27 Jan 2026 15:23:01 -0800 Subject: [PATCH] atransform_search_vector_store_request --- litellm/llms/custom_httpx/llm_http_handler.py | 36 +++++++++++++------ litellm/llms/s3_vectors/__init__.py | 1 + litellm/proxy/rag_endpoints/endpoints.py | 9 +++++ litellm/rag/ingestion/s3_vectors_ingestion.py | 4 ++- litellm/types/router.py | 4 +++ litellm/types/utils.py | 11 +++--- .../test_litellm/llms/s3_vectors/__init__.py | 1 + .../llms/s3_vectors/vector_stores/__init__.py | 1 + 8 files changed, 51 insertions(+), 16 deletions(-) create mode 100644 litellm/llms/s3_vectors/__init__.py create mode 100644 tests/test_litellm/llms/s3_vectors/__init__.py create mode 100644 tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 6a87967c3aa..d2ea7e872a2 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -7033,17 +7033,31 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), ) - ( - url, - request_body, - ) = vector_store_provider_config.transform_search_vector_store_request( - vector_store_id=vector_store_id, - query=query, - vector_store_search_optional_params=vector_store_search_optional_params, - api_base=api_base, - litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), - ) + # Check if provider has async transform method + if hasattr(vector_store_provider_config, "atransform_search_vector_store_request"): + ( + url, + request_body, + ) = await vector_store_provider_config.atransform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + ) + else: + ( + url, + request_body, + ) = vector_store_provider_config.transform_search_vector_store_request( + vector_store_id=vector_store_id, + query=query, + vector_store_search_optional_params=vector_store_search_optional_params, + api_base=api_base, + litellm_logging_obj=logging_obj, + litellm_params=dict(litellm_params), + ) all_optional_params: Dict[str, Any] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) headers, signed_json_body = vector_store_provider_config.sign_request( diff --git a/litellm/llms/s3_vectors/__init__.py b/litellm/llms/s3_vectors/__init__.py new file mode 100644 index 00000000000..e8367949c3e --- /dev/null +++ b/litellm/llms/s3_vectors/__init__.py @@ -0,0 +1 @@ +# S3 Vectors LLM integration diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py index 79f182817f6..2c34457c02e 100644 --- a/litellm/proxy/rag_endpoints/endpoints.py +++ b/litellm/proxy/rag_endpoints/endpoints.py @@ -129,6 +129,14 @@ async def _save_vector_store_to_db_from_rag_ingest( litellm_vector_store_params = ingest_options.get("litellm_vector_store_params", {}) custom_vector_store_name = litellm_vector_store_params.get("vector_store_name") custom_vector_store_description = litellm_vector_store_params.get("vector_store_description") + + # Extract provider-specific params from vector_store_config to save as litellm_params + # This ensures params like aws_region_name, embedding_model, etc. are available for search + provider_specific_params = {} + excluded_keys = {"custom_llm_provider", "vector_store_id"} + for key, value in vector_store_config.items(): + if key not in excluded_keys and value is not None: + provider_specific_params[key] = value # Build file metadata entry using helper file_entry = _build_file_metadata_entry( @@ -167,6 +175,7 @@ async def _save_vector_store_to_db_from_rag_ingest( vector_store_name=vector_store_name, vector_store_description=vector_store_description, vector_store_metadata=initial_metadata, + litellm_params=provider_specific_params if provider_specific_params else None, ) verbose_proxy_logger.info( diff --git a/litellm/rag/ingestion/s3_vectors_ingestion.py b/litellm/rag/ingestion/s3_vectors_ingestion.py index 13964cd89a2..d9abb7aba8d 100644 --- a/litellm/rag/ingestion/s3_vectors_ingestion.py +++ b/litellm/rag/ingestion/s3_vectors_ingestion.py @@ -478,7 +478,9 @@ class S3VectorsRAGIngestion(BaseRAGIngestion, BaseAWSLLM): # Call PutVectors API await self._put_vectors(vectors) - return self.index_name, filename + # Return vector_store_id in format bucket_name:index_name for S3 Vectors search compatibility + vector_store_id = f"{self.vector_bucket_name}:{self.index_name}" + return vector_store_id, filename async def query_vector_store( self, vector_store_id: str, query: str, top_k: int = 5 diff --git a/litellm/types/router.py b/litellm/types/router.py index 8ea7a207535..43943d9e07e 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -404,6 +404,10 @@ class LiteLLMParamsTypedDict(TypedDict, total=False): aws_access_key_id: Optional[str] aws_secret_access_key: Optional[str] aws_region_name: Optional[str] + ## AWS S3 VECTORS ## + vector_bucket_name: Optional[str] + index_name: Optional[str] + embedding_model: Optional[str] ## IBM WATSONX ## watsonx_region_name: Optional[str] ## CUSTOM PRICING ## diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2d67e13e92d..49c903502d1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4,20 +4,22 @@ from enum import Enum from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union from openai._models import BaseModel as OpenAIObject -from openai.types.audio.transcription_create_params import FileTypes as FileTypes # type: ignore +from openai.types.audio.transcription_create_params import ( + FileTypes as FileTypes, # type: ignore +) from openai.types.chat.chat_completion import ChatCompletion as ChatCompletion from openai.types.completion_usage import ( CompletionTokensDetails, CompletionUsage, PromptTokensDetails, ) +from openai.types.moderation import Categories as Categories from openai.types.moderation import ( - Categories as Categories, CategoryAppliedInputTypes as CategoryAppliedInputTypes, - CategoryScores as CategoryScores, ) +from openai.types.moderation import CategoryScores as CategoryScores +from openai.types.moderation_create_response import Moderation as Moderation from openai.types.moderation_create_response import ( - Moderation as Moderation, ModerationCreateResponse as ModerationCreateResponse, ) from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator @@ -3075,6 +3077,7 @@ class LlmProviders(str, Enum): LLAMA = "meta_llama" NSCALE = "nscale" PG_VECTOR = "pg_vector" + S3_VECTORS = "s3_vectors" HELICONE = "helicone" HYPERBOLIC = "hyperbolic" RECRAFT = "recraft" diff --git a/tests/test_litellm/llms/s3_vectors/__init__.py b/tests/test_litellm/llms/s3_vectors/__init__.py new file mode 100644 index 00000000000..d4b0c4d8550 --- /dev/null +++ b/tests/test_litellm/llms/s3_vectors/__init__.py @@ -0,0 +1 @@ +# S3 Vectors tests diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py b/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py new file mode 100644 index 00000000000..231735c1de7 --- /dev/null +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/__init__.py @@ -0,0 +1 @@ +# S3 Vectors vector store tests