fix(bedrock): forward userContext in Knowledge Base Retrieve requests

litellm/llms/bedrock/vector_stores/transformation.py

@@ -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 @@
     BedrockKBResponse,
     BedrockKBRetrievalConfiguration,
     BedrockKBRetrievalQuery,
+    BedrockKBUserContext,
 )
 from litellm.types.router import GenericLiteLLMParams
 from litellm.types.vector_stores import (
@@ -242,10 +244,29 @@ def transform_search_vector_store_request(
             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,

litellm/types/integrations/rag/bedrock_knowledgebase.py

@@ -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,13 +127,18 @@ 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."""
 
     guardrailConfiguration: BedrockKBGuardrailConfiguration | None
     nextToken: str | None
     retrievalConfiguration: BedrockKBRetrievalConfiguration | None
     retrievalQuery: BedrockKBRetrievalQuery
+    userContext: ReadOnly[BedrockKBUserContext | None]
 
 
 #########################################################################

tests/test_litellm/llms/bedrock/vector_stores/test_bedrock_vector_store_transformation.py

@@ -1,3 +1,4 @@
+from typing import Final
 from unittest.mock import MagicMock
 
 from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig
@@ -82,6 +83,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 +154,44 @@ def test_transform_search_request_overrides_filter_without_mutating_extra_body()
         ]["value"]
         == "a"
     )
+
+
+def _search_body(extra_body: dict[str, object] | None, litellm_params: dict[str, object]) -> dict[str, object]:
+    config: Final = BedrockVectorStoreConfig()
+    mock_log: Final = 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"}

tests/test_litellm/vector_stores/test_main.py

@@ -7,6 +7,7 @@
 model_dump() it (the #19550 serialization trap).
 """
 
+import json
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -15,6 +16,7 @@
 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"}