scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

feat(vector-stores): support provider-specific Bedrock retrieval config

Browse source 
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
This commit is contained in:
Sameer Kankute 2026-04-28 15:37:23 +05:30
parent 62920a0cb2
commit cfe4bc678e
No known key found for this signature in database
13 changed files with 66 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/llms/azure_ai/vector_stores/transformation.py
View file

@ -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,

3
litellm/llms/base_llm/vector_store/transformation.py
View file

@ -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,

23
litellm/llms/bedrock/vector_stores/transformation.py
View file

@ -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

3
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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),

1
litellm/llms/gemini/vector_stores/transformation.py
View file

@ -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,

1
litellm/llms/milvus/vector_stores/transformation.py
View file

@ -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,

1
litellm/llms/openai/vector_stores/transformation.py
View file

@ -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,

2
litellm/llms/pg_vector/vector_stores/transformation.py
View file

@ -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,

1
litellm/llms/ragflow/vector_stores/transformation.py
View file

@ -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,

2
litellm/llms/s3_vectors/vector_stores/transformation.py
View file

@ -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,

1
litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py
View file

@ -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,

1
litellm/llms/vertex_ai/vector_stores/search_api/transformation.py
View file

@ -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,

35
tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py
View file

@ -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