Xiaochuang Yuan commited on
Commit
075d71c
ยท
1 Parent(s): 467b8fe

providers: implement BedrockProvider via Converse API

Browse files

The BedrockProvider stub raised NotImplementedError. This commit
implements it as a thin translation layer between the agent's
OpenAI-shape messages + tool schemas and AWS Bedrock Converse, and
back to the same ChatReply the OpenAI path returns โ€” so the agent
and FullPlayback see identical shapes regardless of provider.

* OpenAI system / user / assistant / tool messages map to Bedrock's
top-level system blocks plus alternating user/assistant turns with
toolUse / toolResult content blocks. Adjacent same-role messages
collapse (Bedrock requires strict alternation).
* Multimodal user content (data:image/png;base64) lifts to Bedrock
{image: {format, source: {bytes}}} blocks.
* OpenAI tool schemas translate to toolConfig.tools[].toolSpec; an
empty parameter object is backfilled to type=object.
* Bedrock toolUse / text content blocks parse back to ChatReply
text + tool_calls; reasoningContent (extended thinking) maps onto
ChatReply.reasoning.
* Auth flows through the standard boto3 credential chain โ€” never
hardcoded. boto3 is a soft dep (lazy import).
* ProviderConfig gains bedrock_region (default us-west-2). The
model id is the inference profile id
(us.anthropic.claude-sonnet-4-6); the on-demand model id returns
ValidationException.

Tests in tests/test_providers_bedrock.py exercise every translation
case (system lift, multimodal, toolUse, toolResult, alternation
merge, schema backfill, response parsing, end-to-end stub round-
trip). The existing OpenAI wire tests still pass.

openra_bench/providers.py CHANGED
@@ -6,8 +6,11 @@ Adapters:
6
  * `OpenAICompatibleProvider` โ€” OpenAI Chat Completions wire format. Covers
7
  local **vLLM** (matches Training's rollout path) and **OpenRouter**
8
  (the Phase-0 test target) by base_url alone.
9
- * `BedrockProvider` โ€” AWS Bedrock Converse. Stubbed with a precise
10
- NotImplementedError so the wiring exists before the dependency does.
 
 
 
11
 
12
  Selection is pure config (`ProviderConfig`); no provider-specific code
13
  leaks into the agent.
@@ -47,6 +50,14 @@ _PRESETS: dict[str, dict[str, str]] = {
47
  "base_url": "https://api.together.xyz/v1",
48
  "api_key_env": "TOGETHER_API_KEY",
49
  },
 
 
 
 
 
 
 
 
50
  }
51
 
52
 
@@ -88,6 +99,11 @@ class ProviderConfig:
88
  max_history_turns: int = 16 # sliding wire-history window (0=unbounded)
89
  price_in_per_m: float = 0.0 # USD / 1M prompt tokens
90
  price_out_per_m: float = 0.0 # USD / 1M completion tokens
 
 
 
 
 
91
 
92
  def resolved_base_url(self) -> str:
93
  if self.base_url:
@@ -443,23 +459,376 @@ class OpenAICompatibleProvider(ChatProvider):
443
 
444
 
445
  class BedrockProvider(ChatProvider):
446
- """AWS Bedrock Converse. Wired but not yet implemented."""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
447
 
448
- def __init__(self, cfg: ProviderConfig):
 
449
  self.cfg = cfg
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
450
 
451
  def complete(self, messages: list[dict], tools: list[dict]) -> ChatReply:
452
- raise NotImplementedError(
453
- "BedrockProvider not implemented yet. Use provider='openrouter' "
454
- "or 'vllm' for Phase 0; Bedrock Converse adapter is a tracked "
455
- "follow-up (needs boto3 + message/tool shape translation)."
 
 
 
 
 
 
 
 
 
 
 
 
456
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
457
 
458
 
459
  def make_provider(cfg: ProviderConfig, *, rate_limiter=None,
460
  cost_meter=None) -> ChatProvider:
461
  if cfg.provider == "bedrock":
462
- return BedrockProvider(cfg)
 
 
463
  if cfg.provider in ("openai", "vllm", "openrouter", "together"):
464
  # together.ai's newer Qwen3.x and Llama-3.x families gate on
465
  # streaming (`streaming_required` 400 in non-stream mode); flip
 
6
  * `OpenAICompatibleProvider` โ€” OpenAI Chat Completions wire format. Covers
7
  local **vLLM** (matches Training's rollout path) and **OpenRouter**
8
  (the Phase-0 test target) by base_url alone.
9
+ * `BedrockProvider` โ€” AWS Bedrock Converse. Translates the agent's
10
+ OpenAI-shape messages + tool schemas to Bedrock Converse and back to
11
+ the same `ChatReply` the OpenAI path returns, so the agent stays
12
+ provider-agnostic. Auth comes from the AWS credential chain (env /
13
+ shared config / role) โ€” never hardcoded.
14
 
15
  Selection is pure config (`ProviderConfig`); no provider-specific code
16
  leaks into the agent.
 
50
  "base_url": "https://api.together.xyz/v1",
51
  "api_key_env": "TOGETHER_API_KEY",
52
  },
53
+ # AWS Bedrock โ€” auth via the boto3 credential chain (env, shared
54
+ # config, instance/role). `base_url` is unused (the SDK derives the
55
+ # endpoint from the region). `api_key_env` is unused (left for
56
+ # interface parity); `bedrock_region` on ProviderConfig wins.
57
+ "bedrock": {
58
+ "base_url": "",
59
+ "api_key_env": "",
60
+ },
61
  }
62
 
63
 
 
99
  max_history_turns: int = 16 # sliding wire-history window (0=unbounded)
100
  price_in_per_m: float = 0.0 # USD / 1M prompt tokens
101
  price_out_per_m: float = 0.0 # USD / 1M completion tokens
102
+ # AWS Bedrock: inference region. Sonnet 4.6 is exposed via the
103
+ # `us.anthropic.claude-sonnet-4-6` cross-region inference profile,
104
+ # which routes from `us-west-2` (the on-demand model id returns
105
+ # ValidationException โ€” only the inference profile is callable).
106
+ bedrock_region: str = "us-west-2"
107
 
108
  def resolved_base_url(self) -> str:
109
  if self.base_url:
 
459
 
460
 
461
  class BedrockProvider(ChatProvider):
462
+ """AWS Bedrock Converse adapter.
463
+
464
+ Translates between the agent's OpenAI-shape messages + tool
465
+ schemas and the Bedrock Converse wire format, and translates the
466
+ response back to a `ChatReply` so the agent and FullPlayback see
467
+ the SAME shape they get from the OpenAI-compatible path. Auth
468
+ flows through boto3's standard credential chain โ€” env vars, the
469
+ shared config file, IAM role, etc. The model id is the inference
470
+ profile id (`us.anthropic.claude-sonnet-4-6`), not the on-demand
471
+ model id (which returns ValidationException).
472
+
473
+ Wire-shape mapping:
474
+ * OpenAI `system` messages โ†’ top-level `system: [{text}]`
475
+ * OpenAI text user/assistant โ†’ `content: [{text}]`
476
+ * OpenAI multimodal user content โ†’ `content: [{text}, {image}]`
477
+ * OpenAI assistant `tool_calls` โ†’ `content: [{toolUse}]`
478
+ * OpenAI `tool` reply โ†’ user `[{toolResult}]`
479
+ * OpenAI `tools` (JSON-Schema) โ†’ `toolConfig: {tools: [{toolSpec}]}`
480
+ * Bedrock `output.message.content` โ†’ ChatReply.text + tool_calls
481
+ * Bedrock `usage.{input,output}Tokens` โ†’ usage.{prompt,completion}_tokens
482
+
483
+ Tool-call ids: Bedrock requires a `toolUseId` on every assistant
484
+ `toolUse` and the matching user `toolResult`. The bench agent
485
+ canonicalises these as `c0/c1/...` per turn, so the translation
486
+ passes them straight through.
487
+ """
488
 
489
+ def __init__(self, cfg: ProviderConfig, *, rate_limiter=None,
490
+ cost_meter=None, client=None):
491
  self.cfg = cfg
492
+ self.model_id = cfg.model
493
+ from .resilience import CostMeter, RateLimiter, RetryPolicy
494
+
495
+ self._rl = rate_limiter or RateLimiter(cfg.qps)
496
+ self._cost = cost_meter or CostMeter(
497
+ cfg.price_in_per_m, cfg.price_out_per_m
498
+ )
499
+ self._policy = RetryPolicy(
500
+ max_attempts=max(1, cfg.max_retries),
501
+ base=cfg.retry_base_s,
502
+ cap=cfg.retry_cap_s,
503
+ )
504
+ # Lazy import: keep boto3 a soft dep โ€” only providers='bedrock'
505
+ # forces the dependency, never the OpenRouter / vLLM paths.
506
+ if client is not None:
507
+ self._client = client
508
+ else:
509
+ try:
510
+ import boto3
511
+ except ImportError as e: # pragma: no cover โ€” env-dep
512
+ raise RuntimeError(
513
+ "BedrockProvider needs boto3. Install with "
514
+ "`pip install boto3`."
515
+ ) from e
516
+ self._client = boto3.client(
517
+ "bedrock-runtime", region_name=cfg.bedrock_region
518
+ )
519
+ # Audit hook (parallels OpenAICompatibleProvider): when set to
520
+ # a list, every successful complete() appends a record so
521
+ # FullPlayback can capture literal request + raw response.
522
+ self.request_log: list[dict] | None = None
523
+
524
+ @property
525
+ def cost_meter(self):
526
+ return self._cost
527
+
528
+ # โ”€โ”€ Wire translation: OpenAI โ†’ Bedrock โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
529
+
530
+ @staticmethod
531
+ def _to_bedrock_messages(messages: list[dict]) -> tuple[list[dict], list[dict]]:
532
+ """Pure: split OpenAI messages into (system, conversation).
533
+
534
+ System messages are concatenated into a list of `{text}` blocks
535
+ for Bedrock's top-level `system` parameter. Tool replies
536
+ (`role=tool`) become user-role `toolResult` content blocks; an
537
+ assistant message with `tool_calls` becomes Bedrock `toolUse`
538
+ content blocks (text content, if any, is preserved alongside).
539
+ Adjacent same-role messages are merged because Bedrock REQUIRES
540
+ strictly alternating user/assistant turns โ€” a `tool` reply
541
+ followed by another user briefing must collapse into ONE
542
+ Bedrock user message with multiple content blocks.
543
+ """
544
+ sys_blocks: list[dict] = []
545
+ out: list[dict] = []
546
+ for m in messages:
547
+ role = m.get("role")
548
+ if role == "system":
549
+ txt = m.get("content")
550
+ if isinstance(txt, list):
551
+ txt = "\n".join(
552
+ p.get("text", "") for p in txt
553
+ if isinstance(p, dict) and p.get("type") == "text"
554
+ )
555
+ if txt:
556
+ sys_blocks.append({"text": str(txt)})
557
+ continue
558
+ blocks = BedrockProvider._content_to_blocks(m)
559
+ if not blocks:
560
+ continue
561
+ br_role = "user" if role in ("user", "tool") else "assistant"
562
+ if out and out[-1]["role"] == br_role:
563
+ out[-1]["content"].extend(blocks)
564
+ else:
565
+ out.append({"role": br_role, "content": blocks})
566
+ return sys_blocks, out
567
+
568
+ @staticmethod
569
+ def _content_to_blocks(msg: dict) -> list[dict]:
570
+ """Pure: OpenAI message โ†’ list of Bedrock content blocks."""
571
+ role = msg.get("role")
572
+ # Tool-result reply โ†’ toolResult block.
573
+ if role == "tool":
574
+ tcid = msg.get("tool_call_id") or ""
575
+ content = msg.get("content")
576
+ if isinstance(content, list):
577
+ content = " ".join(
578
+ p.get("text", "") for p in content
579
+ if isinstance(p, dict) and p.get("type") == "text"
580
+ )
581
+ return [{
582
+ "toolResult": {
583
+ "toolUseId": str(tcid),
584
+ "content": [{"text": str(content) if content else "ok"}],
585
+ }
586
+ }]
587
+ blocks: list[dict] = []
588
+ c = msg.get("content")
589
+ if isinstance(c, str):
590
+ if c:
591
+ blocks.append({"text": c})
592
+ elif isinstance(c, list):
593
+ for part in c:
594
+ if not isinstance(part, dict):
595
+ continue
596
+ t = part.get("type")
597
+ if t == "text":
598
+ txt = part.get("text", "")
599
+ if txt:
600
+ blocks.append({"text": txt})
601
+ elif t == "image_url":
602
+ iu = part.get("image_url") or {}
603
+ url = iu.get("url", "") if isinstance(iu, dict) else ""
604
+ img = BedrockProvider._image_block_from_data_url(url)
605
+ if img is not None:
606
+ blocks.append(img)
607
+ # Assistant tool_calls โ†’ toolUse blocks (after any text).
608
+ for tc in msg.get("tool_calls") or []:
609
+ fn = tc.get("function") or {}
610
+ args = fn.get("arguments", {})
611
+ if isinstance(args, str):
612
+ try:
613
+ args = json.loads(args or "{}")
614
+ except json.JSONDecodeError:
615
+ args = {}
616
+ if not isinstance(args, dict):
617
+ args = {}
618
+ blocks.append({
619
+ "toolUse": {
620
+ "toolUseId": str(tc.get("id") or ""),
621
+ "name": fn.get("name", ""),
622
+ "input": args,
623
+ }
624
+ })
625
+ return blocks
626
+
627
+ @staticmethod
628
+ def _image_block_from_data_url(url: str) -> dict | None:
629
+ """Pure: turn a `data:image/png;base64,...` URL into a Bedrock
630
+ `{image: {format, source: {bytes}}}` block. Bedrock accepts
631
+ png / jpeg / gif / webp; the bench only emits png minimaps."""
632
+ import base64
633
+
634
+ if not url.startswith("data:"):
635
+ return None
636
+ try:
637
+ header, b64 = url.split(",", 1)
638
+ except ValueError:
639
+ return None
640
+ fmt = "png"
641
+ if "image/" in header:
642
+ mt = header.split("image/", 1)[1].split(";", 1)[0].lower()
643
+ if mt in ("png", "jpeg", "jpg", "gif", "webp"):
644
+ fmt = "jpeg" if mt == "jpg" else mt
645
+ try:
646
+ raw = base64.b64decode(b64)
647
+ except (ValueError, TypeError):
648
+ return None
649
+ return {"image": {"format": fmt, "source": {"bytes": raw}}}
650
+
651
+ @staticmethod
652
+ def _to_bedrock_tools(tools: list[dict]) -> dict | None:
653
+ """Pure: OpenAI tool list โ†’ Bedrock `toolConfig`. The OpenAI
654
+ schema is `{type: "function", function: {name, description,
655
+ parameters}}`; Bedrock wants `{toolSpec: {name, description,
656
+ inputSchema: {json: <parameters>}}}`. Bedrock additionally
657
+ requires `inputSchema.json.type` (some agents emit empty
658
+ params) โ€” we backfill an empty object schema."""
659
+ if not tools:
660
+ return None
661
+ specs = []
662
+ for t in tools:
663
+ fn = t.get("function") or {}
664
+ params = fn.get("parameters") or {"type": "object", "properties": {}}
665
+ if "type" not in params:
666
+ params = {"type": "object", **params}
667
+ specs.append({
668
+ "toolSpec": {
669
+ "name": fn.get("name", ""),
670
+ "description": fn.get("description", ""),
671
+ "inputSchema": {"json": params},
672
+ }
673
+ })
674
+ return {"tools": specs}
675
+
676
+ # โ”€โ”€ Wire translation: Bedrock โ†’ ChatReply โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
677
+
678
+ @staticmethod
679
+ def _reply_from_bedrock(resp: dict) -> ChatReply:
680
+ """Pure: parse a Bedrock Converse response into a ChatReply.
681
+
682
+ Bedrock emits one assistant message; its content blocks are
683
+ either `{text}` (plain reply) or `{toolUse}` (a function call).
684
+ We concatenate text blocks and lift toolUse blocks into the
685
+ same `[{name, arguments}]` list the OpenAI parser produces."""
686
+ msg = (resp.get("output") or {}).get("message") or {}
687
+ content_blocks = msg.get("content") or []
688
+ text_parts: list[str] = []
689
+ calls: list[dict] = []
690
+ reasoning_parts: list[str] = []
691
+ for blk in content_blocks:
692
+ if not isinstance(blk, dict):
693
+ continue
694
+ if "text" in blk:
695
+ text_parts.append(blk["text"])
696
+ elif "toolUse" in blk:
697
+ tu = blk["toolUse"]
698
+ calls.append({
699
+ "name": tu.get("name", ""),
700
+ "arguments": tu.get("input") or {},
701
+ })
702
+ elif "reasoningContent" in blk:
703
+ # Bedrock surfaces extended thinking under
704
+ # reasoningContent.{reasoningText: {text}} โ€” preserve
705
+ # it on the reply for FullPlayback.
706
+ rc = blk["reasoningContent"] or {}
707
+ rt = rc.get("reasoningText") or {}
708
+ t = rt.get("text") if isinstance(rt, dict) else None
709
+ if t:
710
+ reasoning_parts.append(str(t))
711
+ usage = resp.get("usage") or {}
712
+ return ChatReply(
713
+ text="".join(text_parts),
714
+ tool_calls=calls,
715
+ reasoning="".join(reasoning_parts),
716
+ usage={
717
+ "prompt_tokens": int(usage.get("inputTokens", 0) or 0),
718
+ "completion_tokens": int(usage.get("outputTokens", 0) or 0),
719
+ },
720
+ raw=resp,
721
+ )
722
+
723
+ # โ”€โ”€ Public API โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
724
+
725
+ def _converse_once(self, system_blocks, br_messages, tool_config,
726
+ inference_cfg) -> dict:
727
+ from .resilience import FatalProviderError
728
+ try:
729
+ kwargs = {
730
+ "modelId": self.model_id,
731
+ "messages": br_messages,
732
+ "inferenceConfig": inference_cfg,
733
+ }
734
+ if system_blocks:
735
+ kwargs["system"] = system_blocks
736
+ if tool_config:
737
+ kwargs["toolConfig"] = tool_config
738
+ return self._client.converse(**kwargs)
739
+ except Exception as e: # noqa: BLE001
740
+ # Boto raises ClientError with a `response[Error][Code]`.
741
+ code = ""
742
+ status = 0
743
+ try:
744
+ err = getattr(e, "response", {}) or {}
745
+ meta = err.get("ResponseMetadata") or {}
746
+ status = int(meta.get("HTTPStatusCode", 0) or 0)
747
+ code = (err.get("Error") or {}).get("Code", "")
748
+ except Exception: # noqa: BLE001
749
+ pass
750
+ transient = status in (408, 425, 429, 500, 502, 503, 504) or code in (
751
+ "ThrottlingException",
752
+ "ServiceUnavailableException",
753
+ "ModelTimeoutException",
754
+ "InternalServerException",
755
+ "ModelStreamErrorException",
756
+ )
757
+ cls = RuntimeError if transient else FatalProviderError
758
+ new = cls(f"bedrock {code or status or 'error'}: {e}")
759
+ new.transient = transient # type: ignore[attr-defined]
760
+ new.retry_after = None # type: ignore[attr-defined]
761
+ raise new from e
762
 
763
  def complete(self, messages: list[dict], tools: list[dict]) -> ChatReply:
764
+ from .resilience import retry_call
765
+
766
+ cfg = self.cfg
767
+ sys_blocks, br_messages = self._to_bedrock_messages(messages)
768
+ tool_config = self._to_bedrock_tools(tools)
769
+ inference_cfg = {
770
+ "temperature": cfg.temperature,
771
+ "maxTokens": cfg.max_tokens,
772
+ }
773
+
774
+ self._rl.acquire()
775
+ resp = retry_call(
776
+ lambda: self._converse_once(
777
+ sys_blocks, br_messages, tool_config, inference_cfg,
778
+ ),
779
+ self._policy,
780
  )
781
+ reply = self._reply_from_bedrock(resp)
782
+ u = reply.usage or {}
783
+ self._cost.add(u.get("prompt_tokens", 0), u.get("completion_tokens", 0))
784
+ self._cost.check()
785
+ if self.request_log is not None:
786
+ try:
787
+ # Redact image bytes from the request log (they're
788
+ # huge, and duplicated per turn). Replace with a
789
+ # short placeholder; the rest of the body is small.
790
+ def _redact(b):
791
+ if isinstance(b, dict):
792
+ return {k: _redact(v) for k, v in b.items()}
793
+ if isinstance(b, list):
794
+ return [_redact(x) for x in b]
795
+ if isinstance(b, (bytes, bytearray)):
796
+ return f"<bytes:{len(b)}>"
797
+ return b
798
+
799
+ self.request_log.append({
800
+ "request": {
801
+ "model": self.model_id,
802
+ "system": _redact(sys_blocks),
803
+ "messages": _redact(br_messages),
804
+ "toolConfig": tool_config,
805
+ "inferenceConfig": inference_cfg,
806
+ },
807
+ "response": {
808
+ "raw": _redact(reply.raw),
809
+ "text": reply.text,
810
+ "tool_calls": reply.tool_calls,
811
+ "reasoning": reply.reasoning,
812
+ "usage": dict(reply.usage or {}),
813
+ "finish_reason": resp.get("stopReason"),
814
+ },
815
+ })
816
+ except Exception: # noqa: BLE001 โ€” audit must never break a run
817
+ pass
818
+ return reply
819
+
820
+ def close(self) -> None: # noqa: D401 โ€” interface parity
821
+ # boto3 clients don't need explicit close; provided for
822
+ # symmetry with OpenAICompatibleProvider.
823
+ pass
824
 
825
 
826
  def make_provider(cfg: ProviderConfig, *, rate_limiter=None,
827
  cost_meter=None) -> ChatProvider:
828
  if cfg.provider == "bedrock":
829
+ return BedrockProvider(
830
+ cfg, rate_limiter=rate_limiter, cost_meter=cost_meter,
831
+ )
832
  if cfg.provider in ("openai", "vllm", "openrouter", "together"):
833
  # together.ai's newer Qwen3.x and Llama-3.x families gate on
834
  # streaming (`streaming_required` 400 in non-stream mode); flip
tests/test_providers_bedrock.py ADDED
@@ -0,0 +1,259 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Bedrock Converse adapter โ€” wire-shape translation in both directions.
2
+
3
+ These tests exercise the pure translation helpers (no AWS, no
4
+ network). The end-to-end smoke test against a real `us-west-2`
5
+ inference profile lives in `docs/BEDROCK_SMOKE.md` โ€” runnable but
6
+ not part of CI (avoids non-zero AWS charges in the default suite).
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ import base64
12
+
13
+ from openra_bench.providers import (
14
+ BedrockProvider,
15
+ ChatReply,
16
+ ProviderConfig,
17
+ )
18
+
19
+
20
+ def _png_data_url() -> str:
21
+ """A 1x1 transparent PNG, sufficient to exercise the data-url path."""
22
+ raw = base64.b64decode(
23
+ "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAj"
24
+ "CB0C8AAAAASUVORK5CYII="
25
+ )
26
+ b64 = base64.b64encode(raw).decode("ascii")
27
+ return f"data:image/png;base64,{b64}"
28
+
29
+
30
+ # โ”€โ”€ Outbound: OpenAI messages โ†’ Bedrock Converse โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
31
+
32
+
33
+ def test_system_messages_lift_to_top_level():
34
+ sys, conv = BedrockProvider._to_bedrock_messages([
35
+ {"role": "system", "content": "you are a commander"},
36
+ {"role": "user", "content": "hi"},
37
+ ])
38
+ assert sys == [{"text": "you are a commander"}]
39
+ assert conv == [{"role": "user", "content": [{"text": "hi"}]}]
40
+
41
+
42
+ def test_multimodal_user_message_lifts_image_block():
43
+ msg = {
44
+ "role": "user",
45
+ "content": [
46
+ {"type": "text", "text": "MAP TURN 1"},
47
+ {"type": "image_url",
48
+ "image_url": {"url": _png_data_url()}},
49
+ ],
50
+ }
51
+ _, conv = BedrockProvider._to_bedrock_messages([msg])
52
+ assert len(conv) == 1
53
+ blocks = conv[0]["content"]
54
+ assert blocks[0] == {"text": "MAP TURN 1"}
55
+ assert "image" in blocks[1]
56
+ assert blocks[1]["image"]["format"] == "png"
57
+ assert isinstance(blocks[1]["image"]["source"]["bytes"], (bytes, bytearray))
58
+
59
+
60
+ def test_assistant_tool_calls_become_toolUse_blocks():
61
+ msg = {
62
+ "role": "assistant",
63
+ "content": "moving",
64
+ "tool_calls": [{
65
+ "id": "c0", "type": "function",
66
+ "function": {
67
+ "name": "move_units",
68
+ "arguments": {"unit_ids": [1004], "target_x": 50, "target_y": 50},
69
+ },
70
+ }],
71
+ }
72
+ _, conv = BedrockProvider._to_bedrock_messages([msg])
73
+ blocks = conv[0]["content"]
74
+ assert conv[0]["role"] == "assistant"
75
+ assert blocks[0] == {"text": "moving"}
76
+ tu = blocks[1]["toolUse"]
77
+ assert tu["toolUseId"] == "c0"
78
+ assert tu["name"] == "move_units"
79
+ assert tu["input"] == {"unit_ids": [1004], "target_x": 50, "target_y": 50}
80
+
81
+
82
+ def test_assistant_tool_calls_with_string_arguments_decoded():
83
+ """The OpenAI wire spec stores `arguments` as a JSON STRING; the
84
+ Bedrock toolUse `input` MUST be a dict."""
85
+ msg = {
86
+ "role": "assistant",
87
+ "content": "",
88
+ "tool_calls": [{
89
+ "id": "c0", "type": "function",
90
+ "function": {"name": "observe", "arguments": '{}'},
91
+ }],
92
+ }
93
+ _, conv = BedrockProvider._to_bedrock_messages([msg])
94
+ tu = conv[0]["content"][0]["toolUse"]
95
+ assert tu["input"] == {}
96
+
97
+
98
+ def test_tool_reply_becomes_user_toolResult_block():
99
+ msgs = [
100
+ {"role": "user", "content": "go"},
101
+ {"role": "assistant", "content": "",
102
+ "tool_calls": [{
103
+ "id": "c0", "type": "function",
104
+ "function": {"name": "observe", "arguments": {}},
105
+ }]},
106
+ {"role": "tool", "tool_call_id": "c0", "content": "ok"},
107
+ ]
108
+ _, conv = BedrockProvider._to_bedrock_messages(msgs)
109
+ # 3 turns: user, assistant, user(toolResult).
110
+ assert [m["role"] for m in conv] == ["user", "assistant", "user"]
111
+ tr = conv[2]["content"][0]["toolResult"]
112
+ assert tr["toolUseId"] == "c0"
113
+ assert tr["content"] == [{"text": "ok"}]
114
+
115
+
116
+ def test_consecutive_user_messages_merge_to_satisfy_alternation():
117
+ """Bedrock REQUIRES alternating user/assistant turns. After a
118
+ tool reply (which becomes a user message) the next briefing is
119
+ also a user message โ€” they must collapse into one user turn."""
120
+ msgs = [
121
+ {"role": "tool", "tool_call_id": "c0", "content": "ok"},
122
+ {"role": "user", "content": "next briefing"},
123
+ ]
124
+ _, conv = BedrockProvider._to_bedrock_messages(msgs)
125
+ assert len(conv) == 1
126
+ assert conv[0]["role"] == "user"
127
+ # toolResult + text under one user message
128
+ assert "toolResult" in conv[0]["content"][0]
129
+ assert conv[0]["content"][1] == {"text": "next briefing"}
130
+
131
+
132
+ # โ”€โ”€ Outbound: tool schemas โ†’ toolConfig โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
133
+
134
+
135
+ def test_tool_schema_translation():
136
+ tools = [{
137
+ "type": "function",
138
+ "function": {
139
+ "name": "move_units",
140
+ "description": "move units to a cell",
141
+ "parameters": {
142
+ "type": "object",
143
+ "properties": {
144
+ "unit_ids": {"type": "array",
145
+ "items": {"type": "integer"}},
146
+ "target_x": {"type": "integer"},
147
+ "target_y": {"type": "integer"},
148
+ },
149
+ "required": ["unit_ids", "target_x", "target_y"],
150
+ },
151
+ },
152
+ }]
153
+ cfg = BedrockProvider._to_bedrock_tools(tools)
154
+ assert cfg is not None
155
+ spec = cfg["tools"][0]["toolSpec"]
156
+ assert spec["name"] == "move_units"
157
+ assert spec["description"] == "move units to a cell"
158
+ assert spec["inputSchema"]["json"]["type"] == "object"
159
+ assert "unit_ids" in spec["inputSchema"]["json"]["properties"]
160
+
161
+
162
+ def test_tool_schema_backfills_object_type_for_paramless_tool():
163
+ tools = [{"type": "function", "function": {
164
+ "name": "observe", "description": "noop",
165
+ "parameters": {"properties": {}},
166
+ }}]
167
+ cfg = BedrockProvider._to_bedrock_tools(tools)
168
+ spec = cfg["tools"][0]["toolSpec"]
169
+ assert spec["inputSchema"]["json"]["type"] == "object"
170
+
171
+
172
+ def test_no_tools_returns_none():
173
+ assert BedrockProvider._to_bedrock_tools([]) is None
174
+
175
+
176
+ # โ”€โ”€ Inbound: Bedrock response โ†’ ChatReply โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€
177
+
178
+
179
+ def test_reply_from_bedrock_text_only():
180
+ resp = {
181
+ "output": {"message": {"role": "assistant",
182
+ "content": [{"text": "hello there"}]}},
183
+ "usage": {"inputTokens": 11, "outputTokens": 3, "totalTokens": 14},
184
+ "stopReason": "end_turn",
185
+ }
186
+ reply = BedrockProvider._reply_from_bedrock(resp)
187
+ assert isinstance(reply, ChatReply)
188
+ assert reply.text == "hello there"
189
+ assert reply.tool_calls == []
190
+ assert reply.usage == {"prompt_tokens": 11, "completion_tokens": 3}
191
+
192
+
193
+ def test_reply_from_bedrock_tool_use():
194
+ resp = {
195
+ "output": {"message": {"role": "assistant", "content": [
196
+ {"text": "I'll scout."},
197
+ {"toolUse": {
198
+ "toolUseId": "tooluse_abc",
199
+ "name": "move_units",
200
+ "input": {"unit_ids": [1004], "target_x": 60, "target_y": 60},
201
+ }},
202
+ ]}},
203
+ "usage": {"inputTokens": 200, "outputTokens": 25},
204
+ "stopReason": "tool_use",
205
+ }
206
+ reply = BedrockProvider._reply_from_bedrock(resp)
207
+ assert reply.text == "I'll scout."
208
+ assert reply.tool_calls == [{
209
+ "name": "move_units",
210
+ "arguments": {"unit_ids": [1004], "target_x": 60, "target_y": 60},
211
+ }]
212
+ assert reply.usage["prompt_tokens"] == 200
213
+
214
+
215
+ def test_reply_from_bedrock_reasoning_block():
216
+ resp = {
217
+ "output": {"message": {"role": "assistant", "content": [
218
+ {"reasoningContent": {
219
+ "reasoningText": {"text": "thinkingโ€ฆ", "signature": "s"},
220
+ }},
221
+ {"text": "answer"},
222
+ ]}},
223
+ "usage": {"inputTokens": 5, "outputTokens": 2},
224
+ }
225
+ reply = BedrockProvider._reply_from_bedrock(resp)
226
+ assert reply.text == "answer"
227
+ assert reply.reasoning == "thinkingโ€ฆ"
228
+
229
+
230
+ # โ”€โ”€ Plumbing: make_provider routes bedrock through BedrockProvider โ”€โ”€โ”€โ”€โ”€
231
+
232
+
233
+ def test_make_provider_routes_to_bedrock():
234
+ from openra_bench.providers import make_provider
235
+
236
+ cfg = ProviderConfig(
237
+ provider="bedrock",
238
+ model="us.anthropic.claude-sonnet-4-6",
239
+ )
240
+
241
+ class _StubClient:
242
+ def converse(self, **kwargs):
243
+ return {
244
+ "output": {"message": {"role": "assistant",
245
+ "content": [{"text": "ack"}]}},
246
+ "usage": {"inputTokens": 1, "outputTokens": 1},
247
+ "stopReason": "end_turn",
248
+ }
249
+
250
+ # Bypass make_provider's lazy boto3 import by constructing directly
251
+ # with a stub client; smoke that complete() round-trips.
252
+ p = BedrockProvider(cfg, client=_StubClient())
253
+ reply = p.complete(
254
+ [{"role": "system", "content": "hi"},
255
+ {"role": "user", "content": "ping"}],
256
+ tools=[],
257
+ )
258
+ assert reply.text == "ack"
259
+ assert reply.usage == {"prompt_tokens": 1, "completion_tokens": 1}