Philosophy-RAG-demo/generic_rag/backend/models.py

import logging

from generic_rag.parsers.config import AppSettings, ChatBackend, EmbeddingBackend

# Langchain imports
from langchain_core.embeddings import Embeddings
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_aws import BedrockEmbeddings, ChatBedrock  # Import ChatBedrock
from langchain_google_vertexai import VertexAIEmbeddings, ChatVertexAI  # Import ChatVertexAI
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings, ChatOpenAI, OpenAIEmbeddings  # Import ChatOpenAI

logger = logging.getLogger(__name__)


def get_chat_model(settings: AppSettings) -> BaseChatModel:
    """
    Initializes and returns a chat model based on the backend type and configuration.

    Args:
        settings: The loaded AppSettings object containing configurations.

    Returns:
        An instance of BaseChatModel.

    Raises:
        ValueError: If the backend type is unknown or required configuration is missing.
    """
    logger.info(f"Initializing chat model for backend: {settings.chat_backend.value}")

    if settings.chat_backend == ChatBackend.azure:
        if not settings.azure:
            raise ValueError("Azure chat backend selected, but 'azure' configuration section is missing in config.")
        if (
            not settings.azure.llm_endpoint
            or not settings.azure.llm_deployment_name
            or not settings.azure.llm_api_version
        ):
            raise ValueError(
                "Azure configuration requires 'llm_endpoint', 'llm_deployment_name', and 'llm_api_version'."
            )
        return AzureChatOpenAI(
            azure_endpoint=settings.azure.llm_endpoint,
            azure_deployment=settings.azure.llm_deployment_name,
            openai_api_version=settings.azure.llm_api_version,
            openai_api_key=settings.azure.openai_api_key.get_secret_value() if settings.azure.openai_api_key else None,
        )

    if settings.chat_backend == ChatBackend.openai:
        if not settings.openai:
            raise ValueError("OpenAI chat backend selected, but 'openai' configuration section is missing.")
        if not settings.openai.api_key or not settings.openai.chat_model:
            raise ValueError("OpenAI configuration requires 'api_key' and 'chat_model'.")
        logger.info(f"Using OpenAI model: {model_name}")
        return ChatOpenAI(model=settings.openai.chat_model, openai_api_key=settings.openai.api_key.get_secret_value())

    if settings.chat_backend == ChatBackend.google_vertex:
        if not settings.google_vertex:
            raise ValueError(
                "Google Vertex chat backend selected, but 'google_vertex' configuration section is missing."
            )
        if settings.google_vertex.chat_model:
            model_name = settings.google.chat_model
        logger.info(f"Using Google Vertex model: {model_name}")
        return ChatVertexAI(
            model_name=settings.google_vertex.chat_model,
            project=settings.google_vertex.project_id,
            location=settings.google_vertex.location,
        )

    if settings.chat_backend == ChatBackend.aws:
        if not settings.aws:
            raise ValueError("AWS Bedrock chat backend selected, but 'aws' configuration section is missing.")
        model_name = "anthropic.claude-v2"  # Example default
        if hasattr(settings.aws, "chat_model") and settings.aws.chat_model:
            model_name = settings.aws.chat_model
        logger.info(f"Using AWS Bedrock model: {model_name}")
        return ChatBedrock(
            model_id=model_name,
            region_name=settings.aws.region_name,
        )

    if settings.chat_backend == ChatBackend.local:
        if not settings.local or not settings.local.chat_model:
            raise ValueError("Local chat backend selected, but 'local.chat_model' is missing in config.")
        logger.info(f"Using Local Ollama model: {settings.local.chat_model}")
        # Base URL can also be configured, e.g., base_url=config.local.ollama_base_url
        return ChatOllama(model=settings.local.chat_model)

    # This should not be reached if all Enum members are handled
    raise ValueError(f"Unknown or unhandled chat backend type: {settings.chat_backend}")


def get_embedding_model(settings: AppSettings) -> Embeddings:
    """
    Initializes and returns an embedding model based on the backend type and configuration.

    Args:
        settings: The loaded AppSettings object containing configurations.

    Returns:
        An instance of Embeddings.

    Raises:
        ValueError: If the backend type is unknown or required configuration is missing.
    """
    logger.info(f"Initializing embedding model for backend: {settings.emb_backend.value}")

    if settings.emb_backend == EmbeddingBackend.azure:
        if not settings.azure:
            raise ValueError("Azure embedding backend selected, but 'azure' configuration section is missing.")
        if (
            not settings.azure.emb_endpoint
            or not settings.azure.emb_deployment_name
            or not settings.azure.emb_api_version
        ):
            raise ValueError(
                "Azure configuration requires 'emb_endpoint', 'emb_deployment_name', and 'emb_api_version'."
            )
        return AzureOpenAIEmbeddings(
            azure_endpoint=settings.azure.emb_endpoint,
            azure_deployment=settings.azure.emb_deployment_name,
            openai_api_version=settings.azure.emb_api_version,
            openai_api_key=settings.azure.openai_api_key.get_secret_value() if settings.azure.openai_api_key else None,
        )

    if settings.emb_backend == EmbeddingBackend.openai:
        if not settings.openai:
            raise ValueError("OpenAI embedding backend selected, but 'openai' configuration section is missing.")
        if not settings.openai.api_key:
            raise ValueError("OpenAI configuration requires 'api_key'.")
        model_name = "text-embedding-ada-002"  # Example default
        if hasattr(settings.openai, "emb_model") and settings.openai.emb_model:
            model_name = settings.openai.emb_model
        logger.info(f"Using OpenAI embedding model: {model_name}")
        return OpenAIEmbeddings(model=model_name, openai_api_key=settings.openai.api_key.get_secret_value())

    if settings.emb_backend == EmbeddingBackend.google_vertex:
        if not settings.google_vertex:
            raise ValueError(
                "Google Vertex embedding backend selected, but 'google_vertex' configuration section is missing."
            )
        model_name = "textembedding-gecko@001"  # Example default
        if (
            not settings.google_vertex.emb_model
            or not settings.google_vertex.project_id
            or not settings.google_vertex.location
        ):
            raise ValueError("Google Vertex configuration requires 'emb_model', 'project_id', and 'location'.")
        logger.info(f"Using Google Vertex embedding model: {model_name}")
        return VertexAIEmbeddings(
            model_name=settings.google_vertex.emb_model,
            project=settings.google_vertex.project_id,
            location=settings.google_vertex.location,
        )

    if settings.emb_backend == EmbeddingBackend.aws:
        if not settings.aws:
            raise ValueError("AWS Bedrock embedding backend selected, but 'aws' configuration section is missing.")
        model_name = "amazon.titan-embed-text-v1"  # Example default
        if hasattr(settings.aws, "emb_model") and settings.aws.emb_model:
            model_name = settings.aws.emb_model
        logger.info(f"Using AWS Bedrock embedding model: {model_name}")
        return BedrockEmbeddings(model_id=model_name, region_name=settings.aws.region_name)

    if settings.emb_backend == EmbeddingBackend.local:
        if not settings.local or not settings.local.emb_model:
            raise ValueError("Local embedding backend selected, but 'local.emb_model' is missing in config.")
        logger.info(f"Using Local Ollama embedding model: {settings.local.emb_model}")
        return OllamaEmbeddings(model=settings.local.emb_model)

    if settings.emb_backend == EmbeddingBackend.huggingface:
        if not settings.huggingface or not settings.huggingface.emb_model:
            if settings.local and settings.local.emb_model:
                logger.warning(
                    "HuggingFace backend selected, but 'huggingface.emb_model' missing. Using 'local.emb_model'."
                )
                model_name = settings.local.emb_model
            else:
                raise ValueError(
                    "HuggingFace embedding backend selected, but 'huggingface.emb_model' (or 'local.emb_model') is missing in config."
                )
        else:
            model_name = settings.huggingface.emb_model

        logger.info(f"Using HuggingFace embedding model: {model_name}")
        return HuggingFaceEmbeddings(model_name=model_name)

    raise ValueError(f"Unknown or unhandled embedding backend type: {settings.emb_backend}")