From e5ce332db5e966ade329477ffee86e45bfe76433 Mon Sep 17 00:00:00 2001 From: Sindri Date: Thu, 27 Aug 2026 15:03:56 +0000 Subject: [PATCH] feat(bedrock/agentcore): optionally forward client tool declarations in InvokeAgentRuntime payload AgentCore Runtime is schemaless on the agent side: the agent author's @app.entrypoint handler parses whatever JSON arrives. An agent that implements a capability itself (its own search, its own retrieval, its own code runner) has no way today to learn whether the caller offered that capability for this turn. AmazonAgentCoreConfig reports no supported OpenAI params and map_openai_params ignores non_default_params, so `tools` never reaches transform_request. Add a `forward_tools` litellm param, default off. When truthy and the caller sent a non-empty tools list, it is forwarded verbatim under a new top-level `tools` key, alongside the existing `prompt` and optional `content`. Mirrors `forward_multimodal_content` (#28885) in shape, flag parsing and precedence. Reaching the transform still requires `allowed_openai_params: ["tools"]`. That is deliberate: get_supported_openai_params() keeps reporting what AgentCore natively supports, so /model_group/info never advertises tool support for models that do nothing with tools, and forwarding stays an explicit opt-in rather than a capability claim. Nobody who has not set both keys sees any change. The two flag checks fold into one helper. _should_forward_multimodal_content and _should_forward_tools were identical apart from the key they read, so they become _is_flag_enabled(optional_params, litellm_params, key), taking Mapping[str, object] rather than bare dict, which retires an isinstance guard that could never fire. The forwarded list is copied with tuple() rather than list(): still shallow, so the nested tool dicts stay shared and no large payload is cloned, but the value cannot be mutated at all and json.dumps serializes it to the same JSON array. Against the base the file drops from 228 basedpyright diagnostics to 225 and from 42 LIT001 to 40, so it ends up cleaner than it started despite gaining a feature. Note that get_litellm_params only forwards an allowlist of keys (OPTIONAL_KWARGS_KEYS) and `forward_tools` is not on it, so both deployment-level and per-request flags reach transform_request through optional_params via add_provider_specific_params_to_optional_params. litellm_params is only the fallback for callers that build it themselves; an end-to-end test pins the real path. This is a one-way signal. AgentCore responses carry no tool-call channel, so nothing here implies function-calling support. --- .../bedrock/chat/agentcore/transformation.py | 51 ++- .../test_agentcore_transformation.py | 308 ++++++++++++++++++ 2 files changed, 345 insertions(+), 14 deletions(-) diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 4a2db621421..203026d5f77 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -5,7 +5,7 @@ https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgen """ import json -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Mapping from typing import TYPE_CHECKING, Any, Final, Optional, Union, cast from urllib.parse import quote @@ -226,10 +226,19 @@ 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. + - ``tools`` is added ONLY when the ``forward_tools`` litellm param is truthy + AND the caller declared a non-empty OpenAI ``tools`` list, so a schemaless + agent can see what the caller offered this turn. Same opt-in reasoning as + ``content``. AgentCore reports no supported OpenAI params, so ``tools`` is + dropped (or rejected) before this method runs unless the caller also sets + ``allowed_openai_params=["tools"]``: that keeps + ``get_supported_openai_params()`` honest about what AgentCore natively + supports and makes forwarding a deliberate two-step opt-in. One-way signal + only, since AgentCore responses carry no tool-call channel. Returns: - dict: Payload dict containing the prompt and (optionally) the OpenAI - content list. + dict: Payload dict containing the prompt and, optionally, the OpenAI + content list and tool declarations. """ verbose_logger.debug("AgentCore transform_request - optional_params keys: %s", list(optional_params.keys())) @@ -243,7 +252,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # list verbatim under "content" so an attachment-aware agent can read the raw # blocks (image_url, file, etc.). Default off keeps the payload byte-identical # to the legacy {"prompt": "..."} shape for agents that only read the prompt. - if self._should_forward_multimodal_content(optional_params, litellm_params): + if self._is_flag_enabled(optional_params, litellm_params, "forward_multimodal_content"): last_content: Final = messages[-1].get("content") if isinstance(last_content, list) and any( isinstance(block, dict) and block.get("type") not in (None, "text") for block in last_content @@ -252,6 +261,14 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # not deep, to avoid cloning large base64 media on the request path. payload["content"] = list(last_content) + if self._is_flag_enabled(optional_params, litellm_params, "forward_tools"): + tools: Final = optional_params.get("tools") + if isinstance(tools, list) and tools: + # Immutable shallow copy: the payload can never alias the caller's list, + # the nested dicts stay shared to avoid a deep clone, and json.dumps + # serializes a tuple to the same JSON array a list would. + payload["tools"] = tuple(tools) + # Get or generate session ID - this goes in the header runtime_session_id: Final = self._get_runtime_session_id(optional_params) headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] = runtime_session_id @@ -268,19 +285,25 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): return payload @staticmethod - def _should_forward_multimodal_content(optional_params: dict, litellm_params: dict) -> bool: - """Whether to forward raw OpenAI content blocks under ``payload["content"]``. + def _is_flag_enabled( + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + key: str, + ) -> bool: + """Whether the ``key`` payload-forwarding flag is on, ``optional_params`` first. - Opt-in via the ``forward_multimodal_content`` litellm param (default ``False``) - because AgentCore agents must be explicitly written to read the field. The - value may arrive as a bool or a config/env string ("true", "1", ...). Checks - ``optional_params`` first (where other AgentCore params land), then - ``litellm_params``. + These flags are opt-in (default ``False``) because AgentCore agents must be + explicitly written to read the field they add, and the value may arrive as a + bool or as a config/env string ("true", "1", ...). + + Both a deployment-level and a per-request flag reach this method through + ``optional_params``: ``get_litellm_params`` only forwards an allowlist of keys + and these are not on it, so anything unrecognized is routed to the provider + params instead. ``litellm_params`` is the fallback for direct + ``transform_request`` callers. """ for source in (optional_params, litellm_params): - if not isinstance(source, dict): - continue - value = source.get("forward_multimodal_content") + value = source.get(key) if value is None: continue if isinstance(value, str): diff --git a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index ed8aab8d3d0..aa4fee0cd9f 100644 --- a/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/agentcore/test_agentcore_transformation.py @@ -641,3 +641,311 @@ class TestAgentCoreMultimodalContent: payload = config.transform_request(messages=messages, **kwargs) assert payload["content"] == content assert payload["content"] is not content + + +AGENTCORE_TEST_MODEL = ( + "agentcore/arn:aws:bedrock-agentcore:us-west-2:111111111111:runtime/test_agent" +) +WEB_SEARCH_TOOL = { + "type": "function", + "function": { + "name": "web_search", + "description": "Search the web", + "parameters": { + "type": "object", + "properties": {"query": {"type": "string"}}, + }, + }, +} + + +class TestAgentCoreForwardTools: + """Tests for transform_request forwarding the caller's OpenAI tool declarations. + + AgentCore Runtime is schemaless on the agent side — the agent author's + @app.entrypoint handler parses whatever JSON arrives. An agent that implements + a capability itself (its own search, its own retrieval) has no way today to + learn whether the caller offered that capability for this turn: AgentCore + declares no supported OpenAI params, so ``tools`` never reaches the transform. + + When the ``forward_tools`` litellm param is set, the OpenAI tools list is + forwarded verbatim under a "tools" field. This is opt-in: an agent must be + written to read payload["tools"]. Without the flag, the payload is + byte-identical to the legacy {"prompt": "..."} shape. + + This is a one-way signal — AgentCore responses carry no tool-call channel, so + none of this implies function-calling support. + """ + + @pytest.fixture + def config(self): + return AmazonAgentCoreConfig() + + @pytest.fixture + def messages(self): + return [{"role": "user", "content": "what happened today"}] + + @pytest.fixture + def base_kwargs(self): + """Default kwargs — forwarding is OFF (no opt-in flag).""" + return { + "model": "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:111111111111:runtime/test_agent", + "optional_params": {}, + "litellm_params": {}, + "headers": {}, + } + + def test_tools_not_forwarded_by_default(self, config, messages, base_kwargs): + """Default (no opt-in flag): tools are NOT forwarded — backward compat.""" + base_kwargs["optional_params"] = {"tools": [WEB_SEARCH_TOOL]} + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert payload == {"prompt": "what happened today"} + assert "tools" not in payload + + def test_tools_forwarded_when_opted_in(self, config, messages, base_kwargs): + """Flag on + tools present → forwarded verbatim.""" + tools = [WEB_SEARCH_TOOL] + base_kwargs["optional_params"] = {"tools": tools, "forward_tools": True} + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert payload["tools"] == tuple(tools) + assert payload["prompt"] == "what happened today" + + def test_multiple_tools_forwarded_verbatim(self, config, messages, base_kwargs): + """The list is not filtered or reshaped — the agent decides what it wants.""" + calculator = {"type": "function", "function": {"name": "calculator"}} + tools = [WEB_SEARCH_TOOL, calculator] + base_kwargs["optional_params"] = {"tools": tools, "forward_tools": True} + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert payload["tools"] == tuple(tools) + + def test_no_tools_key_when_opted_in_without_tools( + self, config, messages, base_kwargs + ): + """Flag on but the caller declared nothing → no "tools" field.""" + base_kwargs["optional_params"] = {"forward_tools": True} + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert payload == {"prompt": "what happened today"} + + def test_empty_tools_list_not_forwarded(self, config, messages, base_kwargs): + """Clients send tools=[] when a tool toggle is off — treat it as absent.""" + base_kwargs["optional_params"] = {"tools": [], "forward_tools": True} + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert payload == {"prompt": "what happened today"} + assert "tools" not in payload + + def test_non_list_tools_not_forwarded(self, config, messages, base_kwargs): + """A malformed tools value must not reach the payload.""" + base_kwargs["optional_params"] = { + "tools": {"name": "web_search"}, + "forward_tools": True, + } + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert "tools" not in payload + + @pytest.mark.parametrize("flag", ["true", "True", "1", "yes", "on", " TRUE "]) + def test_truthy_string_flags(self, config, messages, base_kwargs, flag): + """Config/env values arrive as strings.""" + base_kwargs["optional_params"] = { + "tools": [WEB_SEARCH_TOOL], + "forward_tools": flag, + } + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert payload["tools"] == (WEB_SEARCH_TOOL,) + + @pytest.mark.parametrize("flag", ["false", "0", "no", "off", ""]) + def test_falsy_string_flags(self, config, messages, base_kwargs, flag): + base_kwargs["optional_params"] = { + "tools": [WEB_SEARCH_TOOL], + "forward_tools": flag, + } + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert "tools" not in payload + + def test_flag_read_from_litellm_params(self, config, messages, base_kwargs): + """litellm_params is the fallback source for direct transform_request callers. + + Deployment-level and per-request flags both arrive in optional_params (see + TestAgentCoreForwardToolsEndToEnd), so this branch only serves callers that + build litellm_params themselves. + """ + base_kwargs["optional_params"] = {"tools": [WEB_SEARCH_TOOL]} + base_kwargs["litellm_params"] = {"forward_tools": True} + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert payload["tools"] == (WEB_SEARCH_TOOL,) + + def test_optional_params_wins_over_litellm_params( + self, config, messages, base_kwargs + ): + """optional_params is read first, so it wins over the litellm_params fallback.""" + base_kwargs["optional_params"] = { + "tools": [WEB_SEARCH_TOOL], + "forward_tools": False, + } + base_kwargs["litellm_params"] = {"forward_tools": True} + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert "tools" not in payload + + def test_payload_does_not_alias_optional_params( + self, config, messages, base_kwargs + ): + """The payload holds an immutable copy, so it cannot alias the caller's list. + + The forwarded value serializes to the same JSON array either way, so the wire + format is asserted here too rather than just the in-memory type. + """ + tools = [WEB_SEARCH_TOOL] + base_kwargs["optional_params"] = {"tools": tools, "forward_tools": True} + + payload = config.transform_request(messages=messages, **base_kwargs) + + assert isinstance(payload["tools"], tuple) + with pytest.raises(AttributeError): + payload["tools"].append({"type": "function", "function": {"name": "extra"}}) + assert tools == [WEB_SEARCH_TOOL] + assert json.loads(json.dumps(payload))["tools"] == [WEB_SEARCH_TOOL] + + def test_tools_and_multimodal_content_are_independent(self, config, base_kwargs): + """Both opt-ins together → both keys present, neither disturbs the other.""" + content = [ + {"type": "text", "text": "what is this"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,aGVsbG8="}, + }, + ] + tools = [WEB_SEARCH_TOOL] + base_kwargs["optional_params"] = { + "tools": tools, + "forward_tools": True, + "forward_multimodal_content": True, + } + + payload = config.transform_request( + messages=[{"role": "user", "content": content}], **base_kwargs + ) + + assert payload["tools"] == tuple(tools) + assert payload["content"] == content + assert payload["prompt"] == "what is this" + + def test_session_header_still_set_when_forwarding( + self, config, messages, base_kwargs + ): + """Forwarding must not disturb the existing header contract.""" + base_kwargs["optional_params"] = { + "tools": [WEB_SEARCH_TOOL], + "forward_tools": True, + "runtimeSessionId": "session-abc", + } + + config.transform_request(messages=messages, **base_kwargs) + + assert ( + base_kwargs["headers"]["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] + == "session-abc" + ) + + +class TestAgentCoreForwardToolsEndToEnd: + """``tools`` only reaches transform_request via allowed_openai_params. + + AgentCore reports no supported OpenAI params, so get_optional_params drops + ``tools`` (or raises without drop_params). ``allowed_openai_params=["tools"]`` + is the documented escape hatch, and is what makes forwarding reachable. + """ + + def test_tools_dropped_without_allowed_openai_params(self): + optional_params = litellm.utils.get_optional_params( + model=AGENTCORE_TEST_MODEL, + custom_llm_provider="bedrock", + tools=[WEB_SEARCH_TOOL], + drop_params=True, + ) + + assert "tools" not in optional_params + + def test_allowed_openai_params_carries_tools_to_the_transform(self): + optional_params = litellm.utils.get_optional_params( + model=AGENTCORE_TEST_MODEL, + custom_llm_provider="bedrock", + tools=[WEB_SEARCH_TOOL], + allowed_openai_params=["tools"], + ) + + assert optional_params["tools"] == [WEB_SEARCH_TOOL] + + payload = AmazonAgentCoreConfig().transform_request( + model=AGENTCORE_TEST_MODEL, + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={"forward_tools": True}, + headers={}, + ) + + assert payload["tools"] == (WEB_SEARCH_TOOL,) + + def test_flag_reaches_the_transform_via_optional_params(self): + """The flag itself is not an OpenAI param, so it lands in optional_params. + + get_litellm_params only forwards an allowlist of keys, which excludes + forward_tools, so this is the path every real caller takes: both the flag and + the tools list arrive in optional_params and litellm_params stays empty. + """ + optional_params = litellm.utils.get_optional_params( + model=AGENTCORE_TEST_MODEL, + custom_llm_provider="bedrock", + tools=[WEB_SEARCH_TOOL], + allowed_openai_params=["tools"], + forward_tools=True, + ) + + assert optional_params["forward_tools"] is True + + payload = AmazonAgentCoreConfig().transform_request( + model=AGENTCORE_TEST_MODEL, + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert payload["tools"] == (WEB_SEARCH_TOOL,) + + def test_allowed_openai_params_alone_does_not_forward(self): + """Both opt-ins are required; the escape hatch alone changes nothing.""" + optional_params = litellm.utils.get_optional_params( + model=AGENTCORE_TEST_MODEL, + custom_llm_provider="bedrock", + tools=[WEB_SEARCH_TOOL], + allowed_openai_params=["tools"], + ) + + payload = AmazonAgentCoreConfig().transform_request( + model=AGENTCORE_TEST_MODEL, + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert payload == {"prompt": "hi"}