tools.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306
  1. import json
  2. from sqlalchemy import ForeignKey
  3. from sqlalchemy.orm import Mapped, mapped_column
  4. from core.tools.entities.common_entities import I18nObject
  5. from core.tools.entities.tool_bundle import ApiToolBundle
  6. from core.tools.entities.tool_entities import ApiProviderSchemaType, WorkflowToolParameterConfiguration
  7. from extensions.ext_database import db
  8. from models.base import Base
  9. from .model import Account, App, Tenant
  10. from .types import StringUUID
  11. class BuiltinToolProvider(db.Model):
  12. """
  13. This table stores the tool provider information for built-in tools for each tenant.
  14. """
  15. __tablename__ = 'tool_builtin_providers'
  16. __table_args__ = (
  17. db.PrimaryKeyConstraint('id', name='tool_builtin_provider_pkey'),
  18. # one tenant can only have one tool provider with the same name
  19. db.UniqueConstraint('tenant_id', 'provider', name='unique_builtin_tool_provider')
  20. )
  21. # id of the tool provider
  22. id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
  23. # id of the tenant
  24. tenant_id = db.Column(StringUUID, nullable=True)
  25. # who created this tool provider
  26. user_id = db.Column(StringUUID, nullable=False)
  27. # name of the tool provider
  28. provider = db.Column(db.String(40), nullable=False)
  29. # credential of the tool provider
  30. encrypted_credentials = db.Column(db.Text, nullable=True)
  31. created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  32. updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  33. @property
  34. def credentials(self) -> dict:
  35. return json.loads(self.encrypted_credentials)
  36. class PublishedAppTool(db.Model):
  37. """
  38. The table stores the apps published as a tool for each person.
  39. """
  40. __tablename__ = 'tool_published_apps'
  41. __table_args__ = (
  42. db.PrimaryKeyConstraint('id', name='published_app_tool_pkey'),
  43. db.UniqueConstraint('app_id', 'user_id', name='unique_published_app_tool')
  44. )
  45. # id of the tool provider
  46. id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
  47. # id of the app
  48. app_id = db.Column(StringUUID, ForeignKey('apps.id'), nullable=False)
  49. # who published this tool
  50. user_id = db.Column(StringUUID, nullable=False)
  51. # description of the tool, stored in i18n format, for human
  52. description = db.Column(db.Text, nullable=False)
  53. # llm_description of the tool, for LLM
  54. llm_description = db.Column(db.Text, nullable=False)
  55. # query description, query will be seem as a parameter of the tool, to describe this parameter to llm, we need this field
  56. query_description = db.Column(db.Text, nullable=False)
  57. # query name, the name of the query parameter
  58. query_name = db.Column(db.String(40), nullable=False)
  59. # name of the tool provider
  60. tool_name = db.Column(db.String(40), nullable=False)
  61. # author
  62. author = db.Column(db.String(40), nullable=False)
  63. created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  64. updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  65. @property
  66. def description_i18n(self) -> I18nObject:
  67. return I18nObject(**json.loads(self.description))
  68. @property
  69. def app(self) -> App | None:
  70. return db.session.query(App).filter(App.id == self.app_id).first()
  71. class ApiToolProvider(Base):
  72. """
  73. The table stores the api providers.
  74. """
  75. __tablename__ = 'tool_api_providers'
  76. __table_args__ = (
  77. db.PrimaryKeyConstraint('id', name='tool_api_provider_pkey'),
  78. db.UniqueConstraint('name', 'tenant_id', name='unique_api_tool_provider')
  79. )
  80. id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
  81. # name of the api provider
  82. name = db.Column(db.String(40), nullable=False)
  83. # icon
  84. icon = db.Column(db.String(255), nullable=False)
  85. # original schema
  86. schema = db.Column(db.Text, nullable=False)
  87. schema_type_str = db.Column(db.String(40), nullable=False)
  88. # who created this tool
  89. user_id = db.Column(StringUUID, nullable=False)
  90. # tenant id
  91. tenant_id = db.Column(StringUUID, nullable=False)
  92. # description of the provider
  93. description = db.Column(db.Text, nullable=False)
  94. # json format tools
  95. tools_str = db.Column(db.Text, nullable=False)
  96. # json format credentials
  97. credentials_str = db.Column(db.Text, nullable=False)
  98. # privacy policy
  99. privacy_policy = db.Column(db.String(255), nullable=True)
  100. # custom_disclaimer
  101. custom_disclaimer = db.Column(db.String(255), nullable=True)
  102. created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  103. updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  104. @property
  105. def schema_type(self) -> ApiProviderSchemaType:
  106. return ApiProviderSchemaType.value_of(self.schema_type_str)
  107. @property
  108. def tools(self) -> list[ApiToolBundle]:
  109. return [ApiToolBundle(**tool) for tool in json.loads(self.tools_str)]
  110. @property
  111. def credentials(self) -> dict:
  112. return json.loads(self.credentials_str)
  113. @property
  114. def user(self) -> Account:
  115. return db.session.query(Account).filter(Account.id == self.user_id).first()
  116. @property
  117. def tenant(self) -> Tenant:
  118. return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first()
  119. class ToolLabelBinding(db.Model):
  120. """
  121. The table stores the labels for tools.
  122. """
  123. __tablename__ = 'tool_label_bindings'
  124. __table_args__ = (
  125. db.PrimaryKeyConstraint('id', name='tool_label_bind_pkey'),
  126. db.UniqueConstraint('tool_id', 'label_name', name='unique_tool_label_bind'),
  127. )
  128. id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
  129. # tool id
  130. tool_id = db.Column(db.String(64), nullable=False)
  131. # tool type
  132. tool_type = db.Column(db.String(40), nullable=False)
  133. # label name
  134. label_name = db.Column(db.String(40), nullable=False)
  135. class WorkflowToolProvider(db.Model):
  136. """
  137. The table stores the workflow providers.
  138. """
  139. __tablename__ = 'tool_workflow_providers'
  140. __table_args__ = (
  141. db.PrimaryKeyConstraint('id', name='tool_workflow_provider_pkey'),
  142. db.UniqueConstraint('name', 'tenant_id', name='unique_workflow_tool_provider'),
  143. db.UniqueConstraint('tenant_id', 'app_id', name='unique_workflow_tool_provider_app_id'),
  144. )
  145. id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
  146. # name of the workflow provider
  147. name = db.Column(db.String(40), nullable=False)
  148. # label of the workflow provider
  149. label = db.Column(db.String(255), nullable=False, server_default='')
  150. # icon
  151. icon = db.Column(db.String(255), nullable=False)
  152. # app id of the workflow provider
  153. app_id = db.Column(StringUUID, nullable=False)
  154. # version of the workflow provider
  155. version = db.Column(db.String(255), nullable=False, server_default='')
  156. # who created this tool
  157. user_id = db.Column(StringUUID, nullable=False)
  158. # tenant id
  159. tenant_id = db.Column(StringUUID, nullable=False)
  160. # description of the provider
  161. description = db.Column(db.Text, nullable=False)
  162. # parameter configuration
  163. parameter_configuration = db.Column(db.Text, nullable=False, server_default='[]')
  164. # privacy policy
  165. privacy_policy = db.Column(db.String(255), nullable=True, server_default='')
  166. created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  167. updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  168. @property
  169. def schema_type(self) -> ApiProviderSchemaType:
  170. return ApiProviderSchemaType.value_of(self.schema_type_str)
  171. @property
  172. def user(self) -> Account:
  173. return db.session.query(Account).filter(Account.id == self.user_id).first()
  174. @property
  175. def tenant(self) -> Tenant:
  176. return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first()
  177. @property
  178. def parameter_configurations(self) -> list[WorkflowToolParameterConfiguration]:
  179. return [
  180. WorkflowToolParameterConfiguration(**config)
  181. for config in json.loads(self.parameter_configuration)
  182. ]
  183. @property
  184. def app(self) -> App:
  185. return db.session.query(App).filter(App.id == self.app_id).first()
  186. class ToolModelInvoke(db.Model):
  187. """
  188. store the invoke logs from tool invoke
  189. """
  190. __tablename__ = "tool_model_invokes"
  191. __table_args__ = (
  192. db.PrimaryKeyConstraint('id', name='tool_model_invoke_pkey'),
  193. )
  194. id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
  195. # who invoke this tool
  196. user_id = db.Column(StringUUID, nullable=False)
  197. # tenant id
  198. tenant_id = db.Column(StringUUID, nullable=False)
  199. # provider
  200. provider = db.Column(db.String(40), nullable=False)
  201. # type
  202. tool_type = db.Column(db.String(40), nullable=False)
  203. # tool name
  204. tool_name = db.Column(db.String(40), nullable=False)
  205. # invoke parameters
  206. model_parameters = db.Column(db.Text, nullable=False)
  207. # prompt messages
  208. prompt_messages = db.Column(db.Text, nullable=False)
  209. # invoke response
  210. model_response = db.Column(db.Text, nullable=False)
  211. prompt_tokens = db.Column(db.Integer, nullable=False, server_default=db.text('0'))
  212. answer_tokens = db.Column(db.Integer, nullable=False, server_default=db.text('0'))
  213. answer_unit_price = db.Column(db.Numeric(10, 4), nullable=False)
  214. answer_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text('0.001'))
  215. provider_response_latency = db.Column(db.Float, nullable=False, server_default=db.text('0'))
  216. total_price = db.Column(db.Numeric(10, 7))
  217. currency = db.Column(db.String(255), nullable=False)
  218. created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  219. updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  220. class ToolConversationVariables(db.Model):
  221. """
  222. store the conversation variables from tool invoke
  223. """
  224. __tablename__ = "tool_conversation_variables"
  225. __table_args__ = (
  226. db.PrimaryKeyConstraint('id', name='tool_conversation_variables_pkey'),
  227. # add index for user_id and conversation_id
  228. db.Index('user_id_idx', 'user_id'),
  229. db.Index('conversation_id_idx', 'conversation_id'),
  230. )
  231. id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
  232. # conversation user id
  233. user_id = db.Column(StringUUID, nullable=False)
  234. # tenant id
  235. tenant_id = db.Column(StringUUID, nullable=False)
  236. # conversation id
  237. conversation_id = db.Column(StringUUID, nullable=False)
  238. # variables pool
  239. variables_str = db.Column(db.Text, nullable=False)
  240. created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  241. updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)'))
  242. @property
  243. def variables(self) -> dict:
  244. return json.loads(self.variables_str)
  245. class ToolFile(Base):
  246. """
  247. store the file created by agent
  248. """
  249. __tablename__ = "tool_files"
  250. __table_args__ = (
  251. db.PrimaryKeyConstraint('id', name='tool_file_pkey'),
  252. # add index for conversation_id
  253. db.Index('tool_file_conversation_id_idx', 'conversation_id'),
  254. )
  255. id: Mapped[str] = mapped_column(StringUUID, server_default=db.text('uuid_generate_v4()'))
  256. # conversation user id
  257. user_id: Mapped[str] = mapped_column(StringUUID)
  258. # tenant id
  259. tenant_id: Mapped[str] = mapped_column(StringUUID)
  260. # conversation id
  261. conversation_id: Mapped[str] = mapped_column(StringUUID, nullable=True)
  262. # file key
  263. file_key: Mapped[str] = mapped_column(db.String(255), nullable=False)
  264. # mime type
  265. mimetype: Mapped[str] = mapped_column(db.String(255), nullable=False)
  266. # original url
  267. original_url: Mapped[str] = mapped_column(db.String(2048), nullable=True)