eveAI/common/utils/model_utils.py

import os
from typing import Dict, Any, Optional, Tuple

import langcodes
from langchain_core.language_models import BaseChatModel

from common.langchain.llm_metrics_handler import LLMMetricsHandler
from langchain_openai import ChatOpenAI
from langchain_mistralai import ChatMistralAI
from flask import current_app

from common.eveai_model.tracked_mistral_embeddings import TrackedMistralAIEmbeddings
from common.langchain.tracked_transcription import TrackedOpenAITranscription
from common.models.user import Tenant
from config.model_config import MODEL_CONFIG
from common.extensions import cache_manager
from common.models.document import EmbeddingMistral
from common.utils.eveai_exceptions import EveAITenantNotFound, EveAIInvalidEmbeddingModel
from crewai import LLM

embedding_llm_model_cache: Dict[Tuple[str, float], BaseChatModel] = {}
crewai_llm_model_cache: Dict[Tuple[str, float], LLM] = {}
llm_metrics_handler = LLMMetricsHandler()


def create_language_template(template: str, language: str) -> str:
    """
    Replace language placeholder in template with specified language

    Args:
        template: Template string with {language} placeholder
        language: Language code to insert

    Returns:
        str: Template with language placeholder replaced
    """
    try:
        full_language = langcodes.Language.make(language=language)
        language_template = template.replace('{language}', full_language.display_name())
    except ValueError:
        language_template = template.replace('{language}', language)

    return language_template


def replace_variable_in_template(template: str, variable: str, value: str) -> str:
    """
    Replace a variable placeholder in template with specified value

    Args:
        template: Template string with variable placeholder
        variable: Variable placeholder to replace (e.g. "{tenant_context}")
        value: Value to insert

    Returns:
        str: Template with variable placeholder replaced
    """

    modified_template = template.replace(f"{{{variable}}}", value or "")
    return modified_template


def get_embedding_model_and_class(tenant_id, catalog_id, full_embedding_name="mistral.mistral-embed"):
    """
        Retrieve the embedding model and embedding model class to store Embeddings

        Args:
            tenant_id: ID of the tenant
            catalog_id: ID of the catalog
            full_embedding_name: The full name of the embedding model: <provider>.<model>

        Returns:
            embedding_model, embedding_model_class
        """
    embedding_provider, embedding_model_name = full_embedding_name.split('.')

    # Calculate the embedding model to be used
    if embedding_provider == "mistral":
        api_key = current_app.config['MISTRAL_API_KEY']
        embedding_model = TrackedMistralAIEmbeddings(
            model=embedding_model_name
        )
    else:
        raise EveAIInvalidEmbeddingModel(tenant_id, catalog_id)

    # Calculate the Embedding Model Class to be used to store embeddings
    if embedding_model_name == "mistral-embed":
        embedding_model_class = EmbeddingMistral
    else:
        raise EveAIInvalidEmbeddingModel(tenant_id, catalog_id)

    return embedding_model, embedding_model_class


def get_embedding_llm(full_model_name='mistral.mistral-small-latest', temperature=0.3):
    llm = embedding_llm_model_cache.get((full_model_name, temperature))
    if not llm:
        llm_provider, llm_model_name = full_model_name.split('.')
        if llm_provider == "openai":
            llm = ChatOpenAI(
                api_key=current_app.config['OPENAI_API_KEY'],
                model=llm_model_name,
                temperature=temperature,
                callbacks=[llm_metrics_handler]
            )
        elif llm_provider == "mistral":
            llm = ChatMistralAI(
                api_key=current_app.config['MISTRAL_API_KEY'],
                model=llm_model_name,
                temperature=temperature,
                callbacks=[llm_metrics_handler]
            )
        embedding_llm_model_cache[(full_model_name, temperature)] = llm

    return llm


def get_crewai_llm(full_model_name='mistral.mistral-large-latest', temperature=0.3):
    llm = crewai_llm_model_cache.get((full_model_name, temperature))
    if not llm:
        llm_provider, llm_model_name = full_model_name.split('.')
        crew_full_model_name = f"{llm_provider}/{llm_model_name}"
        api_key = None
        if llm_provider == "openai":
            api_key = current_app.config['OPENAI_API_KEY']
        elif llm_provider == "mistral":
            api_key = current_app.config['MISTRAL_API_KEY']

        llm = LLM(
            model=crew_full_model_name,
            temperature=temperature,
            api_key=api_key
        )
        crewai_llm_model_cache[(full_model_name, temperature)] = llm

    return llm


def process_pdf():
    full_model_name = 'mistral-ocr-latest'


def get_template(template_name: str, version: Optional[str] = "1.0", temperature: float = 0.3) -> tuple[
    Any, BaseChatModel | None | ChatOpenAI | ChatMistralAI]:
    """
    Get a prompt template
    """
    prompt = cache_manager.prompts_config_cache.get_config(template_name, version)
    if "llm_model" in prompt:
        llm = get_embedding_llm(full_model_name=prompt["llm_model"], temperature=temperature)
    else:
        llm = get_embedding_llm(temperature=temperature)

    return prompt["content"], llm


def get_transcription_model(model_name: str = "whisper-1") -> TrackedOpenAITranscription:
    """
    Get a transcription model instance
    """
    api_key = os.getenv('OPENAI_API_KEY')
    return TrackedOpenAITranscription(
        api_key=api_key,
        model=model_name
    )