Merge 209602939a into 8f19607259

litellm/caching/caching.py

@@ -428,6 +428,10 @@ class Cache:
         combined_kwargs: Final = ModelParamHelper.get_all_llm_api_params()
         is_semantic_cache: Final = self._is_semantic_cache()
         scope_excluded_params: Final = self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS if is_semantic_cache else frozenset()
+        # A deployment that forwards request params into the provider payload makes their values part of the answer
+        keys_provider_params: Final = litellm.enable_caching_on_provider_specific_optional_params is True or bool(
+            kwargs.get("forward_payload_fields")
+        )
         for param in kwargs:
             if param in scope_excluded_params:
                 continue
@@ -438,7 +442,7 @@ class Cache:
                 if param_value is not None:
                     cache_key += f"{param}: {param_value}"
             elif not is_litellm_owned_kwarg(param):
-                if litellm.enable_caching_on_provider_specific_optional_params is True:  # feature flagged for now
+                if keys_provider_params:
                     if kwargs[param] is None:
                         continue  # ignore None params
                     cache_key += f"{param}: {kwargs[param]}"

litellm/llms/bedrock/chat/agentcore/transformation.py

@@ -6,15 +6,17 @@ https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgen
 
 import json
 from collections.abc import AsyncGenerator, Mapping
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Optional, Union
 from urllib.parse import quote
 
 import httpx
-from pydantic import TypeAdapter
+from pydantic import JsonValue, TypeAdapter, ValidationError
 
 from litellm._logging import verbose_logger
 from litellm._uuid import uuid
 from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
+from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     convert_content_list_to_str,
 )
@@ -53,6 +55,17 @@ else:
     AsyncHTTPHandler = Any
 
 _DECODED_JSON: Final = TypeAdapter(object)
+_PARAMS: Final = TypeAdapter(dict[str, object])
+_PAYLOAD_FIELD_NAMES: Final = TypeAdapter(tuple[str, ...])
+_PAYLOAD_FIELD_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
+_UNFORWARDABLE_PAYLOAD_FIELDS: Final = AWS_CREDENTIAL_KWARGS_KEYS | frozenset({"prompt"})
+
+
+def _payload_field_value(name: str, value: object) -> JsonValue:
+    try:
+        return _PAYLOAD_FIELD_VALUE.validate_python(value)
+    except ValidationError:
+        raise BedrockError(status_code=400, message=f"forward_payload_fields: {name} is not a JSON value") from None
 
 
 class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
@@ -245,10 +258,12 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
           OpenAI-shaped multimodal blocks. This is opt-in because an AgentCore agent
           must be explicitly written to read ``payload["content"]``; by default the
           payload stays byte-identical to the legacy ``{"prompt": "..."}`` shape.
+        - Each request param named in the deployment's ``forward_payload_fields``
+          list is copied in under its own name, never ``prompt`` or AWS credentials.
 
         Returns:
-            dict: Payload dict containing the prompt and (optionally) the OpenAI
-            content list.
+            dict: Payload dict containing the prompt, the forwarded params and
+            (optionally) the OpenAI content list.
         """
         verbose_logger.debug("AgentCore transform_request - optional_params keys: %s", list(optional_params.keys()))
 
@@ -256,7 +271,10 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
         prompt: Final = convert_content_list_to_str(messages[-1])
 
         # Create the payload - this is what goes in the body (raw JSON)
-        payload: Final[dict] = {"prompt": prompt}
+        payload: Final[dict] = {
+            "prompt": prompt,
+            **self._forwarded_payload_fields(_PARAMS.validate_python(optional_params)),
+        }
 
         # Opt-in: when forward_multimodal_content is set, forward the OpenAI content
         # list verbatim under "content" so an attachment-aware agent can read the raw
@@ -307,6 +325,17 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
             return bool(value)
         return False
 
+    @staticmethod
+    def _forwarded_payload_fields(optional_params: Mapping[str, object]) -> Mapping[str, JsonValue]:
+        names: Final = _PAYLOAD_FIELD_NAMES.validate_python(optional_params.get("forward_payload_fields") or ())
+        return MappingProxyType(
+            {
+                name: _payload_field_value(name, optional_params[name])
+                for name in names
+                if name in optional_params and name not in _UNFORWARDABLE_PAYLOAD_FIELDS
+            }
+        )
+
     def _extract_sse_json(self, line: str) -> dict | None:
         """Extract and parse JSON from an SSE data line."""
         if not line.startswith("data:"):

litellm/proxy/auth/auth_utils.py

@@ -463,6 +463,8 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
     # Deployment opt-in: a caller-supplied false would switch off identity
     # forwarding and let the caller choose the `user` Fireworks sees.
     "fireworks_forward_user_id",
+    # Deployment opt-in: a caller-supplied list would copy deployment-set params into the AgentCore payload.
+    "forward_payload_fields",
     # SDK-only field; also rejected outright in is_request_body_safe.
     "model_list",
     "vertex_ai_credentials",

tests/unit/caching/test_caching.py

@@ -495,6 +495,20 @@ def test_cache_key_tells_apart_requests_that_differ_only_in_continuation_token()
     assert cache.get_cache_key(**request, continuation_token="") == first_segment_key
 
 
+def test_cache_key_tells_apart_requests_that_forward_different_payload_fields() -> None:
+    cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
+    request: Final = {
+        "model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:111111111111:runtime/docs_agent",
+        "messages": [{"role": "user", "content": "Summarize the document"}],
+        "forward_payload_fields": ["context"],
+    }
+
+    first_document_key: Final = cache.get_cache_key(**request, context={"document": "Q3 revenue grew 4 percent."})
+
+    assert cache.get_cache_key(**request, context={"document": "Q4 revenue fell 2 percent."}) != first_document_key
+    assert cache.get_cache_key(**request, context={"document": "Q3 revenue grew 4 percent."}) == first_document_key
+
+
 def test_cache_key_tells_apart_requests_that_ask_for_token_ids() -> None:
     cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
     request: Final = {"model": "hosted_vllm/qwen3", "messages": [{"role": "user", "content": "count"}]}

tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py

@@ -13,6 +13,7 @@ import json
 import httpx
 import pytest
 import respx
+from pydantic import TypeAdapter
 
 
 from typing import Final
@@ -1365,3 +1366,68 @@ def test_get_complete_url_sends_a_non_string_runtime_endpoint_to_that_endpoint()
     )
 
     assert url.startswith(f"{vpc_endpoint}/runtimes/")
+
+
+def _payload_sent_to_agentcore(respx_mock: respx.Router, **params: object) -> object:
+    route: Final = respx_mock.post(_AGENTCORE_RUNTIME_URL).mock(
+        return_value=httpx.Response(200, json={"result": "done"}, headers={"content-type": "application/json"})
+    )
+    litellm.completion(
+        model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/"
+        "hosted_agent_13sf6-cALnp38iZD",
+        messages=[{"role": "user", "content": "Summarize the document"}],
+        aws_access_key_id="AKIAAGENTCOREKEY",
+        aws_secret_access_key="agentcore-secret",
+        client=HTTPHandler(),
+        **params,
+    )
+    request: Final[httpx.Request] = route.calls.last.request
+    return TypeAdapter(object).validate_json(request.content)
+
+
+DOCUMENT_CONTEXT: Final = {"document": "Q3 revenue grew 4 percent.", "sections": [{"id": 1, "flags": [True, None]}]}
+
+
+@pytest.mark.respx(assert_all_called=True)
+def test_agentcore_payload_holds_only_the_prompt_when_no_fields_are_forwarded(respx_mock: respx.Router) -> None:
+    assert _payload_sent_to_agentcore(respx_mock, context=DOCUMENT_CONTEXT) == {"prompt": "Summarize the document"}
+
+
+@pytest.mark.respx(assert_all_called=True)
+def test_agentcore_payload_carries_each_forwarded_request_param(respx_mock: respx.Router) -> None:
+    payload: Final = _payload_sent_to_agentcore(
+        respx_mock,
+        forward_payload_fields=["context", "trace_label", "not_sent"],
+        context=DOCUMENT_CONTEXT,
+        trace_label="run-7",
+        unlisted_param="stays out",
+    )
+
+    assert payload == {
+        "prompt": "Summarize the document",
+        "context": DOCUMENT_CONTEXT,
+        "trace_label": "run-7",
+    }
+
+
+@pytest.mark.respx(assert_all_called=True)
+def test_agentcore_never_forwards_the_prompt_or_aws_credentials(respx_mock: respx.Router) -> None:
+    payload: Final = _payload_sent_to_agentcore(
+        respx_mock,
+        forward_payload_fields=["prompt", "aws_secret_access_key", "aws_access_key_id", "context"],
+        prompt="replaced prompt",
+        context=DOCUMENT_CONTEXT,
+    )
+
+    assert payload == {"prompt": "Summarize the document", "context": DOCUMENT_CONTEXT}
+
+
+@pytest.mark.respx(assert_all_called=False)
+def test_agentcore_rejects_a_forwarded_param_that_is_not_json(respx_mock: respx.Router) -> None:
+    secret_marker: Final = "do-not-echo-this-value"
+
+    with pytest.raises(litellm.BadRequestError, match="forward_payload_fields: context") as raised:
+        _payload_sent_to_agentcore(respx_mock, forward_payload_fields=["context"], context={secret_marker: {1, 2}})
+
+    assert secret_marker not in str(raised.value)
+    assert not respx_mock.calls

tests/unit/proxy/auth/test_auth_utils.py

@@ -3780,6 +3780,36 @@ class TestIsRequestBodySafeBlocksFireworksForwardUserId:
         )
 
 
+class TestIsRequestBodySafeBlocksForwardPayloadFields:
+    @pytest.mark.parametrize(
+        "body",
+        [
+            pytest.param({"forward_payload_fields": ["qualifier"]}, id="root"),
+            pytest.param({"extra_body": {"forward_payload_fields": ["qualifier"]}}, id="extra_body"),
+            pytest.param({"metadata": {"forward_payload_fields": ["qualifier"]}}, id="metadata"),
+        ],
+    )
+    def test_forward_payload_fields_in_request_body_is_rejected(self, body: dict[str, object]) -> None:
+        with pytest.raises(ValueError, match="forward_payload_fields"):
+            is_request_body_safe(
+                request_body={"model": "agentcore-model", "context": {"document": "note"}, **body},
+                general_settings={},
+                llm_router=None,
+                model="agentcore-model",
+            )
+
+    def test_forwarded_value_without_the_list_is_allowed(self) -> None:
+        assert (
+            is_request_body_safe(
+                request_body={"model": "agentcore-model", "context": {"document": "note"}},
+                general_settings={},
+                llm_router=None,
+                model="agentcore-model",
+            )
+            is True
+        )
+
+
 class TestIsRequestBodySafeBlocksRustOptIn:
     """``rust`` hands the whole call to the Rust core, which signs and sends
     with its own HTTP client rather than the one the deployment configured, and