diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index b9cef2ea301..631e08ec896 100644 --- a/litellm/caching/caching.py +++ b/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]}" diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 2982ae6cd08..c3f1746b841 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/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:"): diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 24602dec677..c149f984a5d 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/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", diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index bd14b505323..8701086a5c2 100644 --- a/tests/unit/caching/test_caching.py +++ b/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"}]} diff --git a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index acfd954215d..03bc4a7e407 100644 --- a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/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 diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 8651d522ac8..5feac4b26a7 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/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