Model Factory

The ModelFactory class provides a centralized way to create and manage model instances in the Bedrock Swarm framework. It handles model registration, configuration, and instantiation.

Class Documentation

Factory for creating Bedrock model implementations.

Functions

create_model(model_id: str) -> BedrockModel classmethod

Create a model implementation for the given model ID.

PARAMETER	DESCRIPTION
model_id	The Bedrock model ID TYPE: str

RETURNS	DESCRIPTION
BedrockModel	An instance of the appropriate model implementation

RAISES	DESCRIPTION
ValueError	If the model ID is not supported

Source code in src/bedrock_swarm/models/factory.py

@classmethod
def create_model(cls, model_id: str) -> BedrockModel:
    """Create a model implementation for the given model ID.

    Args:
        model_id: The Bedrock model ID

    Returns:
        An instance of the appropriate model implementation

    Raises:
        ValueError: If the model ID is not supported
    """
    # Find matching model family
    family = next(
        (f for f in cls._model_registry.keys() if model_id.startswith(f)), None
    )

    if not family:
        supported = ", ".join(cls._model_registry.keys())
        raise ValueError(
            f"Unsupported model family. Model ID must start with one of: {supported}"
        )

    # Extract version (everything after the family name and a hyphen)
    version = model_id[len(family) + 1 :]

    # Get model implementation
    family_registry = cls._model_registry[family]
    if version not in family_registry:
        versions = ", ".join(family_registry.keys())
        raise ValueError(
            f"Unsupported version '{version}' for model family '{family}'. "
            f"Supported versions: {versions}"
        )

    # Create model instance with its configuration
    model_info = family_registry[version]
    model = model_info["class"](model_id)
    model.set_config(model_info["config"])
    return model

register_model(family: str, version: str, model_class: Type[BedrockModel], config: Dict[str, Any]) -> None classmethod

Register a new model implementation.

PARAMETER	DESCRIPTION
family	Model family (e.g., "us.anthropic.claude-3-5-sonnet") TYPE: str
version	Model version (e.g., "20241022-v2:0") TYPE: str
model_class	Model implementation class TYPE: Type[BedrockModel]
config	Model configuration (max_tokens, default_tokens, etc.) TYPE: Dict[str, Any]

Source code in src/bedrock_swarm/models/factory.py

@classmethod
def register_model(
    cls,
    family: str,
    version: str,
    model_class: Type[BedrockModel],
    config: Dict[str, Any],
) -> None:
    """Register a new model implementation.

    Args:
        family: Model family (e.g., "us.anthropic.claude-3-5-sonnet")
        version: Model version (e.g., "20241022-v2:0")
        model_class: Model implementation class
        config: Model configuration (max_tokens, default_tokens, etc.)
    """
    if family not in cls._model_registry:
        cls._model_registry[family] = {}
    cls._model_registry[family][version] = {"class": model_class, "config": config}

get_supported_models() -> Dict[str, Dict[str, Dict[str, Any]]] classmethod

Get all supported models.

RETURNS	DESCRIPTION
Dict[str, Dict[str, Dict[str, Any]]]	Dictionary of supported model families, versions, and their configurations

Source code in src/bedrock_swarm/models/factory.py

@classmethod
def get_supported_models(cls) -> Dict[str, Dict[str, Dict[str, Any]]]:
    """Get all supported models.

    Returns:
        Dictionary of supported model families, versions, and their configurations
    """
    return cls._model_registry.copy()

Model Registry

The model factory maintains a registry of available models and their configurations:

BEDROCK_MODEL_REGISTRY = {
    "anthropic.claude-3-sonnet-20240229-v1:0": {
        "class": "ClaudeModel",
        "max_tokens": 200000
    },
    "anthropic.claude-v2:1": {
        "class": "ClaudeModel",
        "max_tokens": 100000
    },
    "amazon.titan-text-express-v1": {
        "class": "TitanModel",
        "max_tokens": 8000
    },
    "amazon.titan-text-lite-v1": {
        "class": "TitanModel",
        "max_tokens": 4000
    }
}

Usage Example

from bedrock_swarm.models import ModelFactory

# Create a model instance
model = ModelFactory.create_model("anthropic.claude-3-sonnet-20240229-v1:0")

# Use the model
request = {
    "prompt": "Hello, how can I help you today?",
    "temperature": 0.7
}

async for chunk in model.invoke(request):
    print(chunk)

Error Handling

The factory includes error handling for:

1. Invalid model IDs
2. Missing model configurations
3. Model initialization errors

Implementation Details

The factory implementation includes:

1. Model registration system
2. Configuration validation
3. Dynamic model instantiation
4. Error handling and logging
5. Model configuration management