tools.py 11 KB

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