mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
fix(bedrock): forward userContext in Knowledge Base Retrieve requests
The Bedrock vector store search only lifted retrievalConfiguration out of extra_body, so the caller's userContext (the Retrieve API's ACL identity) never reached Bedrock and ACL-enabled data sources answered with zero results. The transform now forwards userContext, taken from extra_body first and then from the top-level params where the OpenAI SDK's extra_body merge lands, as the caller sent it.
This commit is contained in:
parent
4e996400e2
commit
033aa8ba6d
4 changed files with 94 additions and 1 deletions
|
|
@ -1,3 +1,4 @@
|
|||
from collections.abc import Mapping
|
||||
from copy import deepcopy
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from urllib.parse import urlparse
|
||||
|
|
@ -14,6 +15,7 @@ from litellm.types.integrations.rag.bedrock_knowledgebase import (
|
|||
BedrockKBResponse,
|
||||
BedrockKBRetrievalConfiguration,
|
||||
BedrockKBRetrievalQuery,
|
||||
BedrockKBUserContext,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.vector_stores import (
|
||||
|
|
@ -242,10 +244,29 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
|
|||
retrieval_config.setdefault("vectorSearchConfiguration", {})["filter"] = filters
|
||||
if retrieval_config:
|
||||
request_body["retrievalConfiguration"] = cast(BedrockKBRetrievalConfiguration, retrieval_config)
|
||||
user_context: Final = self._user_context(extra_body=extra_body, litellm_params=litellm_params)
|
||||
if user_context is not None:
|
||||
request_body["userContext"] = user_context
|
||||
|
||||
litellm_logging_obj.model_call_details["query"] = query
|
||||
return url, request_body
|
||||
|
||||
@staticmethod
|
||||
def _user_context(
|
||||
extra_body: Mapping[str, object] | None, litellm_params: Mapping[str, object]
|
||||
) -> BedrockKBUserContext | None:
|
||||
sources: Final = tuple(source for source in (extra_body, litellm_params) if isinstance(source, Mapping))
|
||||
found: Final = next(
|
||||
(
|
||||
source[key]
|
||||
for source in sources
|
||||
for key in ("userContext", "user_context")
|
||||
if source.get(key) is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
return None if found is None else cast(BedrockKBUserContext, found)
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from typing import Any, Literal
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
|
||||
class BedrockKBLocation(TypedDict, total=False):
|
||||
|
|
@ -127,6 +127,10 @@ class BedrockKBGuardrailConfiguration(TypedDict, total=False):
|
|||
guardrailVersion: str | None
|
||||
|
||||
|
||||
class BedrockKBUserContext(TypedDict):
|
||||
userId: ReadOnly[str]
|
||||
|
||||
|
||||
class BedrockKBRequest(TypedDict, total=False):
|
||||
"""Complete request structure for Bedrock Knowledge Base retrieval."""
|
||||
|
||||
|
|
@ -134,6 +138,7 @@ class BedrockKBRequest(TypedDict, total=False):
|
|||
nextToken: str | None
|
||||
retrievalConfiguration: BedrockKBRetrievalConfiguration | None
|
||||
retrievalQuery: BedrockKBRetrievalQuery
|
||||
userContext: ReadOnly[BedrockKBUserContext | None]
|
||||
|
||||
|
||||
#########################################################################
|
||||
|
|
|
|||
|
|
@ -82,6 +82,7 @@ def test_transform_search_request_uses_only_retrieval_config_from_extra_body():
|
|||
== "HYBRID"
|
||||
)
|
||||
assert "unrelatedField" not in body
|
||||
assert "userContext" not in body
|
||||
|
||||
|
||||
def test_transform_search_request_does_not_mutate_extra_body_and_overrides_number_of_results():
|
||||
|
|
@ -152,3 +153,44 @@ def test_transform_search_request_overrides_filter_without_mutating_extra_body()
|
|||
]["value"]
|
||||
== "a"
|
||||
)
|
||||
|
||||
|
||||
def _search_body(extra_body, litellm_params):
|
||||
config = BedrockVectorStoreConfig()
|
||||
mock_log = MagicMock()
|
||||
mock_log.model_call_details = {}
|
||||
_, body = config.transform_search_vector_store_request(
|
||||
vector_store_id="kb123",
|
||||
query="hello",
|
||||
vector_store_search_optional_params={"max_num_results": 3},
|
||||
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
|
||||
litellm_logging_obj=mock_log,
|
||||
litellm_params=litellm_params,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
return body
|
||||
|
||||
|
||||
def test_transform_search_request_forwards_user_context_from_extra_body():
|
||||
body = _search_body(extra_body={"userContext": {"userId": "alice@example.com"}}, litellm_params={})
|
||||
|
||||
assert body["userContext"] == {"userId": "alice@example.com"}
|
||||
assert body["retrievalConfiguration"] == {"vectorSearchConfiguration": {"numberOfResults": 3}}
|
||||
|
||||
|
||||
def test_transform_search_request_forwards_top_level_user_context_from_litellm_params():
|
||||
body = _search_body(
|
||||
extra_body=None,
|
||||
litellm_params={"vector_store_id": "kb123", "user_context": {"userId": "bob@example.com"}},
|
||||
)
|
||||
|
||||
assert body["userContext"] == {"userId": "bob@example.com"}
|
||||
|
||||
|
||||
def test_transform_search_request_prefers_extra_body_user_context_over_top_level():
|
||||
body = _search_body(
|
||||
extra_body={"userContext": {"userId": "alice@example.com"}},
|
||||
litellm_params={"userContext": {"userId": "bob@example.com"}},
|
||||
)
|
||||
|
||||
assert body["userContext"] == {"userId": "alice@example.com"}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ executor, and it must never leak into litellm_params/kwargs where logging would
|
|||
model_dump() it (the #19550 serialization trap).
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -15,6 +16,7 @@ import litellm.vector_stores.main as vector_stores_main
|
|||
from litellm.llms.base_llm.vector_store.transformation import (
|
||||
RouterVectorStoreEmbeddingExecutor,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.vector_stores.main import search
|
||||
|
||||
MOCK_SEARCH_RESPONSE = {
|
||||
|
|
@ -89,3 +91,26 @@ def test_search_router_not_in_litellm_params():
|
|||
litellm_params = mock_handler.call_args.kwargs["litellm_params"]
|
||||
assert "router" not in litellm_params.model_dump(exclude_none=True)
|
||||
assert getattr(litellm_params, "router", None) is None
|
||||
|
||||
|
||||
def test_search_forwards_top_level_user_context_to_bedrock_retrieve():
|
||||
"""Regression (LIT-4415): a top-level userContext, the shape the OpenAI SDK's extra_body
|
||||
produces on the proxy path, reaches the Bedrock Retrieve request body."""
|
||||
client = MagicMock(spec=HTTPHandler)
|
||||
client.post.return_value = MagicMock(status_code=200, json=MagicMock(return_value={"retrievalResults": []}))
|
||||
|
||||
search(
|
||||
vector_store_id="kb123",
|
||||
query="q",
|
||||
custom_llm_provider="bedrock",
|
||||
aws_region_name="us-west-2",
|
||||
aws_access_key_id="test-key-id",
|
||||
aws_secret_access_key="test-secret-key",
|
||||
userContext={"userId": "alice@example.com"},
|
||||
client=client,
|
||||
litellm_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
posted = json.loads(client.post.call_args.kwargs["data"])
|
||||
assert posted["userContext"] == {"userId": "alice@example.com"}
|
||||
assert posted["retrievalQuery"] == {"text": "q"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue