entities.py 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242
  1. from enum import Enum
  2. from typing import Any, Optional
  3. from pydantic import BaseModel
  4. from core.model_runtime.entities.message_entities import PromptMessageRole
  5. from models.model import AppMode
  6. class ModelConfigEntity(BaseModel):
  7. """
  8. Model Config Entity.
  9. """
  10. provider: str
  11. model: str
  12. mode: Optional[str] = None
  13. parameters: dict[str, Any] = {}
  14. stop: list[str] = []
  15. class AdvancedChatMessageEntity(BaseModel):
  16. """
  17. Advanced Chat Message Entity.
  18. """
  19. text: str
  20. role: PromptMessageRole
  21. class AdvancedChatPromptTemplateEntity(BaseModel):
  22. """
  23. Advanced Chat Prompt Template Entity.
  24. """
  25. messages: list[AdvancedChatMessageEntity]
  26. class AdvancedCompletionPromptTemplateEntity(BaseModel):
  27. """
  28. Advanced Completion Prompt Template Entity.
  29. """
  30. class RolePrefixEntity(BaseModel):
  31. """
  32. Role Prefix Entity.
  33. """
  34. user: str
  35. assistant: str
  36. prompt: str
  37. role_prefix: Optional[RolePrefixEntity] = None
  38. class PromptTemplateEntity(BaseModel):
  39. """
  40. Prompt Template Entity.
  41. """
  42. class PromptType(Enum):
  43. """
  44. Prompt Type.
  45. 'simple', 'advanced'
  46. """
  47. SIMPLE = 'simple'
  48. ADVANCED = 'advanced'
  49. @classmethod
  50. def value_of(cls, value: str) -> 'PromptType':
  51. """
  52. Get value of given mode.
  53. :param value: mode value
  54. :return: mode
  55. """
  56. for mode in cls:
  57. if mode.value == value:
  58. return mode
  59. raise ValueError(f'invalid prompt type value {value}')
  60. prompt_type: PromptType
  61. simple_prompt_template: Optional[str] = None
  62. advanced_chat_prompt_template: Optional[AdvancedChatPromptTemplateEntity] = None
  63. advanced_completion_prompt_template: Optional[AdvancedCompletionPromptTemplateEntity] = None
  64. class VariableEntity(BaseModel):
  65. """
  66. Variable Entity.
  67. """
  68. class Type(Enum):
  69. TEXT_INPUT = 'text-input'
  70. SELECT = 'select'
  71. PARAGRAPH = 'paragraph'
  72. NUMBER = 'number'
  73. @classmethod
  74. def value_of(cls, value: str) -> 'VariableEntity.Type':
  75. """
  76. Get value of given mode.
  77. :param value: mode value
  78. :return: mode
  79. """
  80. for mode in cls:
  81. if mode.value == value:
  82. return mode
  83. raise ValueError(f'invalid variable type value {value}')
  84. variable: str
  85. label: str
  86. description: Optional[str] = None
  87. type: Type
  88. required: bool = False
  89. max_length: Optional[int] = None
  90. options: Optional[list[str]] = None
  91. default: Optional[str] = None
  92. hint: Optional[str] = None
  93. class ExternalDataVariableEntity(BaseModel):
  94. """
  95. External Data Variable Entity.
  96. """
  97. variable: str
  98. type: str
  99. config: dict[str, Any] = {}
  100. class DatasetRetrieveConfigEntity(BaseModel):
  101. """
  102. Dataset Retrieve Config Entity.
  103. """
  104. class RetrieveStrategy(Enum):
  105. """
  106. Dataset Retrieve Strategy.
  107. 'single' or 'multiple'
  108. """
  109. SINGLE = 'single'
  110. MULTIPLE = 'multiple'
  111. @classmethod
  112. def value_of(cls, value: str) -> 'RetrieveStrategy':
  113. """
  114. Get value of given mode.
  115. :param value: mode value
  116. :return: mode
  117. """
  118. for mode in cls:
  119. if mode.value == value:
  120. return mode
  121. raise ValueError(f'invalid retrieve strategy value {value}')
  122. query_variable: Optional[str] = None # Only when app mode is completion
  123. retrieve_strategy: RetrieveStrategy
  124. top_k: Optional[int] = None
  125. score_threshold: Optional[float] = None
  126. reranking_model: Optional[dict] = None
  127. class DatasetEntity(BaseModel):
  128. """
  129. Dataset Config Entity.
  130. """
  131. dataset_ids: list[str]
  132. retrieve_config: DatasetRetrieveConfigEntity
  133. class SensitiveWordAvoidanceEntity(BaseModel):
  134. """
  135. Sensitive Word Avoidance Entity.
  136. """
  137. type: str
  138. config: dict[str, Any] = {}
  139. class TextToSpeechEntity(BaseModel):
  140. """
  141. Sensitive Word Avoidance Entity.
  142. """
  143. enabled: bool
  144. voice: Optional[str] = None
  145. language: Optional[str] = None
  146. class FileExtraConfig(BaseModel):
  147. """
  148. File Upload Entity.
  149. """
  150. image_config: Optional[dict[str, Any]] = None
  151. class AppAdditionalFeatures(BaseModel):
  152. file_upload: Optional[FileExtraConfig] = None
  153. opening_statement: Optional[str] = None
  154. suggested_questions: list[str] = []
  155. suggested_questions_after_answer: bool = False
  156. show_retrieve_source: bool = False
  157. more_like_this: bool = False
  158. speech_to_text: bool = False
  159. text_to_speech: Optional[TextToSpeechEntity] = None
  160. class AppConfig(BaseModel):
  161. """
  162. Application Config Entity.
  163. """
  164. tenant_id: str
  165. app_id: str
  166. app_mode: AppMode
  167. additional_features: AppAdditionalFeatures
  168. variables: list[VariableEntity] = []
  169. sensitive_word_avoidance: Optional[SensitiveWordAvoidanceEntity] = None
  170. class EasyUIBasedAppModelConfigFrom(Enum):
  171. """
  172. App Model Config From.
  173. """
  174. ARGS = 'args'
  175. APP_LATEST_CONFIG = 'app-latest-config'
  176. CONVERSATION_SPECIFIC_CONFIG = 'conversation-specific-config'
  177. class EasyUIBasedAppConfig(AppConfig):
  178. """
  179. Easy UI Based App Config Entity.
  180. """
  181. app_model_config_from: EasyUIBasedAppModelConfigFrom
  182. app_model_config_id: str
  183. app_model_config_dict: dict
  184. model: ModelConfigEntity
  185. prompt_template: PromptTemplateEntity
  186. dataset: Optional[DatasetEntity] = None
  187. external_data_variables: list[ExternalDataVariableEntity] = []
  188. class WorkflowUIBasedAppConfig(AppConfig):
  189. """
  190. Workflow UI Based App Config Entity.
  191. """
  192. workflow_id: str