neollm/llm/model_name/aws_model_name.py

from typing import Type

from neollm.llm.llm.claude.aws_claude import (
    AWSClaude3Haiku20240307,
    AWSClaude3Opus20240229,
    AWSClaude3Sonnet20240229,
    AWSClaude35Haiku20241022,
    AWSClaude35Sonnet20240620,
    AWSClaude35SonnetV220241022,
    AWSClaudeLLM,
)
from neollm.llm.model_name._abstract_model_name import AbstractModelName
from neollm.llm.model_name.utils import normalize_claude_model_name
from neollm.llm.platform import Platform
from neollm.types import ClientSettings


class AWSModelName(AbstractModelName):
    # Claude ------------------------------------------------
    # claude 3.5
    CLAUDE_3_5_SONNET_V2_20241022 = "anthropic.claude-3-5-sonnet-20241022-v2:0"
    CLAUDE_3_5_SONNET_20240620 = "anthropic.claude-3-5-sonnet-20240620-v1:0"
    CLAUDE_3_5_HAIKU_20241022 = "anthropic.claude-3-5-haiku-20241022-v1:0"
    # claude 3
    CLAUDE_3_OPUS_20240229 = "anthropic.claude-3-opus-20240229-v1:0"
    CLAUDE_3_SONNET_20240229 = "anthropic.claude-3-sonnet-20240229-v1:0"
    CLAUDE_3_HAIKU_20240307 = "anthropic.claude-3-haiku-20240307-v1:0"

    @classmethod
    def _missing_(cls, value: object) -> str | None:
        model_name = str(value)
        # 日付がない場合は追加, 他のプラットフォーム用のmodel_nameをAWSに合わせる
        model_name = normalize_claude_model_name(model_name, platform=Platform.AWS)
        for member in cls:
            if model_name == member.value:
                return member
        return None

    def to_llm(self, client_settings: ClientSettings, model_name: str | None = None) -> AWSClaudeLLM:
        llm_class_map: dict["AWSModelName", Type[AWSClaudeLLM]] = {
            # Claude
            AWSModelName.CLAUDE_3_5_SONNET_V2_20241022: AWSClaude35SonnetV220241022,
            AWSModelName.CLAUDE_3_5_SONNET_20240620: AWSClaude35Sonnet20240620,
            AWSModelName.CLAUDE_3_5_HAIKU_20241022: AWSClaude35Haiku20241022,
            AWSModelName.CLAUDE_3_OPUS_20240229: AWSClaude3Opus20240229,
            AWSModelName.CLAUDE_3_SONNET_20240229: AWSClaude3Sonnet20240229,
            AWSModelName.CLAUDE_3_HAIKU_20240307: AWSClaude3Haiku20240307,
        }
        return llm_class_map[self](client_settings)