From de28652178e074186af2f0e5f3f7dff0d00e5366 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:16:32 -0700 Subject: [PATCH] refactor(agent/bedrock_adapter): compact docstrings and section comments by hand (keep invariants/why) --- agent/bedrock_adapter.py | 153 +++++++++++++++++---------------------- 1 file changed, 67 insertions(+), 86 deletions(-) diff --git a/agent/bedrock_adapter.py b/agent/bedrock_adapter.py index 4e04433654..9d3581aece 100644 --- a/agent/bedrock_adapter.py +++ b/agent/bedrock_adapter.py @@ -1,9 +1,8 @@ """AWS Bedrock Converse API adapter (boto3, optional dependency). -The AWS credential chain, cross-region inference profiles, guardrails and -control-plane model discovery all work without API keys. OpenAI-format -messages/tools are converted to Converse on the way in and responses normalized -back to OpenAI-shaped objects. +Works without API keys via the AWS credential chain, cross-region inference profiles, guardrails and +control-plane model discovery. OpenAI-format messages/tools are converted to Converse on the way in +and responses normalized back to OpenAI-shaped objects. """ import base64 @@ -33,8 +32,8 @@ except Exception: _bedrock_runtime_client_cache: Dict[str, Any] = {} _bedrock_control_client_cache: Dict[str, Any] = {} -# Bedrock-hosted GPT-5.x models are served from the Bedrock Mantle OpenAI-compatible -# endpoint, not Converse. Narrow allowlist so GPT-OSS models stay on the native path. +# Bedrock-hosted GPT-5.x models are served from the Bedrock Mantle OpenAI-compatible endpoint, not +# Converse. Narrow allowlist so GPT-OSS models stay on the native path. BEDROCK_OPENAI_RESPONSES_MODEL_IDS: Tuple[str, ...] = ( "openai.gpt-5.5", "openai.gpt-5.6-sol", "openai.gpt-5.6-terra", "openai.gpt-5.6-luna", ) @@ -92,7 +91,6 @@ def invalidate_runtime_client(region: str) -> bool: # --- Bedrock Mantle / OpenAI Responses support --- - def is_openai_bedrock_model(model_id: str) -> bool: """True for Bedrock-hosted OpenAI models that require Mantle (GPT-OSS excluded).""" return str(model_id or "").strip().lower() in {m.lower() for m in BEDROCK_OPENAI_RESPONSES_MODEL_IDS} @@ -174,8 +172,8 @@ def build_bedrock_openai_http_client(region: str, *, timeout: Optional[float] = def configure_bedrock_openai_client_kwargs( client_kwargs: Dict[str, Any], *, timeout: Optional[float] = None) -> Dict[str, Any]: - """Install SigV4 auth on OpenAI SDK kwargs for Bedrock Mantle. Real API keys keep the - SDK's bearer auth; the ``aws-sdk``/``no-key-required`` placeholders mean IAM chain auth.""" + """Install SigV4 auth on OpenAI SDK kwargs for Bedrock Mantle; a real API key keeps the SDK's + bearer auth, the ``aws-sdk``/``no-key-required`` placeholders mean IAM chain auth.""" base_url = str(client_kwargs.get("base_url") or "") api_key = client_kwargs.get("api_key") if not is_bedrock_openai_base_url(base_url) or ( @@ -189,9 +187,9 @@ def configure_bedrock_openai_client_kwargs( # --- Stale-connection detection --- -# A pooled connection killed under boto3 (NAT timeout, VPN flap, RST) surfaces as a -# botocore/urllib3 transport error or a bare AssertionError from urllib3's pool -# checks; retrying the same client reproduces it, so the fix is to evict the client. +# A pooled connection killed under boto3 (NAT timeout, VPN flap, RST) surfaces as a botocore/urllib3 +# transport error or a bare AssertionError from urllib3's pool checks; retrying the same client +# reproduces it, so the fix is to evict the client. _STALE_LIB_MODULE_PREFIXES = ("urllib3.", "botocore.", "boto3.") @@ -220,9 +218,8 @@ def is_stale_connection_error(exc: BaseException) -> bool: def is_streaming_access_denied_error(exc: BaseException) -> bool: - """True when IAM denied ``bedrock:InvokeModelWithResponseStream`` (permanent: callers - fall back to non-streaming converse()). Message-based because the AnthropicBedrock SDK - wraps the same AWS response but preserves the action name.""" + """True when IAM denied ``bedrock:InvokeModelWithResponseStream`` (permanent: callers fall back + to converse()). Message-based: the AnthropicBedrock SDK wraps the response but keeps the action name.""" msg = str(exc).lower() if "invokemodelwithresponsestream" not in msg: return False @@ -283,9 +280,8 @@ def resolve_bedrock_region(env: Optional[Dict[str, str]] = None) -> str: def resolve_bedrock_runtime_region(config: Optional[Dict[str, Any]] = None) -> str: - """``bedrock.region`` from config.yaml, else :func:`resolve_bedrock_region`. Every - non-runtime Bedrock endpoint must use this so auxiliary calls never leave the primary - runtime's region when config and ambient AWS env disagree. Pass *config* to skip disk.""" + """``bedrock.region`` from config.yaml, else :func:`resolve_bedrock_region`. Every non-runtime Bedrock + endpoint must use this so auxiliary calls never leave the primary runtime's region. *config* skips disk.""" if config is None: with suppress(Exception): from hermes_cli.config import load_config_readonly @@ -312,8 +308,8 @@ _NON_TOOL_CALLING_PATTERNS = [ "cohere.embed", "amazon.titan-embed", # embeddings ] -# cachePoint allowlist — inverted policy vs tools: unknown models get NO cache markers -# (they reject cachePoint). Claude only reaches build_converse_kwargs under bearer auth. +# cachePoint allowlist — inverted policy vs tools: unknown models get NO cache markers (they reject +# cachePoint). Claude only reaches build_converse_kwargs under bearer auth. _CACHE_POINT_PATTERNS = ["anthropic.claude", "amazon.nova"] @@ -327,10 +323,9 @@ def _model_supports_prompt_cache(model_id: str) -> bool: # --- Server-verdict cachePoint suppression --- -# Bedrock's cachePoint rule is per-family AND per-field (Nova accepts it in -# system/messages but hard-fails on toolConfig.tools) and any static table drifts, so -# when Bedrock names a placement as unpermitted we record the verdict, drop the marker -# there for the rest of the process, and retry once without it. +# Bedrock's cachePoint rule is per-family AND per-field (Nova accepts it in system/messages but hard-fails +# on toolConfig.tools) and any static table drifts, so when Bedrock names a placement as unpermitted we +# record the verdict, drop the marker there for the rest of the process, and retry once without it. CACHE_POINT_PLACEMENTS = ("tools", "system", "messages") @@ -343,9 +338,8 @@ _CACHE_POINT = {"cachePoint": {"type": "default"}} def cache_point_rejection_placement(exc: BaseException) -> Optional[str]: - """Converse section whose cachePoint Bedrock refused, or None. Message-based: the JSON - pointer in the ValidationException is the only thing naming the section, and the same - wording arrives raw or SDK-wrapped. Unlocalisable → "tools" (the only known refusal).""" + """Converse section whose cachePoint Bedrock refused, or None. Message-based: the JSON pointer in the + ValidationException is the only thing naming the section (raw or SDK-wrapped). Unlocalisable → "tools".""" msg = str(exc) lowered = msg.lower() if "cachepoint" not in lowered or ("not permitted" not in lowered and "extraneous" not in lowered): @@ -430,8 +424,7 @@ _ANTHROPIC_BEDROCK_MODEL_RE = re.compile( def is_anthropic_bedrock_model(model_id: str) -> bool: - """True for Claude on Bedrock (``anthropic.claude-*`` with any regional prefix); these - use the AnthropicBedrock SDK path, everything else the Converse path.""" + """True for Claude on Bedrock (``anthropic.claude-*``, any regional prefix): AnthropicBedrock SDK path.""" return _ANTHROPIC_BEDROCK_MODEL_RE.match(model_id) is not None @@ -457,8 +450,8 @@ def _safe_text(text) -> str: def _image_block_from_data_url(url: str) -> Dict: - """``data:;base64,...`` → Converse image block with RAW bytes (boto3 base64-encodes - at the wire layer; passing the string through double-encodes and Bedrock rejects it).""" + """``data:;base64,...`` → Converse image block with RAW bytes (boto3 base64-encodes on the + wire; passing the string through double-encodes and Bedrock rejects it).""" header, _, data = url.partition(",") media_type = (header[5:].split(";")[0] if header.startswith("data:") else "") or "image/jpeg" try: @@ -469,8 +462,7 @@ def _image_block_from_data_url(url: str) -> Dict: def _convert_content_to_converse(content) -> List[Dict]: - """OpenAI message content → Converse blocks. Blank text becomes the placeholder; remote - image URLs (unsupported by Converse) become a text reference.""" + """OpenAI content → Converse blocks; blank text → placeholder, remote image URLs → text reference.""" if not isinstance(content, list): return [{"text": _safe_text(content)}] blocks = [] @@ -510,8 +502,8 @@ def _decode_redacted(encoded) -> Optional[bytes]: def _replay_ordered_blocks(ordered_blocks: List) -> List[Dict]: - """Rebuild the exact Bedrock block sequence captured at normalization time; redacted - reasoning is stored base64 (JSON-safe sidecar) and undecodable entries are skipped.""" + """Rebuild the exact Bedrock block sequence captured at normalization time; redacted reasoning is + stored base64 (JSON-safe sidecar) and undecodable entries are skipped.""" content_blocks: List[Dict] = [] for block in ordered_blocks: if not isinstance(block, dict): @@ -545,9 +537,8 @@ def _parse_tool_args(args) -> Any: def _assistant_blocks(msg: Dict, content) -> List[Dict]: - """Assistant message → Converse blocks. An ordered ``bedrock_content_blocks`` sidecar is - authoritative; otherwise redacted thinking from ``reasoning_details`` (byte-for-byte - round-trip), then text, then tool calls.""" + """Assistant message → Converse blocks. An ordered ``bedrock_content_blocks`` sidecar is authoritative; + otherwise redacted thinking from ``reasoning_details`` (byte-for-byte), then text, then tool calls.""" ordered_blocks = msg.get("bedrock_content_blocks") if isinstance(ordered_blocks, list) and (content_blocks := _replay_ordered_blocks(ordered_blocks)): return content_blocks @@ -567,9 +558,9 @@ def _assistant_blocks(msg: Dict, content) -> List[Dict]: def convert_messages_to_converse(messages: List[Dict]) -> Tuple[Optional[List[Dict]], List[Dict]]: - """OpenAI messages → ``(system_blocks_or_None, converse_messages)``. Tool results become - ``toolResult`` user blocks. Converse requires strict user/assistant alternation with a - user turn first and last: same-role neighbours merge, placeholder user turns pad the ends.""" + """OpenAI messages → ``(system_blocks_or_None, converse_messages)``; tool results become ``toolResult`` + user blocks. Converse needs strict user/assistant alternation with a user turn first and last: + same-role neighbours merge, placeholder user turns pad the ends.""" system_blocks: List[Dict] = [] converse_msgs: List[Dict] = [] @@ -648,8 +639,8 @@ class _ResponseParts: block["redactedContentBase64"] = encoded def build(self, ordered_blocks: List[Dict[str, Any]], usage_data: Dict[str, int], stop_reason: str, model: str) -> SimpleNamespace: - """Assemble the OpenAI-shaped response. Converse's inputTokens EXCLUDES cache - read/write tokens (OpenAI's prompt_tokens includes them), so they are added back.""" + """Assemble the OpenAI-shaped response. Converse's inputTokens EXCLUDES cache read/write tokens + (OpenAI's prompt_tokens includes them), so they are added back.""" msg = SimpleNamespace( role="assistant", content="\n".join(self.text_parts) if self.text_parts else None, tool_calls=self.tool_calls or None, reasoning_details=self.reasoning_details or None, @@ -673,9 +664,8 @@ class _ResponseParts: def normalize_converse_response(response: Dict) -> SimpleNamespace: - """Bedrock Converse response → OpenAI ``ChatCompletion``-shaped SimpleNamespace - (``.choices[0].message.{content,tool_calls,reasoning_content,reasoning_details, - bedrock_content_blocks}``, ``.choices[0].finish_reason``, ``.usage``).""" + """Bedrock Converse response → OpenAI ``ChatCompletion``-shaped SimpleNamespace (``.choices[0].message. + {content,tool_calls,reasoning_content,reasoning_details,bedrock_content_blocks}``, ``finish_reason``, ``.usage``).""" parts = _ResponseParts() ordered_blocks = [] for block in response.get("output", {}).get("message", {}).get("content", []): @@ -707,11 +697,10 @@ def stream_converse_with_callbacks( event_stream, on_text_delta=None, on_tool_start=None, on_reasoning_delta=None, on_interrupt_check=None, on_event=None, ) -> SimpleNamespace: - """Process a boto3 ``converse_stream()`` response with callbacks → the - ``normalize_converse_response()`` shape. ``on_text_delta`` only fires while no toolUse - block has been seen (as on the Anthropic/chat_completions paths); ``on_interrupt_check`` - True stops streaming; ``on_event`` fires for EVERY event before branching and its - exceptions are swallowed so a watchdog hook can never abort the stream.""" + """boto3 ``converse_stream()`` response + callbacks → the ``normalize_converse_response()`` shape. + ``on_text_delta`` only fires while no toolUse block has been seen (as on the Anthropic/chat_completions + paths); ``on_interrupt_check`` True stops streaming; ``on_event`` fires for EVERY event before branching + and its exceptions are swallowed so a watchdog hook can never abort the stream.""" parts = _ResponseParts() stream_blocks: Dict[int, Dict[str, Any]] = {} current_block_index: Optional[int] = None @@ -791,12 +780,10 @@ def build_converse_kwargs( temperature: Optional[float] = None, top_p: Optional[float] = None, stop_sequences: Optional[List[str]] = None, guardrail_config: Optional[Dict] = None, ) -> Dict[str, Any]: - """Build kwargs for ``bedrock-runtime.converse()`` / ``converse_stream()``. - - ``max_tokens=None`` omits ``maxTokens`` (model maximum; default stays 4096). cachePoint - markers go on system, tools and the second-newest message (survives as the tail grows — - mirrors Anthropic system_and_3), each only if the model supports caching and Bedrock has - not rejected that placement.""" + """Build kwargs for ``bedrock-runtime.converse()`` / ``converse_stream()``. ``max_tokens=None`` omits + ``maxTokens`` (model maximum; default stays 4096). cachePoint markers go on system, tools and the + second-newest message (survives as the tail grows — mirrors Anthropic system_and_3), each only if the + model supports caching and Bedrock has not rejected that placement.""" system_prompt, converse_messages = convert_messages_to_converse(messages) cache_enabled = _model_supports_prompt_cache(model) @@ -837,8 +824,8 @@ def call_converse( max_tokens: Optional[int] = 4096, temperature: Optional[float] = None, top_p: Optional[float] = None, stop_sequences: Optional[List[str]] = None, guardrail_config: Optional[Dict] = None, ) -> SimpleNamespace: - """Non-streaming Converse call → OpenAI-compatible response. Retries once without a - rejected cachePoint placement; evicts the cached client on stale-connection errors.""" + """Non-streaming Converse call → OpenAI-compatible response. Retries once without a rejected cachePoint + placement; evicts the cached client on stale-connection errors.""" client = _get_bedrock_runtime_client(region) kwargs = build_converse_kwargs(model, messages, tools, max_tokens, temperature, top_p, stop_sequences, guardrail_config) try: @@ -916,8 +903,8 @@ def _list_inference_profiles(client, filter_set: set, models: List[Dict[str, Any def discover_bedrock_models(region: str, provider_filter: Optional[List[str]] = None) -> List[Dict[str, Any]]: - """Discover foundation models + inference profiles (cached 1h per region/filter), sorted - ``global.`` profiles first then by name; [] when the client cannot be built.""" + """Foundation models + inference profiles (cached 1h per region/filter), ``global.`` profiles first then + by name; [] when the client cannot be built.""" cache_key = f"{region}:{','.join(sorted(provider_filter or []))}" cached = _discovery_cache.get(cache_key) if cached and (time.time() - cached["timestamp"]) < _DISCOVERY_CACHE_TTL_SECONDS: @@ -949,48 +936,43 @@ def _extract_provider_from_arn(arn: str) -> str: # --- Bedrock model context lengths --- -# Static fallback when the live probe is unavailable (agent/model_metadata.py). Keys match -# by longest substring, so versioned entries win over the generic "anthropic.claude-opus-4". +# Static fallback when the live probe is unavailable (agent/model_metadata.py). Keys match by longest +# substring, so versioned entries win over the generic "anthropic.claude-opus-4". BEDROCK_CONTEXT_LENGTHS: Dict[str, int] = { - # Anthropic Claude: 1M GA vs 200K. The 1M entries must match - # agent/model_metadata.py DEFAULT_CONTEXT_LENGTHS or context compresses early. + # Anthropic Claude: 1M GA vs 200K. The 1M entries must match agent/model_metadata.py + # DEFAULT_CONTEXT_LENGTHS or context compresses early. **dict.fromkeys(( - "anthropic.claude-fable-5", "anthropic.claude-fable", "anthropic.claude-sonnet-5", - "anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6", - "anthropic.claude-sonnet-4-6", + "anthropic.claude-fable-5", "anthropic.claude-fable", "anthropic.claude-sonnet-5", "anthropic.claude-opus-4-8", + "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6", "anthropic.claude-sonnet-4-6", ), 1_000_000), **dict.fromkeys(( - "anthropic.claude-sonnet-4-5", "anthropic.claude-haiku-4-5", "anthropic.claude-opus-4", - "anthropic.claude-sonnet-4", "anthropic.claude-3-5-sonnet", "anthropic.claude-3-5-haiku", - "anthropic.claude-3-opus", "anthropic.claude-3-sonnet", "anthropic.claude-3-haiku", + "anthropic.claude-sonnet-4-5", "anthropic.claude-haiku-4-5", "anthropic.claude-opus-4", "anthropic.claude-sonnet-4", + "anthropic.claude-3-5-sonnet", "anthropic.claude-3-5-haiku", "anthropic.claude-3-opus", "anthropic.claude-3-sonnet", + "anthropic.claude-3-haiku", ), 200_000), # Amazon Nova **dict.fromkeys(("amazon.nova-pro", "amazon.nova-lite"), 300_000), "amazon.nova-micro": 128_000, # Meta Llama / Mistral / DeepSeek **dict.fromkeys(( - "meta.llama4-maverick", "meta.llama4-scout", "meta.llama3-3-70b-instruct", - "mistral.mistral-large", "deepseek.v3", + "meta.llama4-maverick", "meta.llama4-scout", "meta.llama3-3-70b-instruct", "mistral.mistral-large", "deepseek.v3", ), 128_000), - # OpenAI on Bedrock (Mantle/Responses route) - # https://docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html + # OpenAI on Bedrock (Mantle/Responses route): docs.aws.amazon.com/bedrock/latest/userguide/model-cards-openai.html **dict.fromkeys(BEDROCK_OPENAI_RESPONSES_MODEL_IDS, 272_000), } -# Default for unknown Bedrock models -BEDROCK_DEFAULT_CONTEXT_LENGTH = 128_000 +BEDROCK_DEFAULT_CONTEXT_LENGTH = 128_000 # unknown Bedrock models -# Probe padding tiers (tokens): a wildly oversized payload yields an opaque -# InternalServerException instead of a clean ValidationException. +# Probe padding tiers (tokens): a wildly oversized payload yields an opaque InternalServerException +# instead of a clean ValidationException. _BEDROCK_PROBE_TIERS = (1_300_000, 2_200_000) _WORDS_PER_TOKEN = 0.9 # conservative: ensures the padded prompt clears the tier def probe_bedrock_context_length(model_id: str, region: str) -> Optional[int]: - """Discover a model's real context window by provoking a length error — the only - authoritative source ("prompt is too long: 1300032 tokens > 1000000 maximum"); length - validation runs before inference so the probe costs nothing. An accepted tier is - returned as a safe lower bound; None (no creds / network / unparseable) → static table.""" + """Discover a model's real context window by provoking a length error — the only authoritative source + ("prompt is too long: 1300032 tokens > 1000000 maximum"); length validation runs before inference so the + probe costs nothing. An accepted tier is a safe lower bound; None (no creds/network/unparseable) → static table.""" from agent.model_metadata import parse_context_limit_from_error try: client = _get_bedrock_runtime_client(region) @@ -1018,9 +1000,8 @@ def probe_bedrock_context_length(model_id: str, region: str) -> Optional[int]: def get_bedrock_context_length(model_id: str, region: str = "", probe: bool = True) -> int: - """Context window: live probe (if ``probe`` and ``region``) → static table → default. The - table is fallback only: a stale substring match silently caps the window (a 1M Opus - pinned to 200K via "opus-4"). ``probe=False`` / empty region skips the network call.""" + """Context window: live probe (if ``probe`` and ``region``) → static table → default. The table is fallback + only: a stale substring match silently caps the window (a 1M Opus pinned to 200K via "opus-4").""" if probe and region and (probed := probe_bedrock_context_length(model_id, region)): return probed matches = [key for key in BEDROCK_CONTEXT_LENGTHS if key in model_id.lower()]