Skip to content

Commit 57411fc

Browse files
kshitijk4poorteknium1
authored andcommitted
feat: add BedrockTransport + wire all Bedrock transport paths
Fourth and final transport — completes the transport layer with all four api_modes covered. Wraps agent/bedrock_adapter.py behind the ProviderTransport ABC, handles both raw boto3 dicts and already-normalized SimpleNamespace. Wires all transport methods to production paths in run_agent.py: - build_kwargs: _build_api_kwargs bedrock branch - validate_response: response validation, new bedrock_converse branch - finish_reason: new bedrock_converse branch in finish_reason extraction Based on PR #13467 by @kshitijk4poor, with one adjustment: the main normalize loop does NOT add a bedrock_converse branch to invoke normalize_response on the already-normalized response. Bedrock's normalize_converse_response runs at the dispatch site (run_agent.py:5189), so the response already has the OpenAI-compatible .choices[0].message shape by the time the main loop sees it. Falling through to the chat_completions else branch is correct and sidesteps a redundant NormalizedResponse rebuild. Transport coverage — complete: | api_mode | Transport | build_kwargs | normalize | validate | |--------------------|--------------------------|:------------:|:---------:|:--------:| | anthropic_messages | AnthropicTransport | ✅ | ✅ | ✅ | | codex_responses | ResponsesApiTransport | ✅ | ✅ | ✅ | | chat_completions | ChatCompletionsTransport | ✅ | ✅ | ✅ | | bedrock_converse | BedrockTransport | ✅ | ✅ | ✅ | 17 new BedrockTransport tests pass. 117 transport tests total pass. 160 bedrock/converse tests across tests/agent/ pass. Full tests/run_agent/ targeted suite passes (885/885 + 15 skipped; the 1 remaining failure is the pre-existing test_concurrent_interrupt flake on origin/main).
1 parent 572e27c commit 57411fc

4 files changed

Lines changed: 352 additions & 13 deletions

File tree

agent/transports/__init__.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,3 +45,7 @@ def _discover_transports() -> None:
4545
import agent.transports.chat_completions # noqa: F401
4646
except ImportError:
4747
pass
48+
try:
49+
import agent.transports.bedrock # noqa: F401
50+
except ImportError:
51+
pass

agent/transports/bedrock.py

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
"""AWS Bedrock Converse API transport.
2+
3+
Delegates to the existing adapter functions in agent/bedrock_adapter.py.
4+
Bedrock uses its own boto3 client (not the OpenAI SDK), so the transport
5+
owns format conversion and normalization, while client construction and
6+
boto3 calls stay on AIAgent.
7+
"""
8+
9+
from typing import Any, Dict, List, Optional
10+
11+
from agent.transports.base import ProviderTransport
12+
from agent.transports.types import NormalizedResponse, ToolCall, Usage
13+
14+
15+
class BedrockTransport(ProviderTransport):
16+
"""Transport for api_mode='bedrock_converse'."""
17+
18+
@property
19+
def api_mode(self) -> str:
20+
return "bedrock_converse"
21+
22+
def convert_messages(self, messages: List[Dict[str, Any]], **kwargs) -> Any:
23+
"""Convert OpenAI messages to Bedrock Converse format."""
24+
from agent.bedrock_adapter import convert_messages_to_converse
25+
return convert_messages_to_converse(messages)
26+
27+
def convert_tools(self, tools: List[Dict[str, Any]]) -> Any:
28+
"""Convert OpenAI tool schemas to Bedrock Converse toolConfig."""
29+
from agent.bedrock_adapter import convert_tools_to_converse
30+
return convert_tools_to_converse(tools)
31+
32+
def build_kwargs(
33+
self,
34+
model: str,
35+
messages: List[Dict[str, Any]],
36+
tools: Optional[List[Dict[str, Any]]] = None,
37+
**params,
38+
) -> Dict[str, Any]:
39+
"""Build Bedrock converse() kwargs.
40+
41+
Calls convert_messages and convert_tools internally.
42+
43+
params:
44+
max_tokens: int — output token limit (default 4096)
45+
temperature: float | None
46+
guardrail_config: dict | None — Bedrock guardrails
47+
region: str — AWS region (default 'us-east-1')
48+
"""
49+
from agent.bedrock_adapter import build_converse_kwargs
50+
51+
region = params.get("region", "us-east-1")
52+
guardrail = params.get("guardrail_config")
53+
54+
kwargs = build_converse_kwargs(
55+
model=model,
56+
messages=messages,
57+
tools=tools,
58+
max_tokens=params.get("max_tokens", 4096),
59+
temperature=params.get("temperature"),
60+
guardrail_config=guardrail,
61+
)
62+
# Sentinel keys for dispatch — agent pops these before the boto3 call
63+
kwargs["__bedrock_converse__"] = True
64+
kwargs["__bedrock_region__"] = region
65+
return kwargs
66+
67+
def normalize_response(self, response: Any, **kwargs) -> NormalizedResponse:
68+
"""Normalize Bedrock response to NormalizedResponse.
69+
70+
Handles two shapes:
71+
1. Raw boto3 dict (from direct converse() calls)
72+
2. Already-normalized SimpleNamespace with .choices (from dispatch site)
73+
"""
74+
from agent.bedrock_adapter import normalize_converse_response
75+
76+
# Normalize to OpenAI-compatible SimpleNamespace
77+
if hasattr(response, "choices") and response.choices:
78+
# Already normalized at dispatch site
79+
ns = response
80+
else:
81+
# Raw boto3 dict
82+
ns = normalize_converse_response(response)
83+
84+
choice = ns.choices[0]
85+
msg = choice.message
86+
finish_reason = choice.finish_reason or "stop"
87+
88+
tool_calls = None
89+
if msg.tool_calls:
90+
tool_calls = [
91+
ToolCall(
92+
id=tc.id,
93+
name=tc.function.name,
94+
arguments=tc.function.arguments,
95+
)
96+
for tc in msg.tool_calls
97+
]
98+
99+
usage = None
100+
if hasattr(ns, "usage") and ns.usage:
101+
u = ns.usage
102+
usage = Usage(
103+
prompt_tokens=getattr(u, "prompt_tokens", 0) or 0,
104+
completion_tokens=getattr(u, "completion_tokens", 0) or 0,
105+
total_tokens=getattr(u, "total_tokens", 0) or 0,
106+
)
107+
108+
reasoning = getattr(msg, "reasoning", None) or getattr(msg, "reasoning_content", None)
109+
110+
return NormalizedResponse(
111+
content=msg.content,
112+
tool_calls=tool_calls,
113+
finish_reason=finish_reason,
114+
reasoning=reasoning,
115+
usage=usage,
116+
)
117+
118+
def validate_response(self, response: Any) -> bool:
119+
"""Check Bedrock response structure.
120+
121+
After normalize_converse_response, the response has OpenAI-compatible
122+
.choices — same check as chat_completions.
123+
"""
124+
if response is None:
125+
return False
126+
# Raw Bedrock dict response — check for 'output' key
127+
if isinstance(response, dict):
128+
return "output" in response
129+
# Already-normalized SimpleNamespace
130+
if hasattr(response, "choices"):
131+
return bool(response.choices)
132+
return False
133+
134+
def map_finish_reason(self, raw_reason: str) -> str:
135+
"""Map Bedrock stop reason to OpenAI finish_reason.
136+
137+
The adapter already does this mapping inside normalize_converse_response,
138+
so this is only used for direct access to raw responses.
139+
"""
140+
_MAP = {
141+
"end_turn": "stop",
142+
"tool_use": "tool_calls",
143+
"max_tokens": "length",
144+
"stop_sequence": "stop",
145+
"guardrail_intervened": "content_filter",
146+
"content_filtered": "content_filter",
147+
}
148+
return _MAP.get(raw_reason, "stop")
149+
150+
151+
# Auto-register on import
152+
from agent.transports import register_transport # noqa: E402
153+
154+
register_transport("bedrock_converse", BedrockTransport)

run_agent.py

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6583,6 +6583,15 @@ def _get_chat_completions_transport(self):
65836583
self._chat_completions_transport = t
65846584
return t
65856585

6586+
def _get_bedrock_transport(self):
6587+
"""Return the cached BedrockTransport instance (lazy singleton)."""
6588+
t = getattr(self, "_bedrock_transport", None)
6589+
if t is None:
6590+
from agent.transports import get_transport
6591+
t = get_transport("bedrock_converse")
6592+
self._bedrock_transport = t
6593+
return t
6594+
65866595
def _prepare_anthropic_messages_for_api(self, api_messages: list) -> list:
65876596
if not any(
65886597
isinstance(msg, dict) and self._content_has_image_parts(msg.get("content"))
@@ -6722,21 +6731,17 @@ def _build_api_kwargs(self, api_messages: list) -> dict:
67226731
# AWS Bedrock native Converse API — bypasses the OpenAI client entirely.
67236732
# The adapter handles message/tool conversion and boto3 calls directly.
67246733
if self.api_mode == "bedrock_converse":
6725-
from agent.bedrock_adapter import build_converse_kwargs
6734+
_bt = self._get_bedrock_transport()
67266735
region = getattr(self, "_bedrock_region", None) or "us-east-1"
67276736
guardrail = getattr(self, "_bedrock_guardrail_config", None)
6728-
return {
6729-
"__bedrock_converse__": True,
6730-
"__bedrock_region__": region,
6731-
**build_converse_kwargs(
6732-
model=self.model,
6733-
messages=api_messages,
6734-
tools=self.tools,
6735-
max_tokens=self.max_tokens or 4096,
6736-
temperature=None, # Let the model use its default
6737-
guardrail_config=guardrail,
6738-
),
6739-
}
6737+
return _bt.build_kwargs(
6738+
model=self.model,
6739+
messages=api_messages,
6740+
tools=self.tools,
6741+
max_tokens=self.max_tokens or 4096,
6742+
region=region,
6743+
guardrail_config=guardrail,
6744+
)
67406745

67416746
if self.api_mode == "codex_responses":
67426747
_ct = self._get_codex_transport()
@@ -9250,6 +9255,14 @@ def _stop_spinner():
92509255
error_details.append("response is None")
92519256
else:
92529257
error_details.append("response.content invalid (not a non-empty list)")
9258+
elif self.api_mode == "bedrock_converse":
9259+
_btv = self._get_bedrock_transport()
9260+
if not _btv.validate_response(response):
9261+
response_invalid = True
9262+
if response is None:
9263+
error_details.append("response is None")
9264+
else:
9265+
error_details.append("Bedrock response invalid (no output or choices)")
92539266
else:
92549267
_ctv = self._get_chat_completions_transport()
92559268
if not _ctv.validate_response(response):
@@ -9413,6 +9426,10 @@ def _stop_spinner():
94139426
elif self.api_mode == "anthropic_messages":
94149427
_tfr = self._get_anthropic_transport()
94159428
finish_reason = _tfr.map_finish_reason(response.stop_reason)
9429+
elif self.api_mode == "bedrock_converse":
9430+
# Bedrock response is already normalized at dispatch — finish_reason
9431+
# is already in OpenAI format via normalize_converse_response()
9432+
finish_reason = response.choices[0].finish_reason if hasattr(response, "choices") and response.choices else "stop"
94169433
else:
94179434
finish_reason = response.choices[0].finish_reason
94189435
assistant_message = response.choices[0].message

0 commit comments

Comments
 (0)