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 178

Fix greptile reviews

Browse source
This commit is contained in:
Sameer Kankute 2026-04-28 15:51:48 +05:30
parent cfe4bc678e
commit 6b86e544e8
No known key found for this signature in database
2 changed files with 92 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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