diff --git a/agent/bedrock_adapter.py b/agent/bedrock_adapter.py index 50f05e3aa29f..c8cff3f76e13 100644 --- a/agent/bedrock_adapter.py +++ b/agent/bedrock_adapter.py @@ -1393,9 +1393,22 @@ BEDROCK_CONTEXT_LENGTHS: Dict[str, int] = { # Default for unknown Bedrock models BEDROCK_DEFAULT_CONTEXT_LENGTH = 128_000 +# Probe tiers (in tokens). We send a request padded just past each tier and +# read the real window from Bedrock's length-validation error. Two reasons +# this is tiered rather than one giant request: +# 1. A wildly oversized payload (e.g. 5M tokens) makes Bedrock return an +# opaque InternalServerException after retries instead of a clean +# ValidationException — so we must stay within a sane overage. +# 2. Stepping up lets us discover larger windows (2M+) without over-padding +# smaller ones. +# Each tier value is the *padding target*; the error reports the true maximum, +# which is what we actually return. +_BEDROCK_PROBE_TIERS = (1_300_000, 2_200_000) +_WORDS_PER_TOKEN = 0.9 # conservative: ensures the padded prompt clears the tier -def get_bedrock_context_length(model_id: str) -> int: - """Look up the context window size for a Bedrock model. + +def _static_bedrock_context_length(model_id: str) -> int: + """Longest-substring-match lookup against the static fallback table. Uses substring matching so versioned IDs like ``anthropic.claude-sonnet-4-6-20250514-v1:0`` resolve correctly. @@ -1408,3 +1421,103 @@ def get_bedrock_context_length(model_id: str) -> int: best_key = key best_val = val return best_val + + +def probe_bedrock_context_length(model_id: str, region: str) -> Optional[int]: + """Discover a Bedrock model's real context window by provoking a length error. + + Bedrock does not expose the context window via any metadata API + (``get-foundation-model`` omits it, ``Converse`` metrics omit it, + ``CountTokens`` is unsupported on several models). The only authoritative + source is the ``ValidationException`` raised when a prompt exceeds the + window: + + "The model returned the following errors: prompt is too long: + 1300032 tokens > 1000000 maximum" + + Length validation happens *before* inference, so an oversized request is + rejected immediately and cheaply — no tokens are generated and no input is + actually processed. We pad a request just past each tier in + ``_BEDROCK_PROBE_TIERS`` and parse the reported ``maximum``. Tiers exist + because (a) a *wildly* oversized payload makes Bedrock fail with an opaque + InternalServerException instead of a clean length error, and (b) stepping + up discovers larger windows without over-padding smaller ones. + + Returns the detected window, or ``None`` if the probe could not run + (missing credentials, network error, or no parseable limit) so the caller + can fall back to the static table. + """ + try: + from agent.model_metadata import parse_context_limit_from_error + except ImportError: # pragma: no cover — same package + return None + + try: + client = _get_bedrock_runtime_client(region) + except Exception as exc: # boto3 missing / credential resolution failure + logger.debug("Bedrock context probe skipped for %s: %s", model_id, exc) + return None + + last_error = "" + for tier_tokens in _BEDROCK_PROBE_TIERS: + pad_words = int(tier_tokens / _WORDS_PER_TOKEN) + oversized = "data " * pad_words + try: + client.converse( + modelId=model_id, + messages=[{"role": "user", "content": [{"text": oversized}]}], + inferenceConfig={"maxTokens": 8}, + ) + # Accepted a prompt this large → the window is at least this tier. + # Returning the tier as a lower bound is safe and avoids inventing + # a number we can't confirm. + logger.debug( + "Bedrock context probe for %s accepted ~%s-token prompt; " + "window is at least that", model_id, f"{tier_tokens:,}", + ) + return tier_tokens + except Exception as exc: + msg = str(exc) + last_error = msg + limit = parse_context_limit_from_error(msg) + if limit and limit >= 1024: + logger.info( + "Probed Bedrock context window for %s: %s tokens", + model_id, f"{limit:,}", + ) + return limit + # No parseable limit at this tier (opaque server error, auth, + # throttle). Try the next, smaller-overage strategy is N/A here — + # tiers ascend — so just continue; if all fail we return None. + continue + + logger.debug( + "Bedrock context probe for %s returned no parseable limit: %s", + model_id, last_error[:200], + ) + return None + + +def get_bedrock_context_length(model_id: str, region: str = "", probe: bool = True) -> int: + """Resolve the context window for a Bedrock model. + + Resolution order: + 1. Live probe against Bedrock (authoritative; cached by the caller). + 2. Static fallback table (longest-substring match). + 3. Conservative default. + + The static table is intentionally a *fallback*, not the primary source: + AWS ships new model versions (opus-4-7, opus-4-8, ...) faster than the + table can track, and a stale entry silently caps the window (e.g. a + 1M-token Opus pinned to 200K via an ``opus-4`` substring match). The + probe asks Bedrock directly so every model — current or future — gets its + real window with no table maintenance. + + ``probe=False`` (or an empty ``region``) skips the network call and uses + the static table only — used by pure-offline/display code paths. + """ + if probe and region: + probed = probe_bedrock_context_length(model_id, region) + if probed: + return probed + return _static_bedrock_context_length(model_id) diff --git a/agent/model_metadata.py b/agent/model_metadata.py index 7e193aab09a6..bd3505c251d3 100644 --- a/agent/model_metadata.py +++ b/agent/model_metadata.py @@ -2287,10 +2287,42 @@ def get_model_context_length( # back to the 128K default before reaching the original step 4b branch. if is_bedrock_context: try: - from agent.bedrock_adapter import get_bedrock_context_length - return get_bedrock_context_length(model) + from agent.bedrock_adapter import ( + get_bedrock_context_length, + resolve_bedrock_region, + ) except ImportError: pass # boto3 not installed — fall through to generic resolution + else: + # Bedrock does not expose the context window via any metadata API, + # so get_bedrock_context_length() probes the live endpoint (one + # fast, pre-inference length rejection) to read the real window. + # Cache the probe result per model so we pay that cost once, not + # every turn — keyed by base_url when present, else a synthetic + # bedrock:// key so display/offline paths share the entry. + cache_key_url = base_url or "bedrock://" + cached = get_cached_context_length(model, cache_key_url) + if cached is not None: + return cached + # Resolve region from the base_url host first, then the standard + # AWS region chain. An empty region disables probing (table only). + region = "" + if base_url: + _m = re.search(r"bedrock-runtime\.([a-z0-9-]+)\.", base_url) + if _m: + region = _m.group(1) + if not region: + try: + region = resolve_bedrock_region() + except Exception: + region = "" + ctx = get_bedrock_context_length(model, region=region, probe=bool(region)) + if ctx and region: + # Only persist probe-derived values (region present); a pure + # table fallback shouldn't poison the cache against a later + # successful probe. + save_context_length(model, cache_key_url, ctx) + return ctx if provider == "novita" or (base_url and base_url_host_matches(base_url, "api.novita.ai")): ctx = _resolve_endpoint_context_length(model, base_url or "https://api.novita.ai/openai/v1", api_key=api_key) diff --git a/tests/agent/test_bedrock_adapter.py b/tests/agent/test_bedrock_adapter.py index 2f8a6dc652fa..56b99838d0be 100644 --- a/tests/agent/test_bedrock_adapter.py +++ b/tests/agent/test_bedrock_adapter.py @@ -1238,6 +1238,71 @@ class TestBedrockContextLength: # "anthropic.claude-3-5-sonnet" should match before "anthropic.claude-3" assert get_bedrock_context_length("anthropic.claude-3-5-sonnet-20240620-v1:0") == 200_000 + def test_no_region_skips_probe_uses_table(self): + # Default call (no region) must NOT hit the network — returns the + # static table value. Guards backward compatibility for callers that + # still invoke get_bedrock_context_length(model_id) with one arg. + from agent.bedrock_adapter import get_bedrock_context_length + with patch("agent.bedrock_adapter.probe_bedrock_context_length") as mock_probe: + assert get_bedrock_context_length("anthropic.claude-opus-4-6") == 200_000 + mock_probe.assert_not_called() + + +class TestBedrockContextProbe: + """Test the live context-window probe that reads the real window from + Bedrock's 'prompt is too long' validation error.""" + + def _client_raising(self, message): + client = MagicMock() + client.converse.side_effect = Exception(message) + return client + + def test_probe_parses_real_window_from_error(self): + from agent.bedrock_adapter import probe_bedrock_context_length + err = ( + "An error occurred (ValidationException) when calling the Converse " + "operation: The model returned the following errors: prompt is too " + "long: 5000032 tokens > 1000000 maximum" + ) + with patch("agent.bedrock_adapter._get_bedrock_runtime_client", + return_value=self._client_raising(err)): + assert probe_bedrock_context_length( + "eu.anthropic.claude-opus-4-8", "eu-central-1") == 1_000_000 + + def test_probe_returns_none_on_unparseable_error(self): + from agent.bedrock_adapter import probe_bedrock_context_length + err = "An error occurred (AccessDeniedException): not authorized" + with patch("agent.bedrock_adapter._get_bedrock_runtime_client", + return_value=self._client_raising(err)): + assert probe_bedrock_context_length( + "eu.anthropic.claude-opus-4-8", "eu-central-1") is None + + def test_probe_returns_none_when_client_unavailable(self): + from agent.bedrock_adapter import probe_bedrock_context_length + with patch("agent.bedrock_adapter._get_bedrock_runtime_client", + side_effect=RuntimeError("boto3 missing")): + assert probe_bedrock_context_length("any.model", "eu-central-1") is None + + def test_probe_result_beats_static_table(self): + # A successful probe (1M) must override the stale table value (200K + # via the 'anthropic.claude-opus-4' substring match). + from agent.bedrock_adapter import get_bedrock_context_length + err = "prompt is too long: 5000032 tokens > 1000000 maximum" + with patch("agent.bedrock_adapter._get_bedrock_runtime_client", + return_value=self._client_raising(err)): + assert get_bedrock_context_length( + "eu.anthropic.claude-opus-4-8", + region="eu-central-1") == 1_000_000 + + def test_probe_failure_falls_back_to_table(self): + from agent.bedrock_adapter import get_bedrock_context_length + err = "AccessDeniedException: nope" + with patch("agent.bedrock_adapter._get_bedrock_runtime_client", + return_value=self._client_raising(err)): + # opus-4-6 is in the table at 200K; probe fails → table wins. + assert get_bedrock_context_length( + "anthropic.claude-opus-4-6", region="eu-central-1") == 200_000 + # --------------------------------------------------------------------------- # Tool-calling capability detection