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:
mateo-berri 2026-09-16 11:50:36 -07:00
parent 4e996400e2
commit 033aa8ba6d
4 changed files with 94 additions and 1 deletions

View file

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

View file

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

View file

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

View file

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