From cfe4bc678e612bcc37831b69582e886bf05e03da Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 28 Apr 2026 15:37:23 +0530 Subject: [PATCH 1/4] feat(vector-stores): support provider-specific Bedrock retrieval config Route vector store search `extra_body` into provider transformers and handle Bedrock `retrievalConfiguration` explicitly so only intended provider-specific fields are forwarded. Made-with: Cursor --- .../azure_ai/vector_stores/transformation.py | 1 + .../base_llm/vector_store/transformation.py | 3 ++ .../bedrock/vector_stores/transformation.py | 23 +++++++----- litellm/llms/custom_httpx/llm_http_handler.py | 3 ++ .../gemini/vector_stores/transformation.py | 1 + .../milvus/vector_stores/transformation.py | 1 + .../openai/vector_stores/transformation.py | 1 + .../pg_vector/vector_stores/transformation.py | 2 ++ .../ragflow/vector_stores/transformation.py | 1 + .../vector_stores/transformation.py | 2 ++ .../vector_stores/rag_api/transformation.py | 1 + .../search_api/transformation.py | 1 + ...est_bedrock_vector_store_transformation.py | 35 +++++++++++++++++++ 13 files changed, 66 insertions(+), 9 deletions(-) diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index b62acb65166..d2c8206ca9a 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -92,6 +92,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 5fbf0a4b19f..49d2f72db7c 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -56,6 +56,7 @@ class BaseVectorStoreConfig: vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, @@ -67,6 +68,7 @@ class BaseVectorStoreConfig: vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, @@ -81,6 +83,7 @@ class BaseVectorStoreConfig: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, + extra_body=extra_body, api_base=api_base, litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 4da0a7c7791..5b0b64f2429 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from urllib.parse import urlparse import httpx @@ -7,8 +7,8 @@ from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreCon from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.integrations.rag.bedrock_knowledgebase import ( BedrockKBContent, - BedrockKBResponse, BedrockKBRetrievalConfiguration, + BedrockKBResponse, BedrockKBRetrievalQuery, ) from litellm.types.router import GenericLiteLLMParams @@ -199,6 +199,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, @@ -213,6 +214,14 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): } retrieval_config: Dict[str, Any] = {} + from litellm import verbose_logger + + if isinstance(extra_body, dict): + retrieval_config = dict( + extra_body.get("retrievalConfiguration") + or extra_body.get("retrieval_configuration") + or {} + ) max_results = vector_store_search_optional_params.get("max_num_results") if max_results is not None: retrieval_config.setdefault("vectorSearchConfiguration", {})[ @@ -224,13 +233,9 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): "filter" ] = filters if retrieval_config: - # Create a properly typed retrieval configuration - typed_retrieval_config: BedrockKBRetrievalConfiguration = {} - if "vectorSearchConfiguration" in retrieval_config: - typed_retrieval_config["vectorSearchConfiguration"] = retrieval_config[ - "vectorSearchConfiguration" - ] - request_body["retrievalConfiguration"] = typed_retrieval_config + request_body["retrievalConfiguration"] = cast( + BedrockKBRetrievalConfiguration, retrieval_config + ) litellm_logging_obj.model_call_details["query"] = query return url, request_body diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index b9ada079f6b..99f748c0c1a 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -8582,6 +8582,7 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, + extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), @@ -8594,6 +8595,7 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, + extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), @@ -8694,6 +8696,7 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, + extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index e6e8369643e..b6cace066c8 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -115,6 +115,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index fcf5d14db7c..8c08b783387 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -127,6 +127,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index c763ed1c8da..b6eae390d93 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -103,6 +103,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index ba87a8f2b01..8261036cae0 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -77,6 +77,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, @@ -86,6 +87,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, + extra_body=extra_body, api_base=api_base, litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index ed5397eef0c..ae28222c3ce 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -99,6 +99,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 19b59769863..0cf86358873 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -76,6 +76,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, @@ -137,6 +138,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 4baa5774c48..b3fcf4b394c 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -97,6 +97,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 179bd7aeff1..47dac8dca32 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -104,6 +104,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index bcd4c620055..5fa48703b12 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -18,6 +18,7 @@ def test_transform_search_request(): vector_store_id="kb123", query="hello", vector_store_search_optional_params={}, + extra_body=None, api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", litellm_logging_obj=mock_log, litellm_params={}, @@ -25,3 +26,37 @@ def test_transform_search_request(): assert url.endswith("/kb123/retrieve") assert body["retrievalQuery"].get("text") == "hello" + + +def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + + url, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={}, + extra_body={ + "retrievalConfiguration": { + "vectorSearchConfiguration": { + "overrideSearchType": "HYBRID", + "numberOfResults": 8, + } + }, + "unrelatedField": {"should_not": "be_forwarded"}, + }, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + ) + + assert url.endswith("/kb123/retrieve") + assert body["retrievalQuery"].get("text") == "hello" + assert ( + body["retrievalConfiguration"]["vectorSearchConfiguration"][ + "overrideSearchType" + ] + == "HYBRID" + ) + assert "unrelatedField" not in body From 6b86e544e889f6a1aa464b559643e95a26b2e899 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 28 Apr 2026 15:51:48 +0530 Subject: [PATCH 2/4] Fix greptile reviews --- .../bedrock/vector_stores/transformation.py | 24 ++++++- ...est_bedrock_vector_store_transformation.py | 70 +++++++++++++++++++ 2 files changed, 92 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 5b0b64f2429..4e81fa4e66e 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -1,8 +1,10 @@ +from copy import deepcopy from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from urllib.parse import urlparse import httpx +from litellm._logging import verbose_logger from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.types.integrations.rag.bedrock_knowledgebase import ( @@ -214,21 +216,39 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): } retrieval_config: Dict[str, Any] = {} - from litellm import verbose_logger if isinstance(extra_body, dict): - retrieval_config = dict( + retrieval_config = deepcopy( extra_body.get("retrievalConfiguration") or extra_body.get("retrieval_configuration") or {} ) max_results = vector_store_search_optional_params.get("max_num_results") if max_results is not None: + existing_number_of_results = retrieval_config.get( + "vectorSearchConfiguration", {} + ).get("numberOfResults") + if ( + existing_number_of_results is not None + and existing_number_of_results != max_results + ): + verbose_logger.debug( + "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.numberOfResults (%s) with max_num_results=%s", + existing_number_of_results, + max_results, + ) retrieval_config.setdefault("vectorSearchConfiguration", {})[ "numberOfResults" ] = max_results filters = vector_store_search_optional_params.get("filters") if filters is not None: + existing_filter = retrieval_config.get("vectorSearchConfiguration", {}).get( + "filter" + ) + if existing_filter is not None and existing_filter != filters: + verbose_logger.debug( + "Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.filter with filters from vector_store_search_optional_params" + ) retrieval_config.setdefault("vectorSearchConfiguration", {})[ "filter" ] = filters diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index 5fa48703b12..c211a3536e3 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -60,3 +60,73 @@ def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): == "HYBRID" ) assert "unrelatedField" not in body + + +def test_transform_search_request_does_not_mutate_extra_body_and_overrides_number_of_results(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + extra_body = { + "retrievalConfiguration": { + "vectorSearchConfiguration": { + "overrideSearchType": "HYBRID", + "numberOfResults": 8, + } + } + } + + _, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={"max_num_results": 10}, + extra_body=extra_body, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + ) + + assert ( + body["retrievalConfiguration"]["vectorSearchConfiguration"]["numberOfResults"] + == 10 + ) + assert ( + extra_body["retrievalConfiguration"]["vectorSearchConfiguration"][ + "numberOfResults" + ] + == 8 + ) + + +def test_transform_search_request_overrides_filter_without_mutating_extra_body(): + config = BedrockVectorStoreConfig() + mock_log = MagicMock() + mock_log.model_call_details = {} + extra_body = { + "retrievalConfiguration": { + "vectorSearchConfiguration": { + "filter": {"equals": {"key": "tenant", "value": "a"}} + } + } + } + new_filter = {"equals": {"key": "tenant", "value": "b"}} + + _, body = config.transform_search_vector_store_request( + vector_store_id="kb123", + query="hello", + vector_store_search_optional_params={"filters": new_filter}, + extra_body=extra_body, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, + ) + + assert ( + body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"] + == new_filter + ) + assert ( + extra_body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"][ + "equals" + ]["value"] + == "a" + ) From 898040fcdda0f3193862ab495905198e26f5c45f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 28 Apr 2026 22:34:14 +0530 Subject: [PATCH 3/4] Fix tests --- .../s3_vectors/vector_stores/test_s3_vectors_transformation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 7c06823d167..9507ff401aa 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -56,6 +56,7 @@ class TestS3VectorsVectorStoreConfig: vector_store_id="invalid-format", query="test query", vector_store_search_optional_params={}, + extra_body=None, api_base="https://s3vectors.us-west-2.api.aws", litellm_logging_obj=mock_logging_obj, litellm_params={}, From cf74f55b7983e6e0fc58c96d9077803afb37c6ae Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 29 Apr 2026 08:34:31 +0530 Subject: [PATCH 4/4] Fix extra body error --- .../llms/azure_ai/vector_stores/transformation.py | 2 +- litellm/llms/base_llm/vector_store/transformation.py | 6 +++--- litellm/llms/bedrock/vector_stores/transformation.py | 2 +- litellm/llms/custom_httpx/llm_http_handler.py | 6 +++--- litellm/llms/gemini/vector_stores/transformation.py | 2 +- litellm/llms/milvus/vector_stores/transformation.py | 2 +- litellm/llms/openai/vector_stores/transformation.py | 2 +- .../llms/pg_vector/vector_stores/transformation.py | 4 ++-- litellm/llms/ragflow/vector_stores/transformation.py | 2 +- .../llms/s3_vectors/vector_stores/transformation.py | 4 ++-- .../vector_stores/rag_api/transformation.py | 2 +- .../vector_stores/search_api/transformation.py | 2 +- .../test_bedrock_knowledgebase_hook.py | 1 + .../test_bedrock_vector_store_transformation.py | 12 ++++++------ .../vector_stores/test_s3_vectors_transformation.py | 2 +- .../vector_store_tests/test_ragflow_vector_store.py | 1 + 16 files changed, 27 insertions(+), 25 deletions(-) diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index d2c8206ca9a..d1b93c9e7a3 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -92,10 +92,10 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Azure AI Search API diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index 49d2f72db7c..85a9c838264 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -56,10 +56,10 @@ class BaseVectorStoreConfig: vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: pass @@ -68,10 +68,10 @@ class BaseVectorStoreConfig: vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """ Optional async version of transform_search_vector_store_request. @@ -83,10 +83,10 @@ class BaseVectorStoreConfig: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, - extra_body=extra_body, api_base=api_base, litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, + extra_body=extra_body, ) @abstractmethod diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index 4e81fa4e66e..f028503c6a2 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -201,10 +201,10 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: if isinstance(query, list): query = " ".join(query) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 99f748c0c1a..d8515d2c4d6 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -8582,10 +8582,10 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, - extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), + extra_body=extra_body, ) else: ( @@ -8595,10 +8595,10 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, - extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), + extra_body=extra_body, ) all_optional_params: Dict[str, Any] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) @@ -8696,10 +8696,10 @@ class BaseLLMHTTPHandler: vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, - extra_body=extra_body, api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=dict(litellm_params), + extra_body=extra_body, ) all_optional_params: Dict[str, Any] = dict(litellm_params) diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index b6cace066c8..35d83bd2adc 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -115,10 +115,10 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """ Transform search request to Gemini's generateContent format. diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 8c08b783387..af78cd8dbda 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -127,10 +127,10 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Azure AI Search API diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index b6eae390d93..2c11d137480 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -103,10 +103,10 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: url = f"{api_base}/{vector_store_id}/search" typed_request_body = VectorStoreSearchRequest( diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index 8261036cae0..7b22edd8676 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -77,19 +77,19 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: url = f"{api_base}/{vector_store_id}/search" _, request_body = super().transform_search_vector_store_request( vector_store_id=vector_store_id, query=query, vector_store_search_optional_params=vector_store_search_optional_params, - extra_body=extra_body, api_base=api_base, litellm_logging_obj=litellm_logging_obj, litellm_params=litellm_params, + extra_body=extra_body, ) return url, request_body diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index ae28222c3ce..3238d3e9c14 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -99,10 +99,10 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """RAGFlow vector stores are management-only, search is not supported.""" raise NotImplementedError( diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 0cf86358873..8270e99d456 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -76,10 +76,10 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """Sync version - generates embedding synchronously.""" # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name @@ -138,10 +138,10 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict]: """Async version - generates embedding asynchronously.""" # For S3 Vectors, vector_store_id should be in format: bucket_name:index_name diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index b3fcf4b394c..d31e1f6c8f2 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -97,10 +97,10 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Vertex AI RAG API diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index 47dac8dca32..6cb7a86bea2 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -104,10 +104,10 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): vector_store_id: str, query: Union[str, List[str]], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - extra_body: Optional[Dict[str, Any]], api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, + extra_body: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, Any]]: """ Transform search request for Vertex AI RAG API diff --git a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py index abe96b2ea2e..3e8d59b2992 100644 --- a/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py +++ b/tests/logging_callback_tests/test_bedrock_knowledgebase_hook.py @@ -354,6 +354,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters( api_base=api_base, litellm_logging_obj=logging_obj, litellm_params=litellm_params_dict, + extra_body=None, ) ) captured_request_body["url"] = url diff --git a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py index c211a3536e3..d60d0487d06 100644 --- a/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py +++ b/tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py @@ -18,10 +18,10 @@ def test_transform_search_request(): vector_store_id="kb123", query="hello", vector_store_search_optional_params={}, - extra_body=None, api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", litellm_logging_obj=mock_log, litellm_params={}, + extra_body=None, ) assert url.endswith("/kb123/retrieve") @@ -37,6 +37,9 @@ def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): vector_store_id="kb123", query="hello", vector_store_search_optional_params={}, + api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", + litellm_logging_obj=mock_log, + litellm_params={}, extra_body={ "retrievalConfiguration": { "vectorSearchConfiguration": { @@ -46,9 +49,6 @@ def test_transform_search_request_uses_only_retrieval_config_from_extra_body(): }, "unrelatedField": {"should_not": "be_forwarded"}, }, - api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", - litellm_logging_obj=mock_log, - litellm_params={}, ) assert url.endswith("/kb123/retrieve") @@ -79,10 +79,10 @@ def test_transform_search_request_does_not_mutate_extra_body_and_overrides_numbe vector_store_id="kb123", query="hello", vector_store_search_optional_params={"max_num_results": 10}, - extra_body=extra_body, api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", litellm_logging_obj=mock_log, litellm_params={}, + extra_body=extra_body, ) assert ( @@ -114,10 +114,10 @@ def test_transform_search_request_overrides_filter_without_mutating_extra_body() vector_store_id="kb123", query="hello", vector_store_search_optional_params={"filters": new_filter}, - extra_body=extra_body, api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases", litellm_logging_obj=mock_log, litellm_params={}, + extra_body=extra_body, ) assert ( diff --git a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py index 9507ff401aa..7085e45cdc3 100644 --- a/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py +++ b/tests/test_litellm/llms/s3_vectors/vector_stores/test_s3_vectors_transformation.py @@ -56,10 +56,10 @@ class TestS3VectorsVectorStoreConfig: vector_store_id="invalid-format", query="test query", vector_store_search_optional_params={}, - extra_body=None, api_base="https://s3vectors.us-west-2.api.aws", litellm_logging_obj=mock_logging_obj, litellm_params={}, + extra_body=None, ) def test_transform_search_response(self): diff --git a/tests/vector_store_tests/test_ragflow_vector_store.py b/tests/vector_store_tests/test_ragflow_vector_store.py index 4ca38233129..cb4cfd75c1f 100644 --- a/tests/vector_store_tests/test_ragflow_vector_store.py +++ b/tests/vector_store_tests/test_ragflow_vector_store.py @@ -267,6 +267,7 @@ class TestRAGFlowVectorStore(BaseVectorStoreTest): api_base="http://localhost:9380", litellm_logging_obj=logging_obj, litellm_params={}, + extra_body=None, ) def test_transform_search_vector_store_response_not_implemented(self):