dify/api/core/model_manager.py

import logging
from collections.abc import Callable, Generator, Iterable, Mapping, Sequence
from copy import deepcopy
from typing import IO, TYPE_CHECKING, Any, Literal, Optional, ParamSpec, TypeVar, cast, overload, override
from uuid import UUID

from configs import dify_config
from core.credit_usage import (
    CreditUsageAppType,
    CreditUsageAppTypeInput,
    CreditUsageCreatedBy,
    CreditUsageCreatedByInput,
    normalize_credit_usage_app_type,
    normalize_credit_usage_created_by,
)
from core.entities import PluginCredentialType
from core.entities.embedding_type import EmbeddingInputType
from core.entities.provider_configuration import ProviderConfiguration, ProviderModelBundle
from core.entities.provider_entities import ModelLoadBalancingConfiguration
from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError
from core.model_billing_profile import ModelBillingProfileService
from core.model_context import get_credit_usage_metadata
from core.plugin.impl.model_runtime_factory import create_plugin_provider_manager
from core.provider_manager import ProviderManager
from extensions.ext_redis import redis_client
from graphon.model_runtime.callbacks.base_callback import Callback
from graphon.model_runtime.entities.llm_entities import LLMResult, LLMUsage
from graphon.model_runtime.entities.message_entities import PromptMessage, PromptMessageTool
from graphon.model_runtime.entities.model_entities import AIModelEntity, ModelFeature, ModelType
from graphon.model_runtime.entities.rerank_entities import MultimodalRerankInput, RerankResult
from graphon.model_runtime.entities.text_embedding_entities import EmbeddingResult
from graphon.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeConnectionError, InvokeRateLimitError
from graphon.model_runtime.model_providers.base.large_language_model import LargeLanguageModel
from graphon.model_runtime.model_providers.base.moderation_model import ModerationModel
from graphon.model_runtime.model_providers.base.rerank_model import RerankModel
from graphon.model_runtime.model_providers.base.speech2text_model import Speech2TextModel
from graphon.model_runtime.model_providers.base.text_embedding_model import TextEmbeddingModel
from graphon.model_runtime.model_providers.base.tts_model import TTSModel
from models.provider import ProviderType

logger = logging.getLogger(__name__)
if TYPE_CHECKING:
    from core.app.llm.message_billing import ClassicMessageBilling
P = ParamSpec("P")
R = TypeVar("R")


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

    def __init__(
        self,
        provider_model_bundle: ProviderModelBundle,
        model: str,
        credentials: dict | None = None,
        request_metadata: Mapping[str, object] | None = None,
    ) -> None:
        self.provider_model_bundle = provider_model_bundle
        self.model_name = model
        self.provider = provider_model_bundle.configuration.provider.provider
        if credentials is None:
            credentials = self._fetch_credentials_from_bundle(provider_model_bundle, model)
        self.credentials = credentials
        self._message_billing: ClassicMessageBilling | None = None
        self._request_metadata = dict(request_metadata) if request_metadata else None
        # Runtime LLM invocation fields.
        self.parameters: Mapping[str, Any] = {}
        self.stop: Sequence[str] = ()
        self.model_type_instance = self.provider_model_bundle.model_type_instance
        self.load_balancing_manager = self._get_load_balancing_manager(
            configuration=provider_model_bundle.configuration,
            model_type=provider_model_bundle.model_type_instance.model_type,
            model=model,
            credentials=self.credentials,
        )

    def _resolve_request_metadata(self, request_metadata: Mapping[str, object] | None) -> Mapping[str, object] | None:
        bound_request_metadata = self._request_metadata
        if bound_request_metadata is None:
            if request_metadata is None:
                return get_credit_usage_metadata()
            return request_metadata
        if request_metadata is None:
            return bound_request_metadata
        return {**bound_request_metadata, **request_metadata}

    def get_model_schema(self) -> AIModelEntity:
        """Return the resolved schema for the current model instance."""
        model_schema = self.model_type_instance.get_model_schema(self.model_name, self.credentials)
        if model_schema is None:
            raise ValueError(f"model schema not found for {self.model_name}")
        return model_schema

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

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

        return credentials

    @staticmethod
    def _get_load_balancing_manager(
        configuration: ProviderConfiguration, model_type: ModelType, model: str, credentials: dict[str, Any]
    ) -> Optional["LBModelManager"]:
        """
        Get load balancing model credentials
        :param configuration: provider configuration
        :param model_type: model type
        :param model: model name
        :param credentials: model credentials
        :return:
        """
        if configuration.model_settings and configuration.using_provider_type == ProviderType.CUSTOM:
            current_model_setting = None
            # check if model is disabled by admin
            for model_setting in configuration.model_settings:
                if model_setting.model_type == model_type and model_setting.model == model:
                    current_model_setting = model_setting
                    break

            # check if load balancing is enabled
            if current_model_setting and current_model_setting.load_balancing_configs:
                # use load balancing proxy to choose credentials
                lb_model_manager = LBModelManager(
                    tenant_id=configuration.tenant_id,
                    provider=configuration.provider.provider,
                    model_type=model_type,
                    model=model,
                    load_balancing_configs=current_model_setting.load_balancing_configs,
                    managed_credentials=credentials if configuration.custom_configuration.provider else None,
                )

                return lb_model_manager

        return None

    @overload
    def invoke_llm(
        self,
        prompt_messages: Sequence[PromptMessage],
        model_parameters: dict[str, Any] | None = None,
        tools: Sequence[PromptMessageTool] | None = None,
        stop: list[str] | None = None,
        stream: Literal[True] = True,
        callbacks: list[Callback] | None = None,
        request_metadata: Mapping[str, object] | None = None,
    ) -> Generator: ...

    @overload
    def invoke_llm(
        self,
        prompt_messages: list[PromptMessage],
        model_parameters: dict[str, Any] | None = None,
        tools: Sequence[PromptMessageTool] | None = None,
        stop: list[str] | None = None,
        stream: Literal[False] = False,
        callbacks: list[Callback] | None = None,
        request_metadata: Mapping[str, object] | None = None,
    ) -> LLMResult: ...

    @overload
    def invoke_llm(
        self,
        prompt_messages: list[PromptMessage],
        model_parameters: dict[str, Any] | None = None,
        tools: Sequence[PromptMessageTool] | None = None,
        stop: list[str] | None = None,
        stream: bool = True,
        callbacks: list[Callback] | None = None,
        request_metadata: Mapping[str, object] | None = None,
    ) -> LLMResult | Generator: ...

    def invoke_llm(
        self,
        prompt_messages: Sequence[PromptMessage],
        model_parameters: dict[str, Any] | None = None,
        tools: Sequence[PromptMessageTool] | None = None,
        stop: Sequence[str] | None = None,
        stream: bool = True,
        callbacks: list[Callback] | None = None,
        request_metadata: Mapping[str, object] | None = None,
    ) -> 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 callbacks: callbacks
        :param request_metadata: optional request metadata
        :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")
        request_metadata = self._resolve_request_metadata(request_metadata)
        return cast(
            LLMResult | Generator,
            self._round_robin_invoke(
                self.model_type_instance.invoke,
                model=self.model_name,
                credentials=self.credentials,
                prompt_messages=list(prompt_messages),
                model_parameters=model_parameters,
                tools=list(tools) if tools else None,
                stop=list(stop) if stop else None,
                stream=stream,
                callbacks=callbacks,
                request_metadata=request_metadata,
            ),
        )

    def get_llm_num_tokens(
        self, prompt_messages: Sequence[PromptMessage], tools: Sequence[PromptMessageTool] | None = None
    ) -> int:
        """
        Get number of tokens for llm

        :param prompt_messages: prompt messages
        :param tools: tools for tool calling
        :return:
        """
        if not isinstance(self.model_type_instance, LargeLanguageModel):
            raise Exception("Model type instance is not LargeLanguageModel")
        return self._round_robin_invoke(
            self.model_type_instance.get_num_tokens,
            model=self.model_name,
            credentials=self.credentials,
            prompt_messages=list(prompt_messages),
            tools=list(tools) if tools else None,
        )

    def invoke_text_embedding(
        self,
        texts: list[str],
        input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> EmbeddingResult:
        """
        Invoke large language model

        :param texts: texts to embed
        :param input_type: input type
        :return: embeddings result
        """
        if not isinstance(self.model_type_instance, TextEmbeddingModel):
            raise Exception("Model type instance is not TextEmbeddingModel")
        return self._round_robin_invoke(
            self.model_type_instance.invoke,
            model=self.model_name,
            credentials=self.credentials,
            texts=texts,
            input_type=input_type,
            request_metadata=self._resolve_request_metadata(request_metadata),
        )

    def invoke_multimodal_embedding(
        self,
        multimodel_documents: list[dict],
        input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> EmbeddingResult:
        """
        Invoke large language model

        :param multimodel_documents: multimodel documents to embed
        :param input_type: input type
        :return: embeddings result
        """
        if not isinstance(self.model_type_instance, TextEmbeddingModel):
            raise Exception("Model type instance is not TextEmbeddingModel")
        return self._round_robin_invoke(
            self.model_type_instance.invoke,
            model=self.model_name,
            credentials=self.credentials,
            multimodel_documents=multimodel_documents,
            input_type=input_type,
            request_metadata=self._resolve_request_metadata(request_metadata),
        )

    def get_text_embedding_num_tokens(self, texts: list[str]) -> list[int]:
        """
        Get number of tokens for text embedding

        :param texts: texts to embed
        :return:
        """
        if not isinstance(self.model_type_instance, TextEmbeddingModel):
            raise Exception("Model type instance is not TextEmbeddingModel")
        return self._round_robin_invoke(
            self.model_type_instance.get_num_tokens,
            model=self.model_name,
            credentials=self.credentials,
            texts=texts,
        )

    def invoke_rerank(
        self,
        query: str,
        docs: list[str],
        score_threshold: float | None = None,
        top_n: int | None = None,
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> RerankResult:
        """
        Invoke rerank model

        :param query: search query
        :param docs: docs for reranking
        :param score_threshold: score threshold
        :param top_n: top n
        :return: rerank result
        """
        if not isinstance(self.model_type_instance, RerankModel):
            raise Exception("Model type instance is not RerankModel")
        return self._round_robin_invoke(
            self.model_type_instance.invoke,
            model=self.model_name,
            credentials=self.credentials,
            query=query,
            docs=docs,
            score_threshold=score_threshold,
            top_n=top_n,
            request_metadata=self._resolve_request_metadata(request_metadata),
        )

    def invoke_multimodal_rerank(
        self,
        query: MultimodalRerankInput,
        docs: list[MultimodalRerankInput],
        score_threshold: float | None = None,
        top_n: int | None = None,
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> RerankResult:
        """
        Invoke rerank model

        :param query: search query
        :param docs: docs for reranking
        :param score_threshold: score threshold
        :param top_n: top n
        :return: rerank result
        """
        if not isinstance(self.model_type_instance, RerankModel):
            raise Exception("Model type instance is not RerankModel")
        return self._round_robin_invoke(
            self.model_type_instance.invoke_multimodal_rerank,
            model=self.model_name,
            credentials=self.credentials,
            query=query,
            docs=docs,
            score_threshold=score_threshold,
            top_n=top_n,
            request_metadata=self._resolve_request_metadata(request_metadata),
        )

    def invoke_moderation(self, text: str, *, request_metadata: Mapping[str, object] | None = None) -> bool:
        """
        Invoke moderation model

        :param text: text to moderate
        :return: false if text is safe, true otherwise
        """
        if not isinstance(self.model_type_instance, ModerationModel):
            raise Exception("Model type instance is not ModerationModel")
        return self._round_robin_invoke(
            self.model_type_instance.invoke,
            model=self.model_name,
            credentials=self.credentials,
            text=text,
            request_metadata=self._resolve_request_metadata(request_metadata),
        )

    def invoke_speech2text(self, file: IO[bytes], *, request_metadata: Mapping[str, object] | None = None) -> str:
        """
        Invoke large language model

        :param file: audio file
        :return: text for given audio file
        """
        if not isinstance(self.model_type_instance, Speech2TextModel):
            raise Exception("Model type instance is not Speech2TextModel")
        return self._round_robin_invoke(
            self.model_type_instance.invoke,
            model=self.model_name,
            credentials=self.credentials,
            file=file,
            request_metadata=self._resolve_request_metadata(request_metadata),
        )

    def invoke_tts(
        self,
        content_text: str,
        voice: str = "",
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> Iterable[bytes]:
        """
        Invoke a text-to-speech model.

        :param content_text: text content to be translated
        :param voice: model timbre
        :return: audio byte stream
        """
        if not isinstance(self.model_type_instance, TTSModel):
            raise Exception("Model type instance is not TTSModel")
        chunks = self._round_robin_invoke(
            self.model_type_instance.invoke,
            model=self.model_name,
            credentials=self.credentials,
            content_text=content_text,
            voice=voice,
            request_metadata=self._resolve_request_metadata(request_metadata),
        )
        return (chunk.data for chunk in chunks)

    def _round_robin_invoke(self, function: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
        """
        Round-robin invoke
        :param function: function to invoke
        :param args: function args
        :param kwargs: function kwargs
        :return:
        """
        message_billing = self._message_billing
        if message_billing is not None:
            message_billing.check_active()
        if not self.load_balancing_manager:
            return function(*args, **kwargs)

        last_exception: InvokeRateLimitError | InvokeAuthorizationError | InvokeConnectionError | None = None
        while True:
            lb_config = self.load_balancing_manager.fetch_next()
            if not lb_config:
                if not last_exception:
                    raise ProviderTokenNotInitError("Model credentials is not initialized.")
                else:
                    raise last_exception

            # Additional policy compliance check as fallback (in case fetch_next didn't catch it)
            try:
                from core.helper.credential_utils import runtime_check_credential_policy_compliance

                if lb_config.credential_id:
                    runtime_check_credential_policy_compliance(
                        credential_id=lb_config.credential_id,
                        provider=self.provider,
                        credential_type=PluginCredentialType.MODEL,
                    )
            except Exception as e:
                logger.warning(
                    "Load balancing config %s failed policy compliance check in round-robin: %s", lb_config.id, str(e)
                )
                self.load_balancing_manager.cooldown(lb_config, expire=60)
                continue

            try:
                kwargs["credentials"] = lb_config.credentials
                return function(*args, **kwargs)
            except InvokeRateLimitError as e:
                # expire in 60 seconds
                self.load_balancing_manager.cooldown(lb_config, expire=60)
                last_exception = e
                continue
            except (InvokeAuthorizationError, InvokeConnectionError) as e:
                # expire in 10 seconds
                self.load_balancing_manager.cooldown(lb_config, expire=10)
                last_exception = e
                continue
            except Exception as e:
                raise e

    def get_tts_voices(self, language: str | None = None):
        """
        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")
        return self.model_type_instance.get_tts_model_voices(
            model=self.model_name, credentials=self.credentials, language=language
        )


class QuotaManagedModelInstance(ModelInstance):
    """A system-hosted model instance that owns quota settlement per invocation."""

    def reserve_quota(
        self,
        *,
        request_id: str | None = None,
        app_type: CreditUsageAppTypeInput = None,
        created_by: CreditUsageCreatedByInput = None,
    ):
        from core.app.llm.quota import reserve_model_quota_for_model
        from core.model_invocation_routing import validate_invocation_admission

        validate_invocation_admission(
            self.credentials,
            tenant_id=self.provider_model_bundle.configuration.tenant_id,
            provider=self.provider,
            model_type=self.model_type_instance.model_type,
            model=self.model_name,
        )

        if app_type is None:
            app_type = self._get_reservation_app_type(self._request_metadata)
        if created_by is None:
            created_by = self._get_reservation_created_by(self._request_metadata)

        return reserve_model_quota_for_model(
            tenant_id=self.provider_model_bundle.configuration.tenant_id,
            provider=self.provider,
            model_type=self.model_type_instance.model_type,
            model=self.model_name,
            request_id=request_id,
            app_type=app_type,
            created_by=created_by,
            provider_configuration=self.provider_model_bundle.configuration,
            invocation_credentials=self.credentials,
        )

    @staticmethod
    def _get_reservation_request_id(request_metadata: Mapping[str, object] | None) -> str | None:
        request_id = request_metadata.get("invocation_id") if request_metadata else None
        if not isinstance(request_id, str) or not request_id:
            return None
        try:
            return str(UUID(request_id))
        except ValueError:
            return None

    @staticmethod
    def _get_reservation_created_by(
        request_metadata: Mapping[str, object] | None,
    ) -> CreditUsageCreatedBy | None:
        created_by = request_metadata.get("created_by") if request_metadata else None
        if created_by is None:
            return None
        return normalize_credit_usage_created_by(created_by)

    @staticmethod
    def _get_reservation_app_type(
        request_metadata: Mapping[str, object] | None,
    ) -> CreditUsageAppType | None:
        app_type = request_metadata.get("app_type") if request_metadata else None
        if app_type is None:
            return None
        return normalize_credit_usage_app_type(app_type)

    def _reserve_quota_for_request(self, request_metadata: Mapping[str, object] | None):
        # Validate the prepared binding, never mutate shared instance credentials
        # between schema/token reads and this invocation's reservation.
        request_id = self._get_reservation_request_id(request_metadata)
        app_type = self._get_reservation_app_type(request_metadata)
        created_by = self._get_reservation_created_by(request_metadata)
        if request_id is None and app_type is None and created_by is None:
            return self.reserve_quota()

        reservation_kwargs: dict[str, str | CreditUsageAppType | CreditUsageCreatedBy] = {}
        if request_id is not None:
            reservation_kwargs["request_id"] = request_id
        if app_type is not None:
            reservation_kwargs["app_type"] = app_type
        if created_by is not None:
            reservation_kwargs["created_by"] = created_by
        return self.reserve_quota(**reservation_kwargs)

    @staticmethod
    def release_quota_safely(reservation) -> None:
        try:
            reservation.release()
        except Exception:
            logger.exception("Failed to release model quota reservation")

    def _invoke_with_quota(self, function: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
        request_metadata = kwargs.get("request_metadata")
        effective_request_metadata = self._resolve_request_metadata(
            request_metadata if isinstance(request_metadata, Mapping) else None
        )
        if effective_request_metadata is not None:
            kwargs["request_metadata"] = effective_request_metadata
        reservation = self._reserve_quota_for_request(effective_request_metadata)
        try:
            response = function(*args, **kwargs)
            reservation.commit()
            return response
        finally:
            self.release_quota_safely(reservation)

    @overload
    def invoke_llm(
        self,
        prompt_messages: Sequence[PromptMessage],
        model_parameters: dict[str, Any] | None = None,
        tools: Sequence[PromptMessageTool] | None = None,
        stop: list[str] | None = None,
        stream: Literal[True] = True,
        callbacks: list[Callback] | None = None,
        request_metadata: Mapping[str, object] | None = None,
    ) -> Generator: ...

    @overload
    def invoke_llm(
        self,
        prompt_messages: list[PromptMessage],
        model_parameters: dict[str, Any] | None = None,
        tools: Sequence[PromptMessageTool] | None = None,
        stop: list[str] | None = None,
        stream: Literal[False] = False,
        callbacks: list[Callback] | None = None,
        request_metadata: Mapping[str, object] | None = None,
    ) -> LLMResult: ...

    @overload
    def invoke_llm(
        self,
        prompt_messages: list[PromptMessage],
        model_parameters: dict[str, Any] | None = None,
        tools: Sequence[PromptMessageTool] | None = None,
        stop: list[str] | None = None,
        stream: bool = True,
        callbacks: list[Callback] | None = None,
        request_metadata: Mapping[str, object] | None = None,
    ) -> LLMResult | Generator: ...

    @override
    def invoke_llm(
        self,
        prompt_messages: Sequence[PromptMessage],
        model_parameters: dict[str, Any] | None = None,
        tools: Sequence[PromptMessageTool] | None = None,
        stop: Sequence[str] | None = None,
        stream: bool = True,
        callbacks: list[Callback] | None = None,
        request_metadata: Mapping[str, object] | None = None,
    ) -> LLMResult | Generator:
        normalized_prompt_messages = list(prompt_messages)
        normalized_stop = list(stop) if stop else None
        if stream:
            return self._invoke_llm_stream(
                prompt_messages=normalized_prompt_messages,
                model_parameters=model_parameters,
                tools=tools,
                stop=normalized_stop,
                callbacks=callbacks,
                request_metadata=request_metadata,
            )

        effective_request_metadata = self._resolve_request_metadata(request_metadata)
        reservation = self._reserve_quota_for_request(effective_request_metadata)
        try:
            response = super().invoke_llm(
                prompt_messages=normalized_prompt_messages,
                model_parameters=model_parameters,
                tools=tools,
                stop=normalized_stop,
                stream=False,
                callbacks=callbacks,
                request_metadata=effective_request_metadata,
            )
            if isinstance(response, Generator):
                raise TypeError("Non-streaming LLM invocation returned a generator.")
            reservation.commit(response.usage)
            return response
        finally:
            self.release_quota_safely(reservation)

    def _invoke_llm_stream(
        self,
        *,
        prompt_messages: list[PromptMessage],
        model_parameters: dict[str, Any] | None,
        tools: Sequence[PromptMessageTool] | None,
        stop: list[str] | None,
        callbacks: list[Callback] | None,
        request_metadata: Mapping[str, object] | None,
    ) -> Generator:
        effective_request_metadata = self._resolve_request_metadata(request_metadata)
        reservation = self._reserve_quota_for_request(effective_request_metadata)
        usage: LLMUsage | None = None
        try:
            response = super().invoke_llm(
                prompt_messages=prompt_messages,
                model_parameters=model_parameters,
                tools=tools,
                stop=stop,
                stream=True,
                callbacks=callbacks,
                request_metadata=effective_request_metadata,
            )
            if not isinstance(response, Generator):
                raise TypeError("Streaming LLM invocation did not return a generator.")

            if reservation.commit_before_delivery:
                for chunk in response:
                    chunk_usage = chunk.delta.usage
                    if chunk_usage is not None:
                        usage = chunk_usage
                    reservation.commit(usage)
                    yield chunk
                return

            buffered_chunks = []
            for chunk in response:
                chunk_usage = chunk.delta.usage
                if chunk_usage is not None:
                    usage = chunk_usage
                buffered_chunks.append(chunk)

            reservation.commit(usage)
            yield from buffered_chunks
        finally:
            self.release_quota_safely(reservation)

    @override
    def invoke_text_embedding(
        self,
        texts: list[str],
        input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> EmbeddingResult:
        return self._invoke_with_quota(
            super().invoke_text_embedding,
            texts=texts,
            input_type=input_type,
            request_metadata=request_metadata,
        )

    @override
    def invoke_multimodal_embedding(
        self,
        multimodel_documents: list[dict],
        input_type: EmbeddingInputType = EmbeddingInputType.DOCUMENT,
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> EmbeddingResult:
        return self._invoke_with_quota(
            super().invoke_multimodal_embedding,
            multimodel_documents=multimodel_documents,
            input_type=input_type,
            request_metadata=request_metadata,
        )

    @override
    def invoke_rerank(
        self,
        query: str,
        docs: list[str],
        score_threshold: float | None = None,
        top_n: int | None = None,
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> RerankResult:
        return self._invoke_with_quota(
            super().invoke_rerank,
            query=query,
            docs=docs,
            score_threshold=score_threshold,
            top_n=top_n,
            request_metadata=request_metadata,
        )

    @override
    def invoke_multimodal_rerank(
        self,
        query: MultimodalRerankInput,
        docs: list[MultimodalRerankInput],
        score_threshold: float | None = None,
        top_n: int | None = None,
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> RerankResult:
        return self._invoke_with_quota(
            super().invoke_multimodal_rerank,
            query=query,
            docs=docs,
            score_threshold=score_threshold,
            top_n=top_n,
            request_metadata=request_metadata,
        )

    @override
    def invoke_moderation(self, text: str, *, request_metadata: Mapping[str, object] | None = None) -> bool:
        return self._invoke_with_quota(super().invoke_moderation, text=text, request_metadata=request_metadata)

    @override
    def invoke_speech2text(self, file: IO[bytes], *, request_metadata: Mapping[str, object] | None = None) -> str:
        return self._invoke_with_quota(super().invoke_speech2text, file=file, request_metadata=request_metadata)

    @override
    def invoke_tts(
        self,
        content_text: str,
        voice: str = "",
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> Iterable[bytes]:
        return self._invoke_tts_stream(content_text=content_text, voice=voice, request_metadata=request_metadata)

    def _invoke_tts_stream(
        self,
        *,
        content_text: str,
        voice: str,
        request_metadata: Mapping[str, object] | None = None,
    ) -> Generator[bytes]:
        effective_request_metadata = self._resolve_request_metadata(request_metadata)
        reservation = self._reserve_quota_for_request(effective_request_metadata)
        try:
            response = super().invoke_tts(
                content_text=content_text,
                voice=voice,
                request_metadata=effective_request_metadata,
            )
            for chunk in response:
                reservation.commit()
                yield chunk
        finally:
            self.release_quota_safely(reservation)


def create_model_instance(
    provider_model_bundle: ProviderModelBundle,
    model: str,
    credentials: dict[str, Any] | None = None,
    request_metadata: Mapping[str, object] | None = None,
    *,
    message_billing: "ClassicMessageBilling | None" = None,
) -> ModelInstance:
    """One factory for production callers, including classic Chat/Agent apps."""
    if message_billing is not None:
        message_billing.check_identity(
            provider_model_bundle.configuration.tenant_id, provider_model_bundle.configuration.provider.provider, model
        )
        # A classic Agent can invoke this model many times, but its legacy
        # accounting unit remains the enclosing message's one durable receipt.
        instance = ModelInstance(provider_model_bundle, model, credentials, request_metadata)
        instance._message_billing = message_billing
        return instance
    instance_type = (
        QuotaManagedModelInstance
        if provider_model_bundle.configuration.using_provider_type == ProviderType.SYSTEM
        else ModelInstance
    )
    return instance_type(provider_model_bundle, model, credentials, request_metadata)


class ModelManager:
    """Resolves :class:`ModelInstance` objects for a tenant and provider.

    When ``enable_credentials_cache`` is ``True``, resolved credentials for each
    ``(tenant_id, provider, model_type, model)`` are stored in
    ``_credentials_cache`` and reused. That can return **stale** credentials after
    API keys or provider settings change, so a manager constructed with
    ``enable_credentials_cache=True`` should not be kept for the lifetime of a
    process or shared across unrelated work. Prefer a new manager per request,
    workflow run, or similar bounded scope.

    The default is ``enable_credentials_cache=False``; in that mode the internal
    credential cache is not populated, and each ``get_model_instance`` call
    loads credentials from the current provider configuration.
    """

    def __init__(
        self,
        provider_manager: ProviderManager,
        *,
        enable_credentials_cache: bool = False,
        request_metadata: Mapping[str, object] | None = None,
    ) -> None:
        self._provider_manager = provider_manager
        self._credentials_cache: dict[tuple[str, str, str, str], Any] = {}
        self._enable_credentials_cache = enable_credentials_cache
        self._request_metadata = dict(request_metadata) if request_metadata else None

    @classmethod
    def for_tenant(
        cls,
        tenant_id: str,
        user_id: str | None = None,
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> "ModelManager":
        if request_metadata is None:
            request_metadata = get_credit_usage_metadata()
        return cls(
            provider_manager=create_plugin_provider_manager(tenant_id=tenant_id, user_id=user_id),
            request_metadata=request_metadata,
        )

    def _resolve_request_metadata(self, request_metadata: Mapping[str, object] | None) -> Mapping[str, object] | None:
        if self._request_metadata is None:
            return request_metadata
        if request_metadata is None:
            return self._request_metadata
        return {**self._request_metadata, **request_metadata}

    @staticmethod
    def _validate_system_model_access(
        provider_model_bundle: ProviderModelBundle,
        *,
        model_type: ModelType,
        model: str,
    ) -> None:
        configuration = provider_model_bundle.configuration
        if configuration.using_provider_type != ProviderType.SYSTEM:
            return

        from core.model_invocation_routing import has_compatibility_route, migration_routing_state

        if ModelBillingProfileService.resolve(configuration.tenant_id).uses_tokener and not has_compatibility_route(
            migration_routing_state(configuration.tenant_id)
        ):
            raise ModelCurrentlyNotSupportError(
                f"Hosted SYSTEM model {model_type.value}/{model} is disabled for this workspace."
            )

        # Hosted allowlists retain the existing comma-separated format. Model names
        # are matched exactly; model-type-specific entries will be introduced later.
        quota_configuration = next(
            (
                quota
                for quota in configuration.system_configuration.quota_configurations
                if quota.quota_type == configuration.system_configuration.current_quota_type
            ),
            None,
        )
        if quota_configuration is None or not quota_configuration.restrict_models:
            return
        if any(restricted_model.model == model for restricted_model in quota_configuration.restrict_models):
            return

        raise ModelCurrentlyNotSupportError(f"System model {model_type.value}/{model} is not allowed.")

    @staticmethod
    def _model_instance_class(provider_model_bundle: ProviderModelBundle, model_type: ModelType) -> type[ModelInstance]:
        if provider_model_bundle.configuration.using_provider_type == ProviderType.SYSTEM:
            return QuotaManagedModelInstance
        return ModelInstance

    def get_model_instance(
        self,
        tenant_id: str,
        provider: str,
        model_type: ModelType,
        model: str,
        *,
        request_metadata: Mapping[str, object] | None = None,
    ) -> ModelInstance:
        """
        Get model instance
        :param tenant_id: tenant id
        :param provider: provider name
        :param model_type: model type
        :param model: model name
        :return:
        """
        if not provider:
            return self.get_default_model_instance(tenant_id, model_type)

        provider_model_bundle = self._provider_manager.get_provider_model_bundle(
            tenant_id=tenant_id, provider=provider, model_type=model_type
        )
        self._validate_system_model_access(provider_model_bundle, model_type=model_type, model=model)
        model_instance_class = self._model_instance_class(provider_model_bundle, model_type)
        effective_request_metadata = self._resolve_request_metadata(request_metadata)

        cred_cache_key = (tenant_id, provider, model_type.value, model)

        if cred_cache_key in self._credentials_cache:
            from core.model_invocation_routing import ModelInvocationReprepare, validate_invocation_admission

            try:
                validate_invocation_admission(
                    self._credentials_cache[cred_cache_key],
                    tenant_id=tenant_id,
                    provider=provider,
                    model_type=model_type,
                    model=model,
                )
            except ModelInvocationReprepare:
                self._credentials_cache.pop(cred_cache_key, None)
                self._provider_manager.clear_configurations_cache(tenant_id)
                return self.get_model_instance(
                    tenant_id, provider, model_type, model, request_metadata=request_metadata
                )
            return model_instance_class(
                provider_model_bundle,
                model,
                deepcopy(self._credentials_cache[cred_cache_key]),
                effective_request_metadata,
            )

        ret = model_instance_class(provider_model_bundle, model, request_metadata=effective_request_metadata)
        if self._enable_credentials_cache:
            self._credentials_cache[cred_cache_key] = deepcopy(ret.credentials)
        return ret

    def get_default_provider_model_name(self, tenant_id: str, model_type: ModelType) -> tuple[str | None, str | None]:
        """
        Return first provider and the first model in the provider
        :param tenant_id: tenant id
        :param model_type: model type
        :return: provider name, model name
        """
        return self._provider_manager.get_first_provider_first_model(tenant_id, model_type)

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

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

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

    def check_model_support_vision(self, tenant_id: str, provider: str, model: str, model_type: ModelType) -> bool:
        """
        Check if model supports vision
        :param tenant_id: tenant id
        :param provider: provider name
        :param model: model name
        :return: True if model supports vision, False otherwise
        """
        model_instance = self.get_model_instance(tenant_id, provider, model_type, model)
        model_type_instance = model_instance.model_type_instance
        match model_type:
            case ModelType.LLM:
                model_type_instance = cast(LargeLanguageModel, model_type_instance)
            case ModelType.TEXT_EMBEDDING:
                model_type_instance = cast(TextEmbeddingModel, model_type_instance)
            case ModelType.RERANK:
                model_type_instance = cast(RerankModel, model_type_instance)
            case _:
                raise ValueError(f"Model type {model_type} is not supported")
        model_schema = model_type_instance.get_model_schema(model, model_instance.credentials)
        if not model_schema:
            return False
        return bool(model_schema.features and ModelFeature.VISION in model_schema.features)


class LBModelManager:
    def __init__(
        self,
        tenant_id: str,
        provider: str,
        model_type: ModelType,
        model: str,
        load_balancing_configs: list[ModelLoadBalancingConfiguration],
        managed_credentials: dict[str, Any] | None = None,
    ):
        """
        Load balancing model manager
        :param tenant_id: tenant_id
        :param provider: provider
        :param model_type: model_type
        :param model: model name
        :param load_balancing_configs: all load balancing configurations
        :param managed_credentials: credentials if load balancing configuration name is __inherit__
        """
        self._tenant_id = tenant_id
        self._provider = provider
        self._model_type = model_type
        self._model = model
        self._load_balancing_configs = load_balancing_configs

        for load_balancing_config in self._load_balancing_configs[:]:  # Iterate over a shallow copy of the list
            if load_balancing_config.name == "__inherit__":
                if not managed_credentials:
                    # remove __inherit__ if managed credentials is not provided
                    self._load_balancing_configs.remove(load_balancing_config)
                else:
                    load_balancing_config.credentials = managed_credentials

    def fetch_next(self) -> ModelLoadBalancingConfiguration | None:
        """
        Get next model load balancing config
        Strategy: Round Robin
        :return:
        """
        cache_key = f"model_lb_index:{self._tenant_id}:{self._provider}:{self._model_type.value}:{self._model}"

        cooldown_load_balancing_configs = []
        max_index = len(self._load_balancing_configs)

        while True:
            current_index = redis_client.incr(cache_key)
            current_index = cast(int, current_index)
            if current_index >= 10000000:
                current_index = 1
                redis_client.set(cache_key, current_index)

            redis_client.expire(cache_key, 3600)
            if current_index > max_index:
                current_index = current_index % max_index

            real_index = current_index - 1
            if real_index > max_index:
                real_index = 0

            config: ModelLoadBalancingConfiguration = self._load_balancing_configs[real_index]

            if self.in_cooldown(config):
                cooldown_load_balancing_configs.append(config)
                if len(cooldown_load_balancing_configs) >= len(self._load_balancing_configs):
                    # all configs are in cooldown
                    return None

                continue

            # Check policy compliance for the selected configuration
            try:
                from core.helper.credential_utils import runtime_check_credential_policy_compliance

                if config.credential_id:
                    runtime_check_credential_policy_compliance(
                        credential_id=config.credential_id,
                        provider=self._provider,
                        credential_type=PluginCredentialType.MODEL,
                    )
            except Exception:
                logger.warning("Load balancing config %s failed policy compliance check", config.id, exc_info=True)
                cooldown_load_balancing_configs.append(config)
                if len(cooldown_load_balancing_configs) >= len(self._load_balancing_configs):
                    # all configs are in cooldown or failed policy compliance
                    return None
                continue

            if dify_config.DEBUG:
                logger.info(
                    """Model LB
id: %s
name:%s
tenant_id: %s
provider: %s
model_type: %s
model: %s""",
                    config.id,
                    config.name,
                    self._tenant_id,
                    self._provider,
                    self._model_type.value,
                    self._model,
                )

            return config

    def cooldown(self, config: ModelLoadBalancingConfiguration, expire: int = 60):
        """
        Cooldown model load balancing config
        :param config: model load balancing config
        :param expire: cooldown time
        :return:
        """
        cooldown_cache_key = (
            f"model_lb_index:cooldown:{self._tenant_id}:{self._provider}:"
            f"{self._model_type.value}:{self._model}:{config.id}"
        )

        redis_client.setex(cooldown_cache_key, expire, "true")

    def in_cooldown(self, config: ModelLoadBalancingConfiguration) -> bool:
        """
        Check if model load balancing config is in cooldown
        :param config: model load balancing config
        :return:
        """
        cooldown_cache_key = (
            f"model_lb_index:cooldown:{self._tenant_id}:{self._provider}:"
            f"{self._model_type.value}:{self._model}:{config.id}"
        )

        res: bool = redis_client.exists(cooldown_cache_key)
        return res

    @staticmethod
    def get_config_in_cooldown_and_ttl(
        tenant_id: str, provider: str, model_type: ModelType, model: str, config_id: str
    ) -> tuple[bool, int]:
        """
        Get model load balancing config is in cooldown and ttl
        :param tenant_id: workspace id
        :param provider: provider name
        :param model_type: model type
        :param model: model name
        :param config_id: model load balancing config id
        :return:
        """
        cooldown_cache_key = f"model_lb_index:cooldown:{tenant_id}:{provider}:{model_type.value}:{model}:{config_id}"

        ttl = redis_client.ttl(cooldown_cache_key)
        if ttl == -2:
            return False, 0

        ttl = cast(int, ttl)
        return True, ttl