Langchain-Chatchat/model-providers/model_providers/core/model_manager.py

from typing import IO, Generator, List, Optional, Union, cast

from model_providers.core.entities.provider_configuration import ProviderModelBundle
from model_providers.core.model_runtime.callbacks.base_callback import Callback
from model_providers.core.model_runtime.entities.llm_entities import LLMResult
from model_providers.core.model_runtime.entities.message_entities import (
    PromptMessage,
    PromptMessageTool,
)
from model_providers.core.model_runtime.entities.model_entities import ModelType
from model_providers.core.model_runtime.entities.rerank_entities import RerankResult
from model_providers.core.model_runtime.entities.text_embedding_entities import (
    TextEmbeddingResult,
)
from model_providers.core.model_runtime.model_providers.__base.large_language_model import (
    LargeLanguageModel,
)
from model_providers.core.model_runtime.model_providers.__base.moderation_model import (
    ModerationModel,
)
from model_providers.core.model_runtime.model_providers.__base.rerank_model import (
    RerankModel,
)
from model_providers.core.model_runtime.model_providers.__base.speech2text_model import (
    Speech2TextModel,
)
from model_providers.core.model_runtime.model_providers.__base.text_embedding_model import (
    TextEmbeddingModel,
)
from model_providers.core.model_runtime.model_providers.__base.tts_model import TTSModel
from model_providers.core.provider_manager import ProviderManager
from model_providers.errors.error import ProviderTokenNotInitError


def _fetch_credentials_from_bundle(
    provider_model_bundle: ProviderModelBundle, model: str
) -> dict:
    """
    Fetch credentials from provider model bundle
    :param provider_model_bundle: provider model bundle
    :param model: model name
    :return:
    """
    credentials = provider_model_bundle.configuration.get_current_credentials(
        model_type=provider_model_bundle.model_type_instance.model_type, model=model
    )

    if credentials is None:
        raise ProviderTokenNotInitError(
            f"Model {model} credentials is not initialized."
        )

    return credentials


class ModelInstance:
    """
    Model instance class
    """

    def __init__(self, provider_model_bundle: ProviderModelBundle, model: str) -> None:
        self._provider_model_bundle = provider_model_bundle
        self.model = model
        self.provider = provider_model_bundle.configuration.provider.provider
        self.credentials = _fetch_credentials_from_bundle(provider_model_bundle, model)
        self.model_type_instance = self._provider_model_bundle.model_type_instance

    def invoke_llm(
        self,
        prompt_messages: List[PromptMessage],
        model_parameters: Optional[dict] = None,
        tools: Optional[List[PromptMessageTool]] = None,
        stop: Optional[List[str]] = None,
        stream: bool = True,
        user: Optional[str] = None,
        callbacks: List[Callback] = None,
    ) -> Union[LLMResult, Generator]:
        """
        Invoke large language model

        :param prompt_messages: prompt messages
        :param model_parameters: model parameters
        :param tools: tools for tool calling
        :param stop: stop words
        :param stream: is stream response
        :param user: unique user id
        :param callbacks: callbacks
        :return: full response or stream response chunk generator result
        """
        if not isinstance(self.model_type_instance, LargeLanguageModel):
            raise Exception("Model type instance is not LargeLanguageModel")

        self.model_type_instance = cast(LargeLanguageModel, self.model_type_instance)
        return self.model_type_instance.invoke(
            model=self.model,
            credentials=self.credentials,
            prompt_messages=prompt_messages,
            model_parameters=model_parameters,
            tools=tools,
            stop=stop,
            stream=stream,
            user=user,
            callbacks=callbacks,
        )

    def invoke_text_embedding(
        self, texts: List[str], user: Optional[str] = None
    ) -> TextEmbeddingResult:
        """
        Invoke large language model

        :param texts: texts to embed
        :param user: unique user id
        :return: embeddings result
        """
        if not isinstance(self.model_type_instance, TextEmbeddingModel):
            raise Exception("Model type instance is not TextEmbeddingModel")

        self.model_type_instance = cast(TextEmbeddingModel, self.model_type_instance)
        return self.model_type_instance.invoke(
            model=self.model, credentials=self.credentials, texts=texts, user=user
        )

    def invoke_rerank(
        self,
        query: str,
        docs: List[str],
        score_threshold: Optional[float] = None,
        top_n: Optional[int] = None,
        user: Optional[str] = None,
    ) -> RerankResult:
        """
        Invoke rerank model

        :param query: search query
        :param docs: docs for reranking
        :param score_threshold: score threshold
        :param top_n: top n
        :param user: unique user id
        :return: rerank result
        """
        if not isinstance(self.model_type_instance, RerankModel):
            raise Exception("Model type instance is not RerankModel")

        self.model_type_instance = cast(RerankModel, self.model_type_instance)
        return self.model_type_instance.invoke(
            model=self.model,
            credentials=self.credentials,
            query=query,
            docs=docs,
            score_threshold=score_threshold,
            top_n=top_n,
            user=user,
        )

    def invoke_moderation(self, text: str, user: Optional[str] = None) -> bool:
        """
        Invoke moderation model

        :param text: text to moderate
        :param user: unique user id
        :return: false if text is safe, true otherwise
        """
        if not isinstance(self.model_type_instance, ModerationModel):
            raise Exception("Model type instance is not ModerationModel")

        self.model_type_instance = cast(ModerationModel, self.model_type_instance)
        return self.model_type_instance.invoke(
            model=self.model, credentials=self.credentials, text=text, user=user
        )

    def invoke_speech2text(self, file: IO[bytes], user: Optional[str] = None) -> str:
        """
        Invoke large language model

        :param file: audio file
        :param user: unique user id
        :return: text for given audio file
        """
        if not isinstance(self.model_type_instance, Speech2TextModel):
            raise Exception("Model type instance is not Speech2TextModel")

        self.model_type_instance = cast(Speech2TextModel, self.model_type_instance)
        return self.model_type_instance.invoke(
            model=self.model, credentials=self.credentials, file=file, user=user
        )

    def invoke_tts(
        self,
        content_text: str,
        tenant_id: str,
        voice: str,
        streaming: bool,
        user: Optional[str] = None,
    ) -> str:
        """
        Invoke large language tts model

        :param content_text: text content to be translated
        :param tenant_id: user tenant id
        :param user: unique user id
        :param voice: model timbre
        :param streaming: output is streaming
        :return: text for given audio file
        """
        if not isinstance(self.model_type_instance, TTSModel):
            raise Exception("Model type instance is not TTSModel")

        self.model_type_instance = cast(TTSModel, self.model_type_instance)
        return self.model_type_instance.invoke(
            model=self.model,
            credentials=self.credentials,
            content_text=content_text,
            user=user,
            tenant_id=tenant_id,
            voice=voice,
            streaming=streaming,
        )

    def get_tts_voices(self, language: str) -> list:
        """
        Invoke large language tts model voices

        :param language: tts language
        :return: tts model voices
        """
        if not isinstance(self.model_type_instance, TTSModel):
            raise Exception("Model type instance is not TTSModel")

        self.model_type_instance = cast(TTSModel, self.model_type_instance)
        return self.model_type_instance.get_tts_model_voices(
            model=self.model, credentials=self.credentials, language=language
        )


class ModelManager:
    def __init__(
        self,
        provider_name_to_provider_records_dict: dict,
        provider_name_to_provider_model_records_dict: dict,
    ) -> None:
        self._provider_manager = ProviderManager(
            provider_name_to_provider_records_dict=provider_name_to_provider_records_dict,
            provider_name_to_provider_model_records_dict=provider_name_to_provider_model_records_dict,
        )

    @property
    def provider_manager(self) -> ProviderManager:
        return self._provider_manager

    def get_model_instance(
        self, provider: str, model_type: ModelType, model: str
    ) -> ModelInstance:
        """
        Get model instance
        :param provider: provider name
        :param model_type: model type
        :param model: model name
        :return:
        """
        if not provider:
            return self.get_default_model_instance(model_type)
        provider_model_bundle = self._provider_manager.get_provider_model_bundle(
            provider=provider, model_type=model_type
        )

        return ModelInstance(provider_model_bundle, model)

    def get_default_model_instance(self, model_type: ModelType) -> ModelInstance:
        """
        Get default model instance
        :param model_type: model type
        :return:
        """
        default_model_entity = self._provider_manager.get_default_model(
            model_type=model_type
        )

        if not default_model_entity:
            raise ProviderTokenNotInitError(f"Default model not found for {model_type}")

        return self.get_model_instance(
            provider=default_model_entity.provider.provider,
            model_type=model_type,
            model=default_model_entity.model,
        )