diff --git a/src/strands/models/__init__.py b/src/strands/models/__init__.py index ead290a35..9ec51b08d 100644 --- a/src/strands/models/__init__.py +++ b/src/strands/models/__init__.py @@ -3,8 +3,86 @@ This package includes an abstract base Model class along with concrete implementations for specific providers. """ +from typing import TYPE_CHECKING + from . import bedrock, model from .bedrock import BedrockModel from .model import Model -__all__ = ["bedrock", "model", "BedrockModel", "Model"] +# Type checking imports for static analysis +if TYPE_CHECKING: + from .anthropic import AnthropicModel + from .gemini import GeminiModel + from .litellm import LiteLLMModel + from .llamaapi import LlamaAPIModel + from .llamacpp import LlamaCppModel + from .mistral import MistralModel + from .ollama import OllamaModel + from .openai import OpenAIModel + from .sagemaker import SageMakerAIModel + from .writer import WriterModel + +__all__ = [ + "bedrock", + "model", + "BedrockModel", + "Model", + "AnthropicModel", + "GeminiModel", + "LiteLLMModel", + "LlamaAPIModel", + "LlamaCppModel", + "MistralModel", + "OllamaModel", + "OpenAIModel", + "SageMakerAIModel", + "WriterModel", +] + + +def __getattr__(name: str): + """Lazy load model implementations only when accessed. + + This defers the import of optional dependencies until actually needed. + """ + if name == "AnthropicModel": + from .anthropic import AnthropicModel + + return AnthropicModel + if name == "GeminiModel": + from .gemini import GeminiModel + + return GeminiModel + if name == "LiteLLMModel": + from .litellm import LiteLLMModel + + return LiteLLMModel + if name == "LlamaAPIModel": + from .llamaapi import LlamaAPIModel + + return LlamaAPIModel + if name == "LlamaCppModel": + from .llamacpp import LlamaCppModel + + return LlamaCppModel + if name == "MistralModel": + from .mistral import MistralModel + + return MistralModel + if name == "OllamaModel": + from .ollama import OllamaModel + + return OllamaModel + if name == "OpenAIModel": + from .openai import OpenAIModel + + return OpenAIModel + if name == "SageMakerAIModel": + from .sagemaker import SageMakerAIModel + + return SageMakerAIModel + if name == "WriterModel": + from .writer import WriterModel + + return WriterModel + raise AttributeError(f"module '{__name__}' has no attribute '{name}'")