drivers

@define
class AmazonBedrockDriversConfig(DriversConfig):
    session: boto3.Session = field(
        default=Factory(lambda: import_optional_dependency("boto3").Session()),
        kw_only=True,
        metadata={"serializable": False},
    )

    @lazy_property()
    def prompt_driver(self) -> AmazonBedrockPromptDriver:
        return AmazonBedrockPromptDriver(session=self.session, model="anthropic.claude-3-7-sonnet-20250219-v1:0")

    @lazy_property()
    def embedding_driver(self) -> AmazonBedrockTitanEmbeddingDriver:
        return AmazonBedrockTitanEmbeddingDriver(session=self.session, model="amazon.titan-embed-text-v2:0")

    @lazy_property()
    def image_generation_driver(self) -> AmazonBedrockImageGenerationDriver:
        return AmazonBedrockImageGenerationDriver(
            session=self.session,
            model="amazon.titan-image-generator-v2:0",
            image_generation_model_driver=BedrockTitanImageGenerationModelDriver(),
        )

    @lazy_property()
    def vector_store_driver(self) -> LocalVectorStoreDriver:
        return LocalVectorStoreDriver(
            embedding_driver=AmazonBedrockTitanEmbeddingDriver(
                session=self.session, model="amazon.titan-embed-text-v2:0"
            )
        )

@lazy_property()
def embedding_driver(self) -> AmazonBedrockTitanEmbeddingDriver:
    return AmazonBedrockTitanEmbeddingDriver(session=self.session, model="amazon.titan-embed-text-v2:0")

@lazy_property()
def image_generation_driver(self) -> AmazonBedrockImageGenerationDriver:
    return AmazonBedrockImageGenerationDriver(
        session=self.session,
        model="amazon.titan-image-generator-v2:0",
        image_generation_model_driver=BedrockTitanImageGenerationModelDriver(),
    )

@lazy_property()
def prompt_driver(self) -> AmazonBedrockPromptDriver:
    return AmazonBedrockPromptDriver(session=self.session, model="anthropic.claude-3-7-sonnet-20250219-v1:0")

@lazy_property()
def vector_store_driver(self) -> LocalVectorStoreDriver:
    return LocalVectorStoreDriver(
        embedding_driver=AmazonBedrockTitanEmbeddingDriver(
            session=self.session, model="amazon.titan-embed-text-v2:0"
        )
    )

@define
class AnthropicDriversConfig(DriversConfig):
    @lazy_property()
    def prompt_driver(self) -> AnthropicPromptDriver:
        return AnthropicPromptDriver(model="claude-3-7-sonnet-latest")

@lazy_property()
def prompt_driver(self) -> AnthropicPromptDriver:
    return AnthropicPromptDriver(model="claude-3-7-sonnet-latest")

@define
class AzureOpenAiDriversConfig(DriversConfig):
    """Azure OpenAI Drivers Configuration.

    Attributes:
        azure_endpoint: The endpoint for the Azure OpenAI instance.
        azure_ad_token: An optional Azure Active Directory token.
        azure_ad_token_provider: An optional Azure Active Directory token provider.
        api_key: An optional Azure API key.
        prompt_driver: An Azure OpenAI Chat Prompt Driver.
        image_generation_driver: An Azure OpenAI Image Generation Driver.
        embedding_driver: An Azure OpenAI Embedding Driver.
        vector_store_driver: A Local Vector Store Driver.
    """

    azure_endpoint: str = field(kw_only=True, metadata={"serializable": True})
    azure_ad_token: Optional[str] = field(kw_only=True, default=None, metadata={"serializable": False})
    azure_ad_token_provider: Optional[Callable[[], str]] = field(
        kw_only=True,
        default=None,
        metadata={"serializable": False},
    )
    api_key: Optional[str] = field(kw_only=True, default=None, metadata={"serializable": False})

    @lazy_property()
    def prompt_driver(self) -> AzureOpenAiChatPromptDriver:
        return AzureOpenAiChatPromptDriver(
            model="gpt-4.1",
            azure_endpoint=self.azure_endpoint,
            api_key=self.api_key,
            azure_ad_token=self.azure_ad_token,
            azure_ad_token_provider=self.azure_ad_token_provider,
        )

    @lazy_property()
    def embedding_driver(self) -> AzureOpenAiEmbeddingDriver:
        return AzureOpenAiEmbeddingDriver(
            model="text-embedding-3-small",
            azure_endpoint=self.azure_endpoint,
            api_key=self.api_key,
            azure_ad_token=self.azure_ad_token,
            azure_ad_token_provider=self.azure_ad_token_provider,
        )

    @lazy_property()
    def image_generation_driver(self) -> AzureOpenAiImageGenerationDriver:
        return AzureOpenAiImageGenerationDriver(
            model="dall-e-2",
            azure_endpoint=self.azure_endpoint,
            api_key=self.api_key,
            azure_ad_token=self.azure_ad_token,
            azure_ad_token_provider=self.azure_ad_token_provider,
            image_size="512x512",
        )

    @lazy_property()
    def vector_store_driver(self) -> LocalVectorStoreDriver:
        return LocalVectorStoreDriver(
            embedding_driver=AzureOpenAiEmbeddingDriver(
                model="text-embedding-3-small",
                azure_endpoint=self.azure_endpoint,
                api_key=self.api_key,
                azure_ad_token=self.azure_ad_token,
                azure_ad_token_provider=self.azure_ad_token_provider,
            )
        )

    @lazy_property()
    def text_to_speech_driver(self) -> AzureOpenAiTextToSpeechDriver:
        return AzureOpenAiTextToSpeechDriver(
            model="tts",
            azure_endpoint=self.azure_endpoint,
            api_key=self.api_key,
            azure_ad_token=self.azure_ad_token,
            azure_ad_token_provider=self.azure_ad_token_provider,
        )

@lazy_property()
def embedding_driver(self) -> AzureOpenAiEmbeddingDriver:
    return AzureOpenAiEmbeddingDriver(
        model="text-embedding-3-small",
        azure_endpoint=self.azure_endpoint,
        api_key=self.api_key,
        azure_ad_token=self.azure_ad_token,
        azure_ad_token_provider=self.azure_ad_token_provider,
    )

@lazy_property()
def image_generation_driver(self) -> AzureOpenAiImageGenerationDriver:
    return AzureOpenAiImageGenerationDriver(
        model="dall-e-2",
        azure_endpoint=self.azure_endpoint,
        api_key=self.api_key,
        azure_ad_token=self.azure_ad_token,
        azure_ad_token_provider=self.azure_ad_token_provider,
        image_size="512x512",
    )

@lazy_property()
def prompt_driver(self) -> AzureOpenAiChatPromptDriver:
    return AzureOpenAiChatPromptDriver(
        model="gpt-4.1",
        azure_endpoint=self.azure_endpoint,
        api_key=self.api_key,
        azure_ad_token=self.azure_ad_token,
        azure_ad_token_provider=self.azure_ad_token_provider,
    )

@lazy_property()
def text_to_speech_driver(self) -> AzureOpenAiTextToSpeechDriver:
    return AzureOpenAiTextToSpeechDriver(
        model="tts",
        azure_endpoint=self.azure_endpoint,
        api_key=self.api_key,
        azure_ad_token=self.azure_ad_token,
        azure_ad_token_provider=self.azure_ad_token_provider,
    )

@lazy_property()
def vector_store_driver(self) -> LocalVectorStoreDriver:
    return LocalVectorStoreDriver(
        embedding_driver=AzureOpenAiEmbeddingDriver(
            model="text-embedding-3-small",
            azure_endpoint=self.azure_endpoint,
            api_key=self.api_key,
            azure_ad_token=self.azure_ad_token,
            azure_ad_token_provider=self.azure_ad_token_provider,
        )
    )

@define
class BaseDriversConfig(ABC, SerializableMixin):
    _prompt_driver: Optional[BasePromptDriver] = field(
        kw_only=True, default=None, metadata={"serializable": True}, alias="prompt_driver"
    )
    _image_generation_driver: Optional[BaseImageGenerationDriver] = field(
        kw_only=True, default=None, metadata={"serializable": True}, alias="image_generation_driver"
    )
    _embedding_driver: Optional[BaseEmbeddingDriver] = field(
        kw_only=True, default=None, metadata={"serializable": True}, alias="embedding_driver"
    )
    _vector_store_driver: Optional[BaseVectorStoreDriver] = field(
        default=None, kw_only=True, metadata={"serializable": True}, alias="vector_store_driver"
    )
    _conversation_memory_driver: Optional[BaseConversationMemoryDriver] = field(
        default=None, kw_only=True, metadata={"serializable": True}, alias="conversation_memory_driver"
    )
    _text_to_speech_driver: Optional[BaseTextToSpeechDriver] = field(
        default=None, kw_only=True, metadata={"serializable": True}, alias="text_to_speech_driver"
    )
    _audio_transcription_driver: Optional[BaseAudioTranscriptionDriver] = field(
        default=None, kw_only=True, metadata={"serializable": True}, alias="audio_transcription_driver"
    )
    _ruleset_driver: Optional[BaseRulesetDriver] = field(
        default=None, kw_only=True, metadata={"serializable": True}, alias="ruleset_driver"
    )

    _last_drivers_config: Optional[BaseDriversConfig] = field(default=None)

    def __enter__(self) -> Self:
        from griptape.configs import Defaults

        self._last_drivers_config = Defaults.drivers_config

        Defaults.drivers_config = self

        return self

    def __exit__(self, type, value, traceback) -> None:  # noqa: ANN001, A002
        from griptape.configs import Defaults

        if self._last_drivers_config is not None:
            Defaults.drivers_config = self._last_drivers_config

        self._last_drivers_config = None

    @lazy_property()
    @abstractmethod
    def prompt_driver(self) -> BasePromptDriver: ...

    @lazy_property()
    @abstractmethod
    def image_generation_driver(self) -> BaseImageGenerationDriver: ...

    @lazy_property()
    @abstractmethod
    def embedding_driver(self) -> BaseEmbeddingDriver: ...

    @lazy_property()
    @abstractmethod
    def vector_store_driver(self) -> BaseVectorStoreDriver: ...

    @lazy_property()
    @abstractmethod
    def conversation_memory_driver(self) -> BaseConversationMemoryDriver: ...

    @lazy_property()
    @abstractmethod
    def text_to_speech_driver(self) -> BaseTextToSpeechDriver: ...

    @lazy_property()
    @abstractmethod
    def audio_transcription_driver(self) -> BaseAudioTranscriptionDriver: ...

    @lazy_property()
    @abstractmethod
    def ruleset_driver(self) -> BaseRulesetDriver: ...

def __enter__(self) -> Self:
    from griptape.configs import Defaults

    self._last_drivers_config = Defaults.drivers_config

    Defaults.drivers_config = self

    return self

def __exit__(self, type, value, traceback) -> None:  # noqa: ANN001, A002
    from griptape.configs import Defaults

    if self._last_drivers_config is not None:
        Defaults.drivers_config = self._last_drivers_config

    self._last_drivers_config = None

@lazy_property()
@abstractmethod
def audio_transcription_driver(self) -> BaseAudioTranscriptionDriver: ...

@lazy_property()
@abstractmethod
def conversation_memory_driver(self) -> BaseConversationMemoryDriver: ...

@lazy_property()
@abstractmethod
def embedding_driver(self) -> BaseEmbeddingDriver: ...

@lazy_property()
@abstractmethod
def image_generation_driver(self) -> BaseImageGenerationDriver: ...

@lazy_property()
@abstractmethod
def prompt_driver(self) -> BasePromptDriver: ...

@lazy_property()
@abstractmethod
def ruleset_driver(self) -> BaseRulesetDriver: ...

@lazy_property()
@abstractmethod
def text_to_speech_driver(self) -> BaseTextToSpeechDriver: ...

@lazy_property()
@abstractmethod
def vector_store_driver(self) -> BaseVectorStoreDriver: ...

@define
class CohereDriversConfig(DriversConfig):
    api_key: str = field(metadata={"serializable": False}, kw_only=True)

    @lazy_property()
    def prompt_driver(self) -> CoherePromptDriver:
        return CoherePromptDriver(model="command-r", api_key=self.api_key)

    @lazy_property()
    def embedding_driver(self) -> CohereEmbeddingDriver:
        return CohereEmbeddingDriver(
            model="embed-english-v3.0",
            api_key=self.api_key,
            input_type="search_document",
        )

    @lazy_property()
    def vector_store_driver(self) -> LocalVectorStoreDriver:
        return LocalVectorStoreDriver(
            embedding_driver=CohereEmbeddingDriver(
                model="embed-english-v3.0",
                api_key=self.api_key,
                input_type="search_document",
            )
        )

@lazy_property()
def embedding_driver(self) -> CohereEmbeddingDriver:
    return CohereEmbeddingDriver(
        model="embed-english-v3.0",
        api_key=self.api_key,
        input_type="search_document",
    )

@lazy_property()
def prompt_driver(self) -> CoherePromptDriver:
    return CoherePromptDriver(model="command-r", api_key=self.api_key)

@lazy_property()
def vector_store_driver(self) -> LocalVectorStoreDriver:
    return LocalVectorStoreDriver(
        embedding_driver=CohereEmbeddingDriver(
            model="embed-english-v3.0",
            api_key=self.api_key,
            input_type="search_document",
        )
    )

@define
class DriversConfig(BaseDriversConfig):
    @lazy_property()
    def prompt_driver(self) -> BasePromptDriver:
        return DummyPromptDriver()

    @lazy_property()
    def image_generation_driver(self) -> BaseImageGenerationDriver:
        return DummyImageGenerationDriver()

    @lazy_property()
    def embedding_driver(self) -> BaseEmbeddingDriver:
        return DummyEmbeddingDriver()

    @lazy_property()
    def vector_store_driver(self) -> BaseVectorStoreDriver:
        return DummyVectorStoreDriver(embedding_driver=self.embedding_driver)

    @lazy_property()
    def conversation_memory_driver(self) -> BaseConversationMemoryDriver:
        return LocalConversationMemoryDriver()

    @lazy_property()
    def text_to_speech_driver(self) -> BaseTextToSpeechDriver:
        return DummyTextToSpeechDriver()

    @lazy_property()
    def audio_transcription_driver(self) -> BaseAudioTranscriptionDriver:
        return DummyAudioTranscriptionDriver()

    @lazy_property()
    def ruleset_driver(self) -> BaseRulesetDriver:
        return LocalRulesetDriver()

@lazy_property()
def audio_transcription_driver(self) -> BaseAudioTranscriptionDriver:
    return DummyAudioTranscriptionDriver()

@lazy_property()
def conversation_memory_driver(self) -> BaseConversationMemoryDriver:
    return LocalConversationMemoryDriver()

@lazy_property()
def embedding_driver(self) -> BaseEmbeddingDriver:
    return DummyEmbeddingDriver()

@lazy_property()
def image_generation_driver(self) -> BaseImageGenerationDriver:
    return DummyImageGenerationDriver()

@lazy_property()
def prompt_driver(self) -> BasePromptDriver:
    return DummyPromptDriver()

@lazy_property()
def ruleset_driver(self) -> BaseRulesetDriver:
    return LocalRulesetDriver()

@lazy_property()
def text_to_speech_driver(self) -> BaseTextToSpeechDriver:
    return DummyTextToSpeechDriver()

@lazy_property()
def vector_store_driver(self) -> BaseVectorStoreDriver:
    return DummyVectorStoreDriver(embedding_driver=self.embedding_driver)

@define
class GoogleDriversConfig(DriversConfig):
    @lazy_property()
    def prompt_driver(self) -> GooglePromptDriver:
        return GooglePromptDriver(model="gemini-2.0-flash")

    @lazy_property()
    def embedding_driver(self) -> GoogleEmbeddingDriver:
        return GoogleEmbeddingDriver(model="models/embedding-004")

    @lazy_property()
    def vector_store_driver(self) -> LocalVectorStoreDriver:
        return LocalVectorStoreDriver(embedding_driver=GoogleEmbeddingDriver(model="models/embedding-004"))

@lazy_property()
def embedding_driver(self) -> GoogleEmbeddingDriver:
    return GoogleEmbeddingDriver(model="models/embedding-004")

@lazy_property()
def prompt_driver(self) -> GooglePromptDriver:
    return GooglePromptDriver(model="gemini-2.0-flash")

@lazy_property()
def vector_store_driver(self) -> LocalVectorStoreDriver:
    return LocalVectorStoreDriver(embedding_driver=GoogleEmbeddingDriver(model="models/embedding-004"))

@define
class OpenAiDriversConfig(DriversConfig):
    @lazy_property()
    def prompt_driver(self) -> OpenAiChatPromptDriver:
        return OpenAiChatPromptDriver(model="gpt-4.1")

    @lazy_property()
    def image_generation_driver(self) -> OpenAiImageGenerationDriver:
        return OpenAiImageGenerationDriver(model="dall-e-2", image_size="512x512")

    @lazy_property()
    def embedding_driver(self) -> OpenAiEmbeddingDriver:
        return OpenAiEmbeddingDriver(model="text-embedding-3-small")

    @lazy_property()
    def vector_store_driver(self) -> LocalVectorStoreDriver:
        return LocalVectorStoreDriver(embedding_driver=OpenAiEmbeddingDriver(model="text-embedding-3-small"))

    @lazy_property()
    def text_to_speech_driver(self) -> OpenAiTextToSpeechDriver:
        return OpenAiTextToSpeechDriver(model="tts-1")

    @lazy_property()
    def audio_transcription_driver(self) -> OpenAiAudioTranscriptionDriver:
        return OpenAiAudioTranscriptionDriver(model="whisper-1")

@lazy_property()
def audio_transcription_driver(self) -> OpenAiAudioTranscriptionDriver:
    return OpenAiAudioTranscriptionDriver(model="whisper-1")

@lazy_property()
def embedding_driver(self) -> OpenAiEmbeddingDriver:
    return OpenAiEmbeddingDriver(model="text-embedding-3-small")

@lazy_property()
def image_generation_driver(self) -> OpenAiImageGenerationDriver:
    return OpenAiImageGenerationDriver(model="dall-e-2", image_size="512x512")

@lazy_property()
def prompt_driver(self) -> OpenAiChatPromptDriver:
    return OpenAiChatPromptDriver(model="gpt-4.1")

@lazy_property()
def text_to_speech_driver(self) -> OpenAiTextToSpeechDriver:
    return OpenAiTextToSpeechDriver(model="tts-1")

@lazy_property()
def vector_store_driver(self) -> LocalVectorStoreDriver:
    return LocalVectorStoreDriver(embedding_driver=OpenAiEmbeddingDriver(model="text-embedding-3-small"))