Model Runtime (#1858)

Co-authored-by: StyleZhang <jasonapring2015@outlook.com>
Co-authored-by: Garfield Dai <dai.hai@foxmail.com>
Co-authored-by: chenhe <guchenhe@gmail.com>
Co-authored-by: jyong <jyong@dify.ai>
Co-authored-by: Joel <iamjoel007@gmail.com>
Co-authored-by: Yeuoly <admin@srmxy.cn>

api/core/model_runtime/entities/common_entities.py

@@ -0,0 +1,16 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class I18nObject(BaseModel):
+    """
+    Model class for i18n object.
+    """
+    zh_Hans: Optional[str] = None
+    en_US: str
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        if not self.zh_Hans:
+            self.zh_Hans = self.en_US

api/core/model_runtime/entities/defaults.py

@@ -0,0 +1,87 @@
+from typing import Dict
+
+from core.model_runtime.entities.model_entities import DefaultParameterName
+
+
+PARAMETER_RULE_TEMPLATE: Dict[DefaultParameterName, dict] = {
+    DefaultParameterName.TEMPERATURE: {
+        'label': {
+            'en_US': 'Temperature',
+            'zh_Hans': '温度',
+        },
+        'type': 'float',
+        'help': {
+            'en_US': 'Controls randomness. Lower temperature results in less random completions. As the temperature approaches zero, the model will become deterministic and repetitive. Higher temperature results in more random completions.',
+            'zh_Hans': '温度控制随机性。较低的温度会导致较少的随机完成。随着温度接近零，模型将变得确定性和重复性。较高的温度会导致更多的随机完成。',
+        },
+        'required': False,
+        'default': 0.0,
+        'min': 0.0,
+        'max': 1.0,
+        'precision': 1,
+    },
+    DefaultParameterName.TOP_P: {
+        'label': {
+            'en_US': 'Top P',
+            'zh_Hans': 'Top P',
+        },
+        'type': 'float',
+        'help': {
+            'en_US': 'Controls diversity via nucleus sampling: 0.5 means half of all likelihood-weighted options are considered.',
+            'zh_Hans': '通过核心采样控制多样性：0.5表示考虑了一半的所有可能性加权选项。',
+        },
+        'required': False,
+        'default': 1.0,
+        'min': 0.0,
+        'max': 1.0,
+        'precision': 1,
+    },
+    DefaultParameterName.PRESENCE_PENALTY: {
+        'label': {
+            'en_US': 'Presence Penalty',
+            'zh_Hans': '存在惩罚',
+        },
+        'type': 'float',
+        'help': {
+            'en_US': 'Applies a penalty to the log-probability of tokens already in the text.',
+            'zh_Hans': '对文本中已有的标记的对数概率施加惩罚。',
+        },
+        'required': False,
+        'default': 0.0,
+        'min': 0.0,
+        'max': 1.0,
+        'precision': 1,
+    },
+    DefaultParameterName.FREQUENCY_PENALTY: {
+        'label': {
+            'en_US': 'Frequency Penalty',
+            'zh_Hans': '频率惩罚',
+        },
+        'type': 'float',
+        'help': {
+            'en_US': 'Applies a penalty to the log-probability of tokens that appear in the text.',
+            'zh_Hans': '对文本中出现的标记的对数概率施加惩罚。',
+        },
+        'required': False,
+        'default': 0.0,
+        'min': 0.0,
+        'max': 1.0,
+        'precision': 1,
+    },
+    DefaultParameterName.MAX_TOKENS: {
+        'label': {
+            'en_US': 'Max Tokens',
+            'zh_Hans': '最大标记',
+        },
+        'type': 'int',
+        'help': {
+            'en_US': 'The maximum number of tokens to generate. Requests can use up to 2048 tokens shared between prompt and completion.',
+            'zh_Hans': '要生成的标记的最大数量。请求可以使用最多2048个标记，这些标记在提示和完成之间共享。',
+        },
+        'required': False,
+        'default': 64,
+        'min': 1,
+        'max': 2048,
+        'precision': 0,
+    }
+}

api/core/model_runtime/entities/llm_entities.py

@@ -0,0 +1,102 @@
+from decimal import Decimal
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel
+
+from core.model_runtime.entities.message_entities import AssistantPromptMessage, PromptMessage
+from core.model_runtime.entities.model_entities import ModelUsage, PriceInfo
+
+
+class LLMMode(Enum):
+    """
+    Enum class for large language model mode.
+    """
+    COMPLETION = "completion"
+    CHAT = "chat"
+
+    @classmethod
+    def value_of(cls, value: str) -> 'LLMMode':
+        """
+        Get value of given mode.
+
+        :param value: mode value
+        :return: mode
+        """
+        for mode in cls:
+            if mode.value == value:
+                return mode
+        raise ValueError(f'invalid mode value {value}')
+
+
+class LLMUsage(ModelUsage):
+    """
+    Model class for llm usage.
+    """
+    prompt_tokens: int
+    prompt_unit_price: Decimal
+    prompt_price_unit: Decimal
+    prompt_price: Decimal
+    completion_tokens: int
+    completion_unit_price: Decimal
+    completion_price_unit: Decimal
+    completion_price: Decimal
+    total_tokens: int
+    total_price: Decimal
+    currency: str
+    latency: float
+
+    @classmethod
+    def empty_usage(cls):
+        return cls(
+            prompt_tokens=0,
+            prompt_unit_price=Decimal('0.0'),
+            prompt_price_unit=Decimal('0.0'),
+            prompt_price=Decimal('0.0'),
+            completion_tokens=0,
+            completion_unit_price=Decimal('0.0'),
+            completion_price_unit=Decimal('0.0'),
+            completion_price=Decimal('0.0'),
+            total_tokens=0,
+            total_price=Decimal('0.0'),
+            currency='USD',
+            latency=0.0
+        )
+
+
+class LLMResult(BaseModel):
+    """
+    Model class for llm result.
+    """
+    model: str
+    prompt_messages: list[PromptMessage]
+    message: AssistantPromptMessage
+    usage: LLMUsage
+    system_fingerprint: Optional[str] = None
+
+
+class LLMResultChunkDelta(BaseModel):
+    """
+    Model class for llm result chunk delta.
+    """
+    index: int
+    message: AssistantPromptMessage
+    usage: Optional[LLMUsage] = None
+    finish_reason: Optional[str] = None
+
+
+class LLMResultChunk(BaseModel):
+    """
+    Model class for llm result chunk.
+    """
+    model: str
+    prompt_messages: list[PromptMessage]
+    system_fingerprint: Optional[str] = None
+    delta: LLMResultChunkDelta
+
+
+class NumTokensResult(PriceInfo):
+    """
+    Model class for number of tokens result.
+    """
+    tokens: int

api/core/model_runtime/entities/message_entities.py

@@ -0,0 +1,134 @@
+from abc import ABC
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel
+
+
+class PromptMessageRole(Enum):
+    """
+    Enum class for prompt message.
+    """
+    SYSTEM = "system"
+    USER = "user"
+    ASSISTANT = "assistant"
+    TOOL = "tool"
+
+    @classmethod
+    def value_of(cls, value: str) -> 'PromptMessageRole':
+        """
+        Get value of given mode.
+
+        :param value: mode value
+        :return: mode
+        """
+        for mode in cls:
+            if mode.value == value:
+                return mode
+        raise ValueError(f'invalid prompt message type value {value}')
+
+
+class PromptMessageTool(BaseModel):
+    """
+    Model class for prompt message tool.
+    """
+    name: str
+    description: str
+    parameters: dict
+
+
+class PromptMessageFunction(BaseModel):
+    """
+    Model class for prompt message function.
+    """
+    type: str = 'function'
+    function: PromptMessageTool
+
+
+class PromptMessageContentType(Enum):
+    """
+    Enum class for prompt message content type.
+    """
+    TEXT = 'text'
+    IMAGE = 'image'
+
+
+class PromptMessageContent(BaseModel):
+    """
+    Model class for prompt message content.
+    """
+    type: PromptMessageContentType
+    data: str
+
+
+class TextPromptMessageContent(PromptMessageContent):
+    """
+    Model class for text prompt message content.
+    """
+    type: PromptMessageContentType = PromptMessageContentType.TEXT
+
+
+class ImagePromptMessageContent(PromptMessageContent):
+    """
+    Model class for image prompt message content.
+    """
+    class DETAIL(Enum):
+        LOW = 'low'
+        HIGH = 'high'
+
+    type: PromptMessageContentType = PromptMessageContentType.IMAGE
+    detail: DETAIL = DETAIL.LOW
+
+
+class PromptMessage(ABC, BaseModel):
+    """
+    Model class for prompt message.
+    """
+    role: PromptMessageRole
+    content: Optional[str | list[PromptMessageContent]] = None
+    name: Optional[str] = None
+
+
+class UserPromptMessage(PromptMessage):
+    """
+    Model class for user prompt message.
+    """
+    role: PromptMessageRole = PromptMessageRole.USER
+
+
+class AssistantPromptMessage(PromptMessage):
+    """
+    Model class for assistant prompt message.
+    """
+    class ToolCall(BaseModel):
+        """
+        Model class for assistant prompt message tool call.
+        """
+        class ToolCallFunction(BaseModel):
+            """
+            Model class for assistant prompt message tool call function.
+            """
+            name: str
+            arguments: str
+
+        id: str
+        type: str
+        function: ToolCallFunction
+
+    role: PromptMessageRole = PromptMessageRole.ASSISTANT
+    tool_calls: list[ToolCall] = []
+
+
+class SystemPromptMessage(PromptMessage):
+    """
+    Model class for system prompt message.
+    """
+    role: PromptMessageRole = PromptMessageRole.SYSTEM
+
+
+class ToolPromptMessage(PromptMessage):
+    """
+    Model class for tool prompt message.
+    """
+    role: PromptMessageRole = PromptMessageRole.TOOL
+    tool_call_id: str

api/core/model_runtime/entities/model_entities.py

@@ -0,0 +1,196 @@
+from decimal import Decimal
+from enum import Enum
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+from core.model_runtime.entities.common_entities import I18nObject
+
+
+class ModelType(Enum):
+    """
+    Enum class for model type.
+    """
+    LLM = "llm"
+    TEXT_EMBEDDING = "text-embedding"
+    RERANK = "rerank"
+    SPEECH2TEXT = "speech2text"
+    MODERATION = "moderation"
+    # TTS = "tts"
+    # TEXT2IMG = "text2img"
+
+    @classmethod
+    def value_of(cls, origin_model_type: str) -> "ModelType":
+        """
+        Get model type from origin model type.
+
+        :return: model type
+        """
+        if origin_model_type == 'text-generation' or origin_model_type == cls.LLM.value:
+            return cls.LLM
+        elif origin_model_type == 'embeddings' or origin_model_type == cls.TEXT_EMBEDDING.value:
+            return cls.TEXT_EMBEDDING
+        elif origin_model_type == 'reranking' or origin_model_type == cls.RERANK.value:
+            return cls.RERANK
+        elif origin_model_type == cls.SPEECH2TEXT.value:
+            return cls.SPEECH2TEXT
+        elif origin_model_type == cls.MODERATION.value:
+            return cls.MODERATION
+        else:
+            raise ValueError(f'invalid origin model type {origin_model_type}')
+
+    def to_origin_model_type(self) -> str:
+        """
+        Get origin model type from model type.
+
+        :return: origin model type
+        """
+        if self == self.LLM:
+            return 'text-generation'
+        elif self == self.TEXT_EMBEDDING:
+            return 'embeddings'
+        elif self == self.RERANK:
+            return 'reranking'
+        elif self == self.SPEECH2TEXT:
+            return 'speech2text'
+        elif self == self.MODERATION:
+            return 'moderation'
+        else:
+            raise ValueError(f'invalid model type {self}')
+
+
+class FetchFrom(Enum):
+    """
+    Enum class for fetch from.
+    """
+    PREDEFINED_MODEL = "predefined-model"
+    CUSTOMIZABLE_MODEL = "customizable-model"
+
+
+class ModelFeature(Enum):
+    """
+    Enum class for llm feature.
+    """
+    TOOL_CALL = "tool-call"
+    MULTI_TOOL_CALL = "multi-tool-call"
+    AGENT_THOUGHT = "agent-thought"
+    VISION = "vision"
+
+
+class DefaultParameterName(Enum):
+    """
+    Enum class for parameter template variable.
+    """
+    TEMPERATURE = "temperature"
+    TOP_P = "top_p"
+    PRESENCE_PENALTY = "presence_penalty"
+    FREQUENCY_PENALTY = "frequency_penalty"
+    MAX_TOKENS = "max_tokens"
+
+    @classmethod
+    def value_of(cls, value: Any) -> 'DefaultParameterName':
+        """
+        Get parameter name from value.
+
+        :param value: parameter value
+        :return: parameter name
+        """
+        for name in cls:
+            if name.value == value:
+                return name
+        raise ValueError(f'invalid parameter name {value}')
+
+
+class ParameterType(Enum):
+    """
+    Enum class for parameter type.
+    """
+    FLOAT = "float"
+    INT = "int"
+    STRING = "string"
+    BOOLEAN = "boolean"
+
+
+class ModelPropertyKey(Enum):
+    """
+    Enum class for model property key.
+    """
+    MODE = "mode"
+    CONTEXT_SIZE = "context_size"
+    MAX_CHUNKS = "max_chunks"
+    FILE_UPLOAD_LIMIT = "file_upload_limit"
+    SUPPORTED_FILE_EXTENSIONS = "supported_file_extensions"
+    MAX_CHARACTERS_PER_CHUNK = "max_characters_per_chunk"
+
+
+class ProviderModel(BaseModel):
+    """
+    Model class for provider model.
+    """
+    model: str
+    label: I18nObject
+    model_type: ModelType
+    features: Optional[list[ModelFeature]] = None
+    fetch_from: FetchFrom
+    model_properties: dict[ModelPropertyKey, Any]
+    deprecated: bool = False
+
+    class Config:
+        protected_namespaces = ()
+
+
+class ParameterRule(BaseModel):
+    """
+    Model class for parameter rule.
+    """
+    name: str
+    use_template: Optional[str] = None
+    label: I18nObject
+    type: ParameterType
+    help: Optional[I18nObject] = None
+    required: bool = False
+    default: Optional[Any] = None
+    min: Optional[float | int] = None
+    max: Optional[float | int] = None
+    precision: Optional[int] = None
+    options: list[str] = []
+
+
+class PriceConfig(BaseModel):
+    """
+    Model class for pricing info.
+    """
+    input: Decimal
+    output: Optional[Decimal] = None
+    unit: Decimal
+    currency: str
+
+
+class AIModelEntity(ProviderModel):
+    """
+    Model class for AI model.
+    """
+    parameter_rules: list[ParameterRule] = []
+    pricing: Optional[PriceConfig] = None
+
+
+class ModelUsage(BaseModel):
+    pass
+
+
+class PriceType(Enum):
+    """
+    Enum class for price type.
+    """
+    INPUT = "input"
+    OUTPUT = "output"
+
+
+class PriceInfo(BaseModel):
+    """
+    Model class for price info.
+    """
+    unit_price: Decimal
+    unit: Decimal
+    total_amount: Decimal
+    currency: str

api/core/model_runtime/entities/provider_entities.py

@@ -0,0 +1,149 @@
+from enum import Enum
+from typing import Optional
+
+from pydantic import BaseModel
+
+from core.model_runtime.entities.common_entities import I18nObject
+from core.model_runtime.entities.model_entities import ModelType, ProviderModel, AIModelEntity
+
+
+class ConfigurateMethod(Enum):
+    """
+    Enum class for configurate method of provider model.
+    """
+    PREDEFINED_MODEL = "predefined-model"
+    CUSTOMIZABLE_MODEL = "customizable-model"
+
+
+class FormType(Enum):
+    """
+    Enum class for form type.
+    """
+    TEXT_INPUT = "text-input"
+    SECRET_INPUT = "secret-input"
+    SELECT = "select"
+    RADIO = "radio"
+    SWITCH = "switch"
+
+
+class FormShowOnObject(BaseModel):
+    """
+    Model class for form show on.
+    """
+    variable: str
+    value: str
+
+
+class FormOption(BaseModel):
+    """
+    Model class for form option.
+    """
+    label: I18nObject
+    value: str
+    show_on: list[FormShowOnObject] = []
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        if not self.label:
+            self.label = I18nObject(
+                en_US=self.value
+            )
+
+
+class CredentialFormSchema(BaseModel):
+    """
+    Model class for credential form schema.
+    """
+    variable: str
+    label: I18nObject
+    type: FormType
+    required: bool = True
+    default: Optional[str] = None
+    options: Optional[list[FormOption]] = None
+    placeholder: Optional[I18nObject] = None
+    max_length: int = 0
+    show_on: list[FormShowOnObject] = []
+
+
+class ProviderCredentialSchema(BaseModel):
+    """
+    Model class for provider credential schema.
+    """
+    credential_form_schemas: list[CredentialFormSchema]
+
+
+class FieldModelSchema(BaseModel):
+    label: I18nObject
+    placeholder: Optional[I18nObject] = None
+
+
+class ModelCredentialSchema(BaseModel):
+    """
+    Model class for model credential schema.
+    """
+    model: FieldModelSchema
+    credential_form_schemas: list[CredentialFormSchema]
+
+
+class SimpleProviderEntity(BaseModel):
+    """
+    Simple model class for provider.
+    """
+    provider: str
+    label: I18nObject
+    icon_small: Optional[I18nObject] = None
+    icon_large: Optional[I18nObject] = None
+    supported_model_types: list[ModelType]
+    models: list[AIModelEntity] = []
+
+
+class ProviderHelpEntity(BaseModel):
+    """
+    Model class for provider help.
+    """
+    title: I18nObject
+    url: I18nObject
+
+
+class ProviderEntity(BaseModel):
+    """
+    Model class for provider.
+    """
+    provider: str
+    label: I18nObject
+    description: Optional[I18nObject] = None
+    icon_small: Optional[I18nObject] = None
+    icon_large: Optional[I18nObject] = None
+    background: Optional[str] = None
+    help: Optional[ProviderHelpEntity] = None
+    supported_model_types: list[ModelType]
+    configurate_methods: list[ConfigurateMethod]
+    models: list[ProviderModel] = []
+    provider_credential_schema: Optional[ProviderCredentialSchema] = None
+    model_credential_schema: Optional[ModelCredentialSchema] = None
+
+    class Config:
+        protected_namespaces = ()
+
+    def to_simple_provider(self) -> SimpleProviderEntity:
+        """
+        Convert to simple provider.
+
+        :return: simple provider
+        """
+        return SimpleProviderEntity(
+            provider=self.provider,
+            label=self.label,
+            icon_small=self.icon_small,
+            icon_large=self.icon_large,
+            supported_model_types=self.supported_model_types,
+            models=self.models
+        )
+
+
+class ProviderConfig(BaseModel):
+    """
+    Model class for provider config.
+    """
+    provider: str
+    credentials: dict

api/core/model_runtime/entities/rerank_entities.py

@@ -0,0 +1,18 @@
+from pydantic import BaseModel
+
+
+class RerankDocument(BaseModel):
+    """
+    Model class for rerank document.
+    """
+    index: int
+    text: str
+    score: float
+
+
+class RerankResult(BaseModel):
+    """
+    Model class for rerank result.
+    """
+    model: str
+    docs: list[RerankDocument]

api/core/model_runtime/entities/text_embedding_entities.py

@@ -0,0 +1,28 @@
+from decimal import Decimal
+
+from pydantic import BaseModel
+
+from core.model_runtime.entities.model_entities import ModelUsage
+
+
+class EmbeddingUsage(ModelUsage):
+    """
+    Model class for embedding usage.
+    """
+    tokens: int
+    total_tokens: int
+    unit_price: Decimal
+    price_unit: Decimal
+    total_price: Decimal
+    currency: str
+    latency: float
+
+
+class TextEmbeddingResult(BaseModel):
+    """
+    Model class for text embedding result.
+    """
+    model: str
+    embeddings: list[list[float]]
+    usage: EmbeddingUsage
+