tool_entities.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473
  1. import base64
  2. from enum import Enum
  3. from typing import Any, Optional, Union, cast
  4. from pydantic import BaseModel, Field, field_serializer, field_validator
  5. from core.entities.parameter_entities import AppSelectorScope, CommonParameterType, ModelConfigScope
  6. from core.tools.entities.common_entities import I18nObject
  7. class ToolLabelEnum(Enum):
  8. SEARCH = "search"
  9. IMAGE = "image"
  10. VIDEOS = "videos"
  11. WEATHER = "weather"
  12. FINANCE = "finance"
  13. DESIGN = "design"
  14. TRAVEL = "travel"
  15. SOCIAL = "social"
  16. NEWS = "news"
  17. MEDICAL = "medical"
  18. PRODUCTIVITY = "productivity"
  19. EDUCATION = "education"
  20. BUSINESS = "business"
  21. ENTERTAINMENT = "entertainment"
  22. UTILITIES = "utilities"
  23. OTHER = "other"
  24. class ToolProviderType(str, Enum):
  25. """
  26. Enum class for tool provider
  27. """
  28. BUILT_IN = "builtin"
  29. WORKFLOW = "workflow"
  30. API = "api"
  31. APP = "app"
  32. DATASET_RETRIEVAL = "dataset-retrieval"
  33. @classmethod
  34. def value_of(cls, value: str) -> "ToolProviderType":
  35. """
  36. Get value of given mode.
  37. :param value: mode value
  38. :return: mode
  39. """
  40. for mode in cls:
  41. if mode.value == value:
  42. return mode
  43. raise ValueError(f"invalid mode value {value}")
  44. class ApiProviderSchemaType(Enum):
  45. """
  46. Enum class for api provider schema type.
  47. """
  48. OPENAPI = "openapi"
  49. SWAGGER = "swagger"
  50. OPENAI_PLUGIN = "openai_plugin"
  51. OPENAI_ACTIONS = "openai_actions"
  52. @classmethod
  53. def value_of(cls, value: str) -> "ApiProviderSchemaType":
  54. """
  55. Get value of given mode.
  56. :param value: mode value
  57. :return: mode
  58. """
  59. for mode in cls:
  60. if mode.value == value:
  61. return mode
  62. raise ValueError(f"invalid mode value {value}")
  63. class ApiProviderAuthType(Enum):
  64. """
  65. Enum class for api provider auth type.
  66. """
  67. NONE = "none"
  68. API_KEY = "api_key"
  69. @classmethod
  70. def value_of(cls, value: str) -> "ApiProviderAuthType":
  71. """
  72. Get value of given mode.
  73. :param value: mode value
  74. :return: mode
  75. """
  76. for mode in cls:
  77. if mode.value == value:
  78. return mode
  79. raise ValueError(f"invalid mode value {value}")
  80. class ToolInvokeMessage(BaseModel):
  81. class TextMessage(BaseModel):
  82. text: str
  83. class JsonMessage(BaseModel):
  84. json_object: dict
  85. class BlobMessage(BaseModel):
  86. blob: bytes
  87. class VariableMessage(BaseModel):
  88. variable_name: str = Field(..., description="The name of the variable")
  89. variable_value: str = Field(..., description="The value of the variable")
  90. stream: bool = Field(default=False, description="Whether the variable is streamed")
  91. @field_validator("variable_value", mode="before")
  92. @classmethod
  93. def transform_variable_value(cls, value, values) -> Any:
  94. """
  95. Only basic types and lists are allowed.
  96. """
  97. if not isinstance(value, dict | list | str | int | float | bool):
  98. raise ValueError("Only basic types and lists are allowed.")
  99. # if stream is true, the value must be a string
  100. if values.get('stream'):
  101. if not isinstance(value, str):
  102. raise ValueError("When 'stream' is True, 'variable_value' must be a string.")
  103. return value
  104. @field_validator("variable_name", mode="before")
  105. @classmethod
  106. def transform_variable_name(cls, value) -> str:
  107. """
  108. The variable name must be a string.
  109. """
  110. if value in {"json", "text", "files"}:
  111. raise ValueError(f"The variable name '{value}' is reserved.")
  112. return value
  113. class MessageType(Enum):
  114. TEXT = "text"
  115. IMAGE = "image"
  116. LINK = "link"
  117. BLOB = "blob"
  118. JSON = "json"
  119. IMAGE_LINK = "image_link"
  120. FILE_VAR = "file_var"
  121. VARIABLE = "variable"
  122. type: MessageType = MessageType.TEXT
  123. """
  124. plain text, image url or link url
  125. """
  126. message: JsonMessage | TextMessage | BlobMessage | VariableMessage | None
  127. meta: dict[str, Any] | None = None
  128. save_as: str = ""
  129. @field_validator('message', mode='before')
  130. @classmethod
  131. def decode_blob_message(cls, v):
  132. if isinstance(v, dict) and 'blob' in v:
  133. try:
  134. v['blob'] = base64.b64decode(v['blob'])
  135. except Exception:
  136. pass
  137. return v
  138. @field_serializer('message')
  139. def serialize_message(self, v):
  140. if isinstance(v, self.BlobMessage):
  141. return {
  142. 'blob': base64.b64encode(v.blob).decode('utf-8')
  143. }
  144. return v
  145. class ToolInvokeMessageBinary(BaseModel):
  146. mimetype: str = Field(..., description="The mimetype of the binary")
  147. url: str = Field(..., description="The url of the binary")
  148. save_as: str = ""
  149. file_var: Optional[dict[str, Any]] = None
  150. class ToolParameterOption(BaseModel):
  151. value: str = Field(..., description="The value of the option")
  152. label: I18nObject = Field(..., description="The label of the option")
  153. @field_validator("value", mode="before")
  154. @classmethod
  155. def transform_id_to_str(cls, value) -> str:
  156. if not isinstance(value, str):
  157. return str(value)
  158. else:
  159. return value
  160. class ToolParameter(BaseModel):
  161. class ToolParameterType(str, Enum):
  162. STRING = CommonParameterType.STRING.value
  163. NUMBER = CommonParameterType.NUMBER.value
  164. BOOLEAN = CommonParameterType.BOOLEAN.value
  165. SELECT = CommonParameterType.SELECT.value
  166. SECRET_INPUT = CommonParameterType.SECRET_INPUT.value
  167. FILE = CommonParameterType.FILE.value
  168. class ToolParameterForm(Enum):
  169. SCHEMA = "schema" # should be set while adding tool
  170. FORM = "form" # should be set before invoking tool
  171. LLM = "llm" # will be set by LLM
  172. name: str = Field(..., description="The name of the parameter")
  173. label: I18nObject = Field(..., description="The label presented to the user")
  174. human_description: Optional[I18nObject] = Field(default=None, description="The description presented to the user")
  175. placeholder: Optional[I18nObject] = Field(default=None, description="The placeholder presented to the user")
  176. type: ToolParameterType = Field(..., description="The type of the parameter")
  177. scope: AppSelectorScope | ModelConfigScope | None = None
  178. form: ToolParameterForm = Field(..., description="The form of the parameter, schema/form/llm")
  179. llm_description: Optional[str] = None
  180. required: Optional[bool] = False
  181. default: Optional[Union[float, int, str]] = None
  182. min: Optional[Union[float, int]] = None
  183. max: Optional[Union[float, int]] = None
  184. options: list[ToolParameterOption] = Field(default_factory=list)
  185. @classmethod
  186. def get_simple_instance(
  187. cls,
  188. name: str,
  189. llm_description: str,
  190. type: ToolParameterType,
  191. required: bool,
  192. options: Optional[list[str]] = None,
  193. ) -> "ToolParameter":
  194. """
  195. get a simple tool parameter
  196. :param name: the name of the parameter
  197. :param llm_description: the description presented to the LLM
  198. :param type: the type of the parameter
  199. :param required: if the parameter is required
  200. :param options: the options of the parameter
  201. """
  202. # convert options to ToolParameterOption
  203. if options:
  204. option_objs = [
  205. ToolParameterOption(value=option, label=I18nObject(en_US=option, zh_Hans=option)) for option in options
  206. ]
  207. else:
  208. option_objs = []
  209. return cls(
  210. name=name,
  211. label=I18nObject(en_US='', zh_Hans=''),
  212. placeholder=None,
  213. human_description=I18nObject(en_US='', zh_Hans=''),
  214. type=type,
  215. form=cls.ToolParameterForm.LLM,
  216. llm_description=llm_description,
  217. required=required,
  218. options=option_objs,
  219. )
  220. class ToolProviderIdentity(BaseModel):
  221. author: str = Field(..., description="The author of the tool")
  222. name: str = Field(..., description="The name of the tool")
  223. description: I18nObject = Field(..., description="The description of the tool")
  224. icon: str = Field(..., description="The icon of the tool")
  225. label: I18nObject = Field(..., description="The label of the tool")
  226. tags: Optional[list[ToolLabelEnum]] = Field(
  227. default=[],
  228. description="The tags of the tool",
  229. )
  230. class ToolDescription(BaseModel):
  231. human: I18nObject = Field(..., description="The description presented to the user")
  232. llm: str = Field(..., description="The description presented to the LLM")
  233. class ToolIdentity(BaseModel):
  234. author: str = Field(..., description="The author of the tool")
  235. name: str = Field(..., description="The name of the tool")
  236. label: I18nObject = Field(..., description="The label of the tool")
  237. provider: str = Field(..., description="The provider of the tool")
  238. icon: Optional[str] = None
  239. class ToolRuntimeVariableType(Enum):
  240. TEXT = "text"
  241. IMAGE = "image"
  242. class ToolRuntimeVariable(BaseModel):
  243. type: ToolRuntimeVariableType = Field(..., description="The type of the variable")
  244. name: str = Field(..., description="The name of the variable")
  245. position: int = Field(..., description="The position of the variable")
  246. tool_name: str = Field(..., description="The name of the tool")
  247. class ToolRuntimeTextVariable(ToolRuntimeVariable):
  248. value: str = Field(..., description="The value of the variable")
  249. class ToolRuntimeImageVariable(ToolRuntimeVariable):
  250. value: str = Field(..., description="The path of the image")
  251. class ToolRuntimeVariablePool(BaseModel):
  252. conversation_id: str = Field(..., description="The conversation id")
  253. user_id: str = Field(..., description="The user id")
  254. tenant_id: str = Field(..., description="The tenant id of assistant")
  255. pool: list[ToolRuntimeVariable] = Field(..., description="The pool of variables")
  256. def __init__(self, **data: Any):
  257. pool = data.get("pool", [])
  258. # convert pool into correct type
  259. for index, variable in enumerate(pool):
  260. if variable["type"] == ToolRuntimeVariableType.TEXT.value:
  261. pool[index] = ToolRuntimeTextVariable(**variable)
  262. elif variable["type"] == ToolRuntimeVariableType.IMAGE.value:
  263. pool[index] = ToolRuntimeImageVariable(**variable)
  264. super().__init__(**data)
  265. def dict(self) -> dict:
  266. return {
  267. "conversation_id": self.conversation_id,
  268. "user_id": self.user_id,
  269. "tenant_id": self.tenant_id,
  270. "pool": [variable.model_dump() for variable in self.pool],
  271. }
  272. def set_text(self, tool_name: str, name: str, value: str) -> None:
  273. """
  274. set a text variable
  275. """
  276. for variable in self.pool:
  277. if variable.name == name:
  278. if variable.type == ToolRuntimeVariableType.TEXT:
  279. variable = cast(ToolRuntimeTextVariable, variable)
  280. variable.value = value
  281. return
  282. variable = ToolRuntimeTextVariable(
  283. type=ToolRuntimeVariableType.TEXT,
  284. name=name,
  285. position=len(self.pool),
  286. tool_name=tool_name,
  287. value=value,
  288. )
  289. self.pool.append(variable)
  290. def set_file(self, tool_name: str, value: str, name: Optional[str] = None) -> None:
  291. """
  292. set an image variable
  293. :param tool_name: the name of the tool
  294. :param value: the id of the file
  295. """
  296. # check how many image variables are there
  297. image_variable_count = 0
  298. for variable in self.pool:
  299. if variable.type == ToolRuntimeVariableType.IMAGE:
  300. image_variable_count += 1
  301. if name is None:
  302. name = f"file_{image_variable_count}"
  303. for variable in self.pool:
  304. if variable.name == name:
  305. if variable.type == ToolRuntimeVariableType.IMAGE:
  306. variable = cast(ToolRuntimeImageVariable, variable)
  307. variable.value = value
  308. return
  309. variable = ToolRuntimeImageVariable(
  310. type=ToolRuntimeVariableType.IMAGE,
  311. name=name,
  312. position=len(self.pool),
  313. tool_name=tool_name,
  314. value=value,
  315. )
  316. self.pool.append(variable)
  317. class ModelToolPropertyKey(Enum):
  318. IMAGE_PARAMETER_NAME = "image_parameter_name"
  319. class ModelToolConfiguration(BaseModel):
  320. """
  321. Model tool configuration
  322. """
  323. type: str = Field(..., description="The type of the model tool")
  324. model: str = Field(..., description="The model")
  325. label: I18nObject = Field(..., description="The label of the model tool")
  326. properties: dict[ModelToolPropertyKey, Any] = Field(..., description="The properties of the model tool")
  327. class ModelToolProviderConfiguration(BaseModel):
  328. """
  329. Model tool provider configuration
  330. """
  331. provider: str = Field(..., description="The provider of the model tool")
  332. models: list[ModelToolConfiguration] = Field(..., description="The models of the model tool")
  333. label: I18nObject = Field(..., description="The label of the model tool")
  334. class WorkflowToolParameterConfiguration(BaseModel):
  335. """
  336. Workflow tool configuration
  337. """
  338. name: str = Field(..., description="The name of the parameter")
  339. description: str = Field(..., description="The description of the parameter")
  340. form: ToolParameter.ToolParameterForm = Field(..., description="The form of the parameter")
  341. class ToolInvokeMeta(BaseModel):
  342. """
  343. Tool invoke meta
  344. """
  345. time_cost: float = Field(..., description="The time cost of the tool invoke")
  346. error: Optional[str] = None
  347. tool_config: Optional[dict] = None
  348. @classmethod
  349. def empty(cls) -> "ToolInvokeMeta":
  350. """
  351. Get an empty instance of ToolInvokeMeta
  352. """
  353. return cls(time_cost=0.0, error=None, tool_config={})
  354. @classmethod
  355. def error_instance(cls, error: str) -> "ToolInvokeMeta":
  356. """
  357. Get an instance of ToolInvokeMeta with error
  358. """
  359. return cls(time_cost=0.0, error=error, tool_config={})
  360. def to_dict(self) -> dict:
  361. return {
  362. "time_cost": self.time_cost,
  363. "error": self.error,
  364. "tool_config": self.tool_config,
  365. }
  366. class ToolLabel(BaseModel):
  367. """
  368. Tool label
  369. """
  370. name: str = Field(..., description="The name of the tool")
  371. label: I18nObject = Field(..., description="The label of the tool")
  372. icon: str = Field(..., description="The icon of the tool")
  373. class ToolInvokeFrom(Enum):
  374. """
  375. Enum class for tool invoke
  376. """
  377. WORKFLOW = "workflow"
  378. AGENT = "agent"