feat: add context_window_limit to model configs (#2176)

Author: opieter-aws

src/strands/models/__init__.py

@@ -7,11 +7,12 @@
 
 from . import bedrock, model
 from .bedrock import BedrockModel
-from .model import CacheConfig, Model
+from .model import BaseModelConfig, CacheConfig, Model
 
 __all__ = [
     "bedrock",
     "model",
+    "BaseModelConfig",
     "BedrockModel",
     "CacheConfig",
     "Model",

src/strands/models/anthropic.py

@@ -8,7 +8,7 @@
 import logging
 import mimetypes
 from collections.abc import AsyncGenerator
-from typing import Any, TypedDict, TypeVar, cast
+from typing import Any, TypeVar, cast
 
 import anthropic
 from pydantic import BaseModel
@@ -21,7 +21,7 @@
 from ..types.streaming import StreamEvent
 from ..types.tools import ToolChoice, ToolChoiceToolDict, ToolSpec
 from ._validation import _has_location_source, validate_config_keys
-from .model import Model
+from .model import BaseModelConfig, Model
 
 logger = logging.getLogger(__name__)
 
@@ -46,7 +46,7 @@ class AnthropicModel(Model):
         "input and output tokens exceed your context limit",
     }
 
-    class AnthropicConfig(TypedDict, total=False):
+    class AnthropicConfig(BaseModelConfig, total=False):
         """Configuration options for Anthropic models.
 
         Attributes:

src/strands/models/bedrock.py

@@ -15,7 +15,7 @@
 from botocore.config import Config as BotocoreConfig
 from botocore.exceptions import ClientError
 from pydantic import BaseModel
-from typing_extensions import TypedDict, Unpack, override
+from typing_extensions import Unpack, override
 
 from strands.types.media import S3Location, SourceLocation
 
@@ -31,7 +31,7 @@
 from ..types.streaming import CitationsDelta, StreamEvent
 from ..types.tools import ToolChoice, ToolSpec
 from ._validation import validate_config_keys
-from .model import CacheConfig, Model
+from .model import BaseModelConfig, CacheConfig, Model
 
 logger = logging.getLogger(__name__)
 
@@ -69,7 +69,7 @@ class BedrockModel(Model):
     - Context window overflow detection
     """
 
-    class BedrockConfig(TypedDict, total=False):
+    class BedrockConfig(BaseModelConfig, total=False):
         """Configuration options for Bedrock models.
 
         Attributes:

src/strands/models/gemini.py

@@ -9,7 +9,7 @@
 import mimetypes
 import secrets
 from collections.abc import AsyncGenerator
-from typing import Any, TypedDict, TypeVar, cast
+from typing import Any, TypeVar, cast
 
 import pydantic
 from google import genai
@@ -20,7 +20,7 @@
 from ..types.streaming import StreamEvent
 from ..types.tools import ToolChoice, ToolSpec
 from ._validation import _has_location_source, validate_config_keys
-from .model import Model
+from .model import BaseModelConfig, Model
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +33,7 @@ class GeminiModel(Model):
     - Docs: https://ai.google.dev/api
     """
 
-    class GeminiConfig(TypedDict, total=False):
+    class GeminiConfig(BaseModelConfig, total=False):
         """Configuration options for Gemini models.
 
         Attributes:

src/strands/models/litellm.py

@@ -6,7 +6,7 @@
 import json
 import logging
 from collections.abc import AsyncGenerator
-from typing import Any, TypedDict, TypeVar, cast
+from typing import Any, TypeVar, cast
 
 import litellm
 from litellm.exceptions import ContextWindowExceededError
@@ -21,6 +21,7 @@
 from ..types.streaming import MetadataEvent, StreamEvent
 from ..types.tools import ToolChoice, ToolSpec, ToolUse
 from ._validation import validate_config_keys
+from .model import BaseModelConfig
 from .openai import OpenAIModel
 
 logger = logging.getLogger(__name__)
@@ -35,7 +36,7 @@
 class LiteLLMModel(OpenAIModel):
     """LiteLLM model provider implementation."""
 
-    class LiteLLMConfig(TypedDict, total=False):
+    class LiteLLMConfig(BaseModelConfig, total=False):
         """Configuration options for LiteLLM models.
 
         Attributes:

src/strands/models/llamaapi.py

@@ -14,14 +14,14 @@
 import llama_api_client
 from llama_api_client import LlamaAPIClient
 from pydantic import BaseModel
-from typing_extensions import TypedDict, Unpack, override
+from typing_extensions import Unpack, override
 
 from ..types.content import ContentBlock, Messages
 from ..types.exceptions import ModelThrottledException
 from ..types.streaming import StreamEvent, Usage
 from ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse
 from ._validation import _has_location_source, validate_config_keys, warn_on_tool_choice_not_supported
-from .model import Model
+from .model import BaseModelConfig, Model
 
 logger = logging.getLogger(__name__)
 
@@ -31,7 +31,7 @@
 class LlamaAPIModel(Model):
     """Llama API model provider implementation."""
 
-    class LlamaConfig(TypedDict, total=False):
+    class LlamaConfig(BaseModelConfig, total=False):
         """Configuration options for Llama API models.
 
         Attributes:

src/strands/models/llamacpp.py

@@ -17,7 +17,6 @@
 from collections.abc import AsyncGenerator
 from typing import (
     Any,
-    TypedDict,
     TypeVar,
     cast,
 )
@@ -31,7 +30,7 @@
 from ..types.streaming import StreamEvent
 from ..types.tools import ToolChoice, ToolSpec
 from ._validation import _has_location_source, validate_config_keys, warn_on_tool_choice_not_supported
-from .model import Model
+from .model import BaseModelConfig, Model
 
 logger = logging.getLogger(__name__)
 
@@ -86,7 +85,7 @@ class LlamaCppModel(Model):
         >>> response = agent(image_content)
     """
 
-    class LlamaCppConfig(TypedDict, total=False):
+    class LlamaCppConfig(BaseModelConfig, total=False):
         """Configuration options for llama.cpp models.
 
         Attributes:

src/strands/models/mistral.py

@@ -11,14 +11,14 @@
 
 import mistralai
 from pydantic import BaseModel
-from typing_extensions import TypedDict, Unpack, override
+from typing_extensions import Unpack, override
 
 from ..types.content import ContentBlock, Messages
 from ..types.exceptions import ModelThrottledException
 from ..types.streaming import StopReason, StreamEvent
 from ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse
 from ._validation import _has_location_source, validate_config_keys, warn_on_tool_choice_not_supported
-from .model import Model
+from .model import BaseModelConfig, Model
 
 logger = logging.getLogger(__name__)
 
@@ -36,7 +36,7 @@ class MistralModel(Model):
     - System prompts
     """
 
-    class MistralConfig(TypedDict, total=False):
+    class MistralConfig(BaseModelConfig, total=False):
         """Configuration parameters for Mistral models.
 
         Attributes:

src/strands/models/model.py

@@ -4,7 +4,7 @@
 import logging
 from collections.abc import AsyncGenerator, AsyncIterable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Literal, TypeVar
+from typing import TYPE_CHECKING, Any, Literal, TypedDict, TypeVar
 
 from pydantic import BaseModel
 
@@ -22,6 +22,17 @@
 T = TypeVar("T", bound=BaseModel)
 
 
+class BaseModelConfig(TypedDict, total=False):
+    """Base configuration shared by all model providers.
+
+    Attributes:
+        context_window_limit: Maximum context window size in tokens for the model.
+            This value represents the total token capacity shared between input and output.
+    """
+
+    context_window_limit: int | None
+
+
 @dataclass
 class CacheConfig:
     """Configuration for prompt caching.
@@ -51,6 +62,16 @@ def stateful(self) -> bool:
         """
         return False
 
+    @property
+    def context_window_limit(self) -> int | None:
+        """Maximum context window size in tokens, or None if not configured."""
+        config = self.get_config()
+        return (
+            config.get("context_window_limit")
+            if isinstance(config, dict)
+            else getattr(config, "context_window_limit", None)
+        )
+
     @abc.abstractmethod
     # pragma: no cover
     def update_config(self, **model_config: Any) -> None:

src/strands/models/ollama.py

@@ -10,13 +10,13 @@
 
 import ollama
 from pydantic import BaseModel
-from typing_extensions import TypedDict, Unpack, override
+from typing_extensions import Unpack, override
 
 from ..types.content import ContentBlock, Messages
 from ..types.streaming import StopReason, StreamEvent
 from ..types.tools import ToolChoice, ToolSpec
 from ._validation import _has_location_source, validate_config_keys, warn_on_tool_choice_not_supported
-from .model import Model
+from .model import BaseModelConfig, Model
 
 logger = logging.getLogger(__name__)
 
@@ -33,7 +33,7 @@ class OllamaModel(Model):
     - Tool/function calling
     """
 
-    class OllamaConfig(TypedDict, total=False):
+    class OllamaConfig(BaseModelConfig, total=False):
         """Configuration parameters for Ollama models.
 
         Attributes: