Underground-Digital/Workflow-Engine

1from decimal import Decimal2from enum import Enum3from typing import Any, Optional4 5from pydantic import BaseModel, ConfigDict6 7from core.model_runtime.entities.common_entities import I18nObject8 9 10class ModelType(Enum):11    """12    Enum class for model type.13    """14 15    LLM = "llm"16    TEXT_EMBEDDING = "text-embedding"17    RERANK = "rerank"18    SPEECH2TEXT = "speech2text"19    MODERATION = "moderation"20    TTS = "tts"21    TEXT2IMG = "text2img"22 23    @classmethod24    def value_of(cls, origin_model_type: str) -> "ModelType":25        """26        Get model type from origin model type.27 28        :return: model type29        """30        if origin_model_type in {"text-generation", cls.LLM.value}:31            return cls.LLM32        elif origin_model_type in {"embeddings", cls.TEXT_EMBEDDING.value}:33            return cls.TEXT_EMBEDDING34        elif origin_model_type in {"reranking", cls.RERANK.value}:35            return cls.RERANK36        elif origin_model_type in {"speech2text", cls.SPEECH2TEXT.value}:37            return cls.SPEECH2TEXT38        elif origin_model_type in {"tts", cls.TTS.value}:39            return cls.TTS40        elif origin_model_type in {"text2img", cls.TEXT2IMG.value}:41            return cls.TEXT2IMG42        elif origin_model_type == cls.MODERATION.value:43            return cls.MODERATION44        else:45            raise ValueError(f"invalid origin model type {origin_model_type}")46 47    def to_origin_model_type(self) -> str:48        """49        Get origin model type from model type.50 51        :return: origin model type52        """53        if self == self.LLM:54            return "text-generation"55        elif self == self.TEXT_EMBEDDING:56            return "embeddings"57        elif self == self.RERANK:58            return "reranking"59        elif self == self.SPEECH2TEXT:60            return "speech2text"61        elif self == self.TTS:62            return "tts"63        elif self == self.MODERATION:64            return "moderation"65        elif self == self.TEXT2IMG:66            return "text2img"67        else:68            raise ValueError(f"invalid model type {self}")69 70 71class FetchFrom(Enum):72    """73    Enum class for fetch from.74    """75 76    PREDEFINED_MODEL = "predefined-model"77    CUSTOMIZABLE_MODEL = "customizable-model"78 79 80class ModelFeature(Enum):81    """82    Enum class for llm feature.83    """84 85    TOOL_CALL = "tool-call"86    MULTI_TOOL_CALL = "multi-tool-call"87    AGENT_THOUGHT = "agent-thought"88    VISION = "vision"89    STREAM_TOOL_CALL = "stream-tool-call"90 91 92class DefaultParameterName(str, Enum):93    """94    Enum class for parameter template variable.95    """96 97    TEMPERATURE = "temperature"98    TOP_P = "top_p"99    TOP_K = "top_k"100    PRESENCE_PENALTY = "presence_penalty"101    FREQUENCY_PENALTY = "frequency_penalty"102    MAX_TOKENS = "max_tokens"103    RESPONSE_FORMAT = "response_format"104    JSON_SCHEMA = "json_schema"105 106    @classmethod107    def value_of(cls, value: Any) -> "DefaultParameterName":108        """109        Get parameter name from value.110 111        :param value: parameter value112        :return: parameter name113        """114        for name in cls:115            if name.value == value:116                return name117        raise ValueError(f"invalid parameter name {value}")118 119 120class ParameterType(Enum):121    """122    Enum class for parameter type.123    """124 125    FLOAT = "float"126    INT = "int"127    STRING = "string"128    BOOLEAN = "boolean"129    TEXT = "text"130 131 132class ModelPropertyKey(Enum):133    """134    Enum class for model property key.135    """136 137    MODE = "mode"138    CONTEXT_SIZE = "context_size"139    MAX_CHUNKS = "max_chunks"140    FILE_UPLOAD_LIMIT = "file_upload_limit"141    SUPPORTED_FILE_EXTENSIONS = "supported_file_extensions"142    MAX_CHARACTERS_PER_CHUNK = "max_characters_per_chunk"143    DEFAULT_VOICE = "default_voice"144    VOICES = "voices"145    WORD_LIMIT = "word_limit"146    AUDIO_TYPE = "audio_type"147    MAX_WORKERS = "max_workers"148 149 150class ProviderModel(BaseModel):151    """152    Model class for provider model.153    """154 155    model: str156    label: I18nObject157    model_type: ModelType158    features: Optional[list[ModelFeature]] = None159    fetch_from: FetchFrom160    model_properties: dict[ModelPropertyKey, Any]161    deprecated: bool = False162    model_config = ConfigDict(protected_namespaces=())163 164 165class ParameterRule(BaseModel):166    """167    Model class for parameter rule.168    """169 170    name: str171    use_template: Optional[str] = None172    label: I18nObject173    type: ParameterType174    help: Optional[I18nObject] = None175    required: bool = False176    default: Optional[Any] = None177    min: Optional[float] = None178    max: Optional[float] = None179    precision: Optional[int] = None180    options: list[str] = []181 182 183class PriceConfig(BaseModel):184    """185    Model class for pricing info.186    """187 188    input: Decimal189    output: Optional[Decimal] = None190    unit: Decimal191    currency: str192 193 194class AIModelEntity(ProviderModel):195    """196    Model class for AI model.197    """198 199    parameter_rules: list[ParameterRule] = []200    pricing: Optional[PriceConfig] = None201 202 203class ModelUsage(BaseModel):204    pass205 206 207class PriceType(Enum):208    """209    Enum class for price type.210    """211 212    INPUT = "input"213    OUTPUT = "output"214 215 216class PriceInfo(BaseModel):217    """218    Model class for price info.219    """220 221    unit_price: Decimal222    unit: Decimal223    total_amount: Decimal224    currency: str225