Skip to content

Commit 46091e3

Browse files
committed
feat: add context window limit lookup table
1 parent a245e6d commit 46091e3

12 files changed

Lines changed: 376 additions & 6 deletions

File tree

src/strands/models/_defaults.py

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
"""Default model metadata lookup tables.
2+
3+
Provides context window limits for known model IDs across all providers.
4+
Values sourced from provider documentation and
5+
https://github.com/BerriAI/litellm/blob/litellm_internal_staging/model_prices_and_context_window.json
6+
"""
7+
8+
from typing import Any, TypeVar
9+
10+
_T = TypeVar("_T", bound=dict[str, Any])
11+
12+
# Context window limits (in tokens) for known model IDs.
13+
#
14+
# Best-effort lookup table — unknown models return None and callers
15+
# fall back gracefully (e.g. proactive compression is disabled).
16+
# Users can always override with an explicit context_window_limit in their model config.
17+
#
18+
# For Bedrock models with cross-region prefixes (e.g. us., eu., global.),
19+
# get_context_window_limit strips the prefix before lookup so only the base model ID is needed here.
20+
_CONTEXT_WINDOW_LIMITS: dict[str, int] = {
21+
# Anthropic (direct API)
22+
"claude-sonnet-4-6": 1_000_000,
23+
"claude-sonnet-4-20250514": 1_000_000,
24+
"claude-sonnet-4-5": 200_000,
25+
"claude-sonnet-4-5-20250929": 200_000,
26+
"claude-opus-4-6": 1_000_000,
27+
"claude-opus-4-6-20260205": 1_000_000,
28+
"claude-opus-4-7": 1_000_000,
29+
"claude-opus-4-7-20260416": 1_000_000,
30+
"claude-opus-4-5": 200_000,
31+
"claude-opus-4-5-20251101": 200_000,
32+
"claude-opus-4-20250514": 200_000,
33+
"claude-opus-4-1": 200_000,
34+
"claude-opus-4-1-20250805": 200_000,
35+
"claude-haiku-4-5": 200_000,
36+
"claude-haiku-4-5-20251001": 200_000,
37+
"claude-3-7-sonnet-20250219": 200_000,
38+
"claude-3-5-sonnet-20241022": 200_000,
39+
"claude-3-5-sonnet-20240620": 200_000,
40+
"claude-3-5-haiku-20241022": 200_000,
41+
"claude-3-opus-20240229": 200_000,
42+
"claude-3-haiku-20240307": 200_000,
43+
# Bedrock Anthropic (base model IDs — cross-region prefixes stripped by get_context_window_limit)
44+
"anthropic.claude-sonnet-4-6": 1_000_000,
45+
"anthropic.claude-sonnet-4-20250514-v1:0": 1_000_000,
46+
"anthropic.claude-sonnet-4-5-20250929-v1:0": 200_000,
47+
"anthropic.claude-opus-4-6-v1": 1_000_000,
48+
"anthropic.claude-opus-4-7": 1_000_000,
49+
"anthropic.claude-opus-4-5-20251101-v1:0": 200_000,
50+
"anthropic.claude-opus-4-20250514-v1:0": 200_000,
51+
"anthropic.claude-opus-4-1-20250805-v1:0": 200_000,
52+
"anthropic.claude-haiku-4-5-20251001-v1:0": 200_000,
53+
"anthropic.claude-haiku-4-5@20251001": 200_000,
54+
"anthropic.claude-3-7-sonnet-20250219-v1:0": 200_000,
55+
"anthropic.claude-3-7-sonnet-20240620-v1:0": 200_000,
56+
"anthropic.claude-3-5-sonnet-20241022-v2:0": 200_000,
57+
"anthropic.claude-3-5-sonnet-20240620-v1:0": 200_000,
58+
"anthropic.claude-3-5-haiku-20241022-v1:0": 200_000,
59+
"anthropic.claude-3-opus-20240229-v1:0": 200_000,
60+
"anthropic.claude-3-haiku-20240307-v1:0": 200_000,
61+
"anthropic.claude-3-sonnet-20240229-v1:0": 200_000,
62+
"anthropic.claude-mythos-preview": 1_000_000,
63+
# Bedrock Amazon Nova
64+
"amazon.nova-pro-v1:0": 300_000,
65+
"amazon.nova-lite-v1:0": 300_000,
66+
"amazon.nova-micro-v1:0": 128_000,
67+
"amazon.nova-premier-v1:0": 1_000_000,
68+
"amazon.nova-2-lite-v1:0": 1_000_000,
69+
"amazon.nova-2-pro-preview-20251202-v1:0": 1_000_000,
70+
# OpenAI
71+
"gpt-5.5": 1_050_000,
72+
"gpt-5.5-pro": 1_050_000,
73+
"gpt-5.4": 1_050_000,
74+
"gpt-5.4-pro": 1_050_000,
75+
"gpt-5.4-mini": 272_000,
76+
"gpt-5.4-nano": 272_000,
77+
"gpt-5.2": 272_000,
78+
"gpt-5.2-pro": 272_000,
79+
"gpt-5.1": 272_000,
80+
"gpt-5": 272_000,
81+
"gpt-5-mini": 272_000,
82+
"gpt-5-nano": 272_000,
83+
"gpt-5-pro": 128_000,
84+
"gpt-4.1": 1_047_576,
85+
"gpt-4.1-mini": 1_047_576,
86+
"gpt-4.1-nano": 1_047_576,
87+
"gpt-4o": 128_000,
88+
"gpt-4o-mini": 128_000,
89+
"gpt-4-turbo": 128_000,
90+
"o3": 200_000,
91+
"o3-mini": 200_000,
92+
"o3-pro": 200_000,
93+
"o4-mini": 200_000,
94+
"o1": 200_000,
95+
# Google Gemini
96+
"gemini-2.5-flash": 1_048_576,
97+
"gemini-2.5-flash-lite": 1_048_576,
98+
"gemini-2.5-pro": 1_048_576,
99+
"gemini-2.0-flash": 1_048_576,
100+
"gemini-2.0-flash-lite": 1_048_576,
101+
"gemini-3-pro-preview": 1_048_576,
102+
"gemini-3-flash-preview": 1_048_576,
103+
"gemini-3.1-pro-preview": 1_048_576,
104+
"gemini-3.1-flash-lite-preview": 1_048_576,
105+
}
106+
107+
108+
def get_context_window_limit(model_id: str) -> int | None:
109+
"""Look up the context window limit for a model ID.
110+
111+
For Bedrock cross-region model IDs (e.g. ``us.anthropic.claude-sonnet-4-6``),
112+
the region prefix is stripped as a fallback if the direct lookup fails.
113+
114+
Args:
115+
model_id: The model ID to look up.
116+
117+
Returns:
118+
The context window limit in tokens, or None if not found.
119+
"""
120+
direct = _CONTEXT_WINDOW_LIMITS.get(model_id)
121+
if direct is not None:
122+
return direct
123+
124+
# Fallback: strip prefix before first dot and retry (handles cross-region prefixes)
125+
dot_index = model_id.find(".")
126+
if dot_index != -1:
127+
return _CONTEXT_WINDOW_LIMITS.get(model_id[dot_index + 1 :])
128+
129+
return None
130+
131+
132+
def resolve_config_metadata(config: _T, model_id: str) -> _T:
133+
"""Resolve model metadata fields on a config dict from built-in lookup tables.
134+
135+
When ``context_window_limit`` is not explicitly set, looks it up from the built-in table.
136+
Explicit values pass through unchanged. Returns a new dict only when resolution adds a field;
137+
otherwise returns the original config to avoid unnecessary allocation.
138+
139+
Args:
140+
config: The stored model config dict.
141+
model_id: The model ID to look up.
142+
143+
Returns:
144+
The config with resolved metadata, or the original config if nothing to resolve.
145+
"""
146+
if "context_window_limit" in config:
147+
return config
148+
149+
limit = get_context_window_limit(model_id)
150+
if limit is None:
151+
return config
152+
153+
return {**config, "context_window_limit": limit} # type: ignore[return-value]

src/strands/models/anthropic.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException
2121
from ..types.streaming import StreamEvent
2222
from ..types.tools import ToolChoice, ToolChoiceToolDict, ToolSpec
23+
from ._defaults import resolve_config_metadata
2324
from ._validation import _has_location_source, validate_config_keys
2425
from .model import BaseModelConfig, Model
2526

@@ -95,7 +96,7 @@ def get_config(self) -> AnthropicConfig:
9596
Returns:
9697
The Anthropic model configuration.
9798
"""
98-
return self.config
99+
return resolve_config_metadata(self.config, self.config["model_id"])
99100

100101
def _format_request_message_content(self, content: ContentBlock) -> dict[str, Any]:
101102
"""Format an Anthropic content block.

src/strands/models/bedrock.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
)
3232
from ..types.streaming import CitationsDelta, StreamEvent
3333
from ..types.tools import ToolChoice, ToolSpec
34+
from ._defaults import resolve_config_metadata
3435
from ._strict_schema import ensure_strict_json_schema
3536
from ._validation import validate_config_keys
3637
from .model import BaseModelConfig, CacheConfig, Model
@@ -217,7 +218,7 @@ def get_config(self) -> BedrockConfig:
217218
Returns:
218219
The Bedrock model configuration.
219220
"""
220-
return self.config
221+
return resolve_config_metadata(self.config, self.config.get("model_id", ""))
221222

222223
def _format_request(
223224
self,

src/strands/models/gemini.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException, ProviderTokenCountError
2020
from ..types.streaming import StreamEvent
2121
from ..types.tools import ToolChoice, ToolSpec
22+
from ._defaults import resolve_config_metadata
2223
from ._validation import _has_location_source, validate_config_keys
2324
from .model import BaseModelConfig, Model
2425

@@ -115,7 +116,7 @@ def get_config(self) -> GeminiConfig:
115116
Returns:
116117
The Gemini model configuration.
117118
"""
118-
return self.config
119+
return resolve_config_metadata(self.config, self.config["model_id"])
119120

120121
def _get_client(self) -> genai.Client:
121122
"""Get a Gemini client for making requests.

src/strands/models/openai.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException
2222
from ..types.streaming import StreamEvent
2323
from ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse
24+
from ._defaults import resolve_config_metadata
2425
from ._openai_bedrock import BedrockMantleConfig, resolve_bedrock_client_args
2526
from ._validation import _has_location_source, validate_config_keys
2627
from .model import BaseModelConfig, Model
@@ -150,7 +151,9 @@ def get_config(self) -> OpenAIConfig:
150151
Returns:
151152
The OpenAI model configuration.
152153
"""
153-
return cast(OpenAIModel.OpenAIConfig, self.config)
154+
return cast(
155+
OpenAIModel.OpenAIConfig, resolve_config_metadata(self.config, str(self.config.get("model_id", "")))
156+
)
154157

155158
@classmethod
156159
def format_request_message_content(cls, content: ContentBlock, **kwargs: Any) -> dict[str, Any]:

src/strands/models/openai_responses.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException # noqa: E402
5959
from ..types.streaming import StreamEvent # noqa: E402
6060
from ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse # noqa: E402
61+
from ._defaults import resolve_config_metadata # noqa: E402
6162
from ._openai_bedrock import BedrockMantleConfig, resolve_bedrock_client_args # noqa: E402
6263
from ._validation import validate_config_keys # noqa: E402
6364
from .model import BaseModelConfig, Model # noqa: E402
@@ -210,7 +211,10 @@ def get_config(self) -> OpenAIResponsesConfig:
210211
Returns:
211212
The OpenAI Responses API model configuration.
212213
"""
213-
return cast(OpenAIResponsesModel.OpenAIResponsesConfig, self.config)
214+
return cast(
215+
OpenAIResponsesModel.OpenAIResponsesConfig,
216+
resolve_config_metadata(self.config, str(self.config.get("model_id", ""))),
217+
)
214218

215219
@override
216220
async def count_tokens(

tests/strands/models/test_anthropic.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,30 @@ def test__init__model_configs(anthropic_client, model_id, max_tokens):
8282
assert tru_temperature == exp_temperature
8383

8484

85+
def test__init__auto_populates_context_window_limit(anthropic_client):
86+
_ = anthropic_client
87+
88+
model = AnthropicModel(model_id="claude-sonnet-4-20250514", max_tokens=1)
89+
90+
assert model.get_config().get("context_window_limit") == 1_000_000
91+
92+
93+
def test__init__explicit_context_window_limit_not_overridden(anthropic_client):
94+
_ = anthropic_client
95+
96+
model = AnthropicModel(model_id="claude-sonnet-4-20250514", max_tokens=1, context_window_limit=100_000)
97+
98+
assert model.get_config().get("context_window_limit") == 100_000
99+
100+
101+
def test__init__unknown_model_no_context_window_limit(anthropic_client):
102+
_ = anthropic_client
103+
104+
model = AnthropicModel(model_id="unknown-model", max_tokens=1)
105+
106+
assert model.get_config().get("context_window_limit") is None
107+
108+
85109
def test_update_config(model, model_id):
86110
model.update_config(model_id=model_id)
87111

tests/strands/models/test_bedrock.py

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,46 @@ def test__init__context_window_limit(bedrock_client):
296296
assert model.context_window_limit == 200_000
297297

298298

299+
def test__init__auto_populates_context_window_limit(bedrock_client):
300+
_ = bedrock_client
301+
302+
model = BedrockModel(model_id="anthropic.claude-sonnet-4-20250514-v1:0")
303+
304+
assert model.get_config().get("context_window_limit") == 1_000_000
305+
306+
307+
def test__init__auto_populates_context_window_limit_cross_region(bedrock_client):
308+
_ = bedrock_client
309+
310+
model = BedrockModel(model_id="us.anthropic.claude-sonnet-4-6")
311+
312+
assert model.get_config().get("context_window_limit") == 1_000_000
313+
314+
315+
def test__init__auto_populates_context_window_limit_default_model(bedrock_client):
316+
_ = bedrock_client
317+
318+
model = BedrockModel()
319+
320+
assert model.get_config().get("context_window_limit") == 1_000_000
321+
322+
323+
def test__init__explicit_context_window_limit_not_overridden(bedrock_client):
324+
_ = bedrock_client
325+
326+
model = BedrockModel(model_id="anthropic.claude-sonnet-4-20250514-v1:0", context_window_limit=100_000)
327+
328+
assert model.get_config().get("context_window_limit") == 100_000
329+
330+
331+
def test__init__unknown_model_no_context_window_limit(bedrock_client):
332+
_ = bedrock_client
333+
334+
model = BedrockModel(model_id="unknown.model-v1:0")
335+
336+
assert model.get_config().get("context_window_limit") is None
337+
338+
299339
def test_update_config(model, model_id):
300340
model.update_config(model_id=model_id)
301341

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
"""Tests for model metadata lookup tables."""
2+
3+
4+
from strands.models._defaults import get_context_window_limit, resolve_config_metadata
5+
6+
7+
class TestGetContextWindowLimit:
8+
"""Tests for get_context_window_limit."""
9+
10+
def test_known_anthropic_direct_api(self):
11+
assert get_context_window_limit("claude-sonnet-4-6") == 1_000_000
12+
assert get_context_window_limit("claude-opus-4-6") == 1_000_000
13+
assert get_context_window_limit("claude-opus-4-5") == 200_000
14+
assert get_context_window_limit("claude-haiku-4-5") == 200_000
15+
16+
def test_known_bedrock_anthropic(self):
17+
assert get_context_window_limit("anthropic.claude-sonnet-4-6") == 1_000_000
18+
assert get_context_window_limit("anthropic.claude-haiku-4-5-20251001-v1:0") == 200_000
19+
20+
def test_known_bedrock_nova(self):
21+
assert get_context_window_limit("amazon.nova-pro-v1:0") == 300_000
22+
assert get_context_window_limit("amazon.nova-micro-v1:0") == 128_000
23+
24+
def test_known_openai(self):
25+
assert get_context_window_limit("gpt-5.4") == 1_050_000
26+
assert get_context_window_limit("gpt-4o") == 128_000
27+
assert get_context_window_limit("o3") == 200_000
28+
assert get_context_window_limit("o4-mini") == 200_000
29+
30+
def test_known_gemini(self):
31+
assert get_context_window_limit("gemini-2.5-flash") == 1_048_576
32+
assert get_context_window_limit("gemini-2.5-pro") == 1_048_576
33+
34+
def test_strips_bedrock_cross_region_prefix(self):
35+
assert get_context_window_limit("us.anthropic.claude-sonnet-4-6") == 1_000_000
36+
assert get_context_window_limit("global.anthropic.claude-sonnet-4-6") == 1_000_000
37+
assert get_context_window_limit("eu.anthropic.claude-sonnet-4-6") == 1_000_000
38+
assert get_context_window_limit("ap.anthropic.claude-sonnet-4-6") == 1_000_000
39+
40+
def test_strips_any_prefix_as_fallback(self):
41+
# Any prefix before the first dot is stripped if direct lookup fails
42+
assert get_context_window_limit("custom.anthropic.claude-sonnet-4-6") == 1_000_000
43+
44+
def test_unknown_model_returns_none(self):
45+
assert get_context_window_limit("unknown-model-xyz") is None
46+
assert get_context_window_limit("foo.unknown-model-xyz") is None
47+
48+
49+
class TestResolveConfigMetadata:
50+
"""Tests for resolve_config_metadata."""
51+
52+
def test_resolves_context_window_limit(self):
53+
config: dict = {"model_id": "claude-sonnet-4-6"}
54+
result = resolve_config_metadata(config, "claude-sonnet-4-6")
55+
assert result["context_window_limit"] == 1_000_000
56+
57+
def test_preserves_explicit_context_window_limit(self):
58+
config: dict = {"model_id": "claude-sonnet-4-6", "context_window_limit": 100_000}
59+
result = resolve_config_metadata(config, "claude-sonnet-4-6")
60+
assert result["context_window_limit"] == 100_000
61+
62+
def test_returns_original_config_when_explicit(self):
63+
config: dict = {"model_id": "claude-sonnet-4-6", "context_window_limit": 100_000}
64+
result = resolve_config_metadata(config, "claude-sonnet-4-6")
65+
assert result is config
66+
67+
def test_returns_original_config_when_unknown_model(self):
68+
config: dict = {"model_id": "unknown-model"}
69+
result = resolve_config_metadata(config, "unknown-model")
70+
assert result is config
71+
assert "context_window_limit" not in result
72+
73+
def test_returns_new_dict_when_resolved(self):
74+
config: dict = {"model_id": "claude-sonnet-4-6"}
75+
result = resolve_config_metadata(config, "claude-sonnet-4-6")
76+
assert result is not config
77+
assert "context_window_limit" not in config

0 commit comments

Comments
 (0)