model_entities.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203
  1. from decimal import Decimal
  2. from enum import Enum
  3. from typing import Any, Optional
  4. from core.model_runtime.entities.common_entities import I18nObject
  5. from pydantic import BaseModel
  6. class ModelType(Enum):
  7. """
  8. Enum class for model type.
  9. """
  10. LLM = "llm"
  11. TEXT_EMBEDDING = "text-embedding"
  12. RERANK = "rerank"
  13. SPEECH2TEXT = "speech2text"
  14. MODERATION = "moderation"
  15. TTS = "tts"
  16. # TEXT2IMG = "text2img"
  17. @classmethod
  18. def value_of(cls, origin_model_type: str) -> "ModelType":
  19. """
  20. Get model type from origin model type.
  21. :return: model type
  22. """
  23. if origin_model_type == 'text-generation' or origin_model_type == cls.LLM.value:
  24. return cls.LLM
  25. elif origin_model_type == 'embeddings' or origin_model_type == cls.TEXT_EMBEDDING.value:
  26. return cls.TEXT_EMBEDDING
  27. elif origin_model_type == 'reranking' or origin_model_type == cls.RERANK.value:
  28. return cls.RERANK
  29. elif origin_model_type == 'speech2text' or origin_model_type == cls.SPEECH2TEXT.value:
  30. return cls.SPEECH2TEXT
  31. elif origin_model_type == 'tts' or origin_model_type == cls.TTS.value:
  32. return cls.TTS
  33. elif origin_model_type == cls.MODERATION.value:
  34. return cls.MODERATION
  35. else:
  36. raise ValueError(f'invalid origin model type {origin_model_type}')
  37. def to_origin_model_type(self) -> str:
  38. """
  39. Get origin model type from model type.
  40. :return: origin model type
  41. """
  42. if self == self.LLM:
  43. return 'text-generation'
  44. elif self == self.TEXT_EMBEDDING:
  45. return 'embeddings'
  46. elif self == self.RERANK:
  47. return 'reranking'
  48. elif self == self.SPEECH2TEXT:
  49. return 'speech2text'
  50. elif self == self.TTS:
  51. return 'tts'
  52. elif self == self.MODERATION:
  53. return 'moderation'
  54. else:
  55. raise ValueError(f'invalid model type {self}')
  56. class FetchFrom(Enum):
  57. """
  58. Enum class for fetch from.
  59. """
  60. PREDEFINED_MODEL = "predefined-model"
  61. CUSTOMIZABLE_MODEL = "customizable-model"
  62. class ModelFeature(Enum):
  63. """
  64. Enum class for llm feature.
  65. """
  66. TOOL_CALL = "tool-call"
  67. MULTI_TOOL_CALL = "multi-tool-call"
  68. AGENT_THOUGHT = "agent-thought"
  69. VISION = "vision"
  70. class DefaultParameterName(Enum):
  71. """
  72. Enum class for parameter template variable.
  73. """
  74. TEMPERATURE = "temperature"
  75. TOP_P = "top_p"
  76. PRESENCE_PENALTY = "presence_penalty"
  77. FREQUENCY_PENALTY = "frequency_penalty"
  78. MAX_TOKENS = "max_tokens"
  79. @classmethod
  80. def value_of(cls, value: Any) -> 'DefaultParameterName':
  81. """
  82. Get parameter name from value.
  83. :param value: parameter value
  84. :return: parameter name
  85. """
  86. for name in cls:
  87. if name.value == value:
  88. return name
  89. raise ValueError(f'invalid parameter name {value}')
  90. class ParameterType(Enum):
  91. """
  92. Enum class for parameter type.
  93. """
  94. FLOAT = "float"
  95. INT = "int"
  96. STRING = "string"
  97. BOOLEAN = "boolean"
  98. class ModelPropertyKey(Enum):
  99. """
  100. Enum class for model property key.
  101. """
  102. MODE = "mode"
  103. CONTEXT_SIZE = "context_size"
  104. MAX_CHUNKS = "max_chunks"
  105. FILE_UPLOAD_LIMIT = "file_upload_limit"
  106. SUPPORTED_FILE_EXTENSIONS = "supported_file_extensions"
  107. MAX_CHARACTERS_PER_CHUNK = "max_characters_per_chunk"
  108. DEFAULT_VOICE = "default_voice"
  109. WORD_LIMIT = "word_limit"
  110. AUDOI_TYPE = "audio_type"
  111. MAX_WORKERS = "max_workers"
  112. class ProviderModel(BaseModel):
  113. """
  114. Model class for provider model.
  115. """
  116. model: str
  117. label: I18nObject
  118. model_type: ModelType
  119. features: Optional[list[ModelFeature]] = None
  120. fetch_from: FetchFrom
  121. model_properties: dict[ModelPropertyKey, Any]
  122. deprecated: bool = False
  123. class Config:
  124. protected_namespaces = ()
  125. class ParameterRule(BaseModel):
  126. """
  127. Model class for parameter rule.
  128. """
  129. name: str
  130. use_template: Optional[str] = None
  131. label: I18nObject
  132. type: ParameterType
  133. help: Optional[I18nObject] = None
  134. required: bool = False
  135. default: Optional[Any] = None
  136. min: Optional[float] = None
  137. max: Optional[float] = None
  138. precision: Optional[int] = None
  139. options: list[str] = []
  140. class PriceConfig(BaseModel):
  141. """
  142. Model class for pricing info.
  143. """
  144. input: Decimal
  145. output: Optional[Decimal] = None
  146. unit: Decimal
  147. currency: str
  148. class AIModelEntity(ProviderModel):
  149. """
  150. Model class for AI model.
  151. """
  152. parameter_rules: list[ParameterRule] = []
  153. pricing: Optional[PriceConfig] = None
  154. class ModelUsage(BaseModel):
  155. pass
  156. class PriceType(Enum):
  157. """
  158. Enum class for price type.
  159. """
  160. INPUT = "input"
  161. OUTPUT = "output"
  162. class PriceInfo(BaseModel):
  163. """
  164. Model class for price info.
  165. """
  166. unit_price: Decimal
  167. unit: Decimal
  168. total_amount: Decimal
  169. currency: str