estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

replace db with sa to get typing support (#23240)

Browse Source
This commit is contained in:
Asuka Minato
2025-08-03 00:54:23 +09:00
committed by GitHub
parent ff9fd0cdb2
commit 58608f51da
17 changed files with 623 additions and 613 deletions
Download Patch File Download Diff File Expand all files Collapse all files

152
api/models/tools.py
View File

@@ -25,33 +25,33 @@ from .types import StringUUID
class ToolOAuthSystemClient(Base):
__tablename__ = "tool_oauth_system_clients"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="tool_oauth_system_client_pkey"),
db.UniqueConstraint("plugin_id", "provider", name="tool_oauth_system_client_plugin_id_provider_idx"),
sa.PrimaryKeyConstraint("id", name="tool_oauth_system_client_pkey"),
sa.UniqueConstraint("plugin_id", "provider", name="tool_oauth_system_client_plugin_id_provider_idx"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
plugin_id = mapped_column(String(512), nullable=False)
provider: Mapped[str] = mapped_column(String(255), nullable=False)
# oauth params of the tool provider
encrypted_oauth_params: Mapped[str] = mapped_column(db.Text, nullable=False)
encrypted_oauth_params: Mapped[str] = mapped_column(sa.Text, nullable=False)
# tenant level tool oauth client params (client_id, client_secret, etc.)
class ToolOAuthTenantClient(Base):
__tablename__ = "tool_oauth_tenant_clients"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="tool_oauth_tenant_client_pkey"),
db.UniqueConstraint("tenant_id", "plugin_id", "provider", name="unique_tool_oauth_tenant_client"),
sa.PrimaryKeyConstraint("id", name="tool_oauth_tenant_client_pkey"),
sa.UniqueConstraint("tenant_id", "plugin_id", "provider", name="unique_tool_oauth_tenant_client"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
# tenant id
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
plugin_id: Mapped[str] = mapped_column(String(512), nullable=False)
provider: Mapped[str] = mapped_column(String(255), nullable=False)
enabled: Mapped[bool] = mapped_column(db.Boolean, nullable=False, server_default=db.text("true"))
enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("true"))
# oauth params of the tool provider
encrypted_oauth_params: Mapped[str] = mapped_column(db.Text, nullable=False)
encrypted_oauth_params: Mapped[str] = mapped_column(sa.Text, nullable=False)
@property
def oauth_params(self) -> dict:
@@ -65,14 +65,14 @@ class BuiltinToolProvider(Base):
__tablename__ = "tool_builtin_providers"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="tool_builtin_provider_pkey"),
db.UniqueConstraint("tenant_id", "provider", "name", name="unique_builtin_tool_provider"),
sa.PrimaryKeyConstraint("id", name="tool_builtin_provider_pkey"),
sa.UniqueConstraint("tenant_id", "provider", "name", name="unique_builtin_tool_provider"),
)
# id of the tool provider
id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
name: Mapped[str] = mapped_column(
String(256), nullable=False, server_default=db.text("'API KEY 1'::character varying")
String(256), nullable=False, server_default=sa.text("'API KEY 1'::character varying")
)
# id of the tenant
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=True)
@@ -81,19 +81,19 @@ class BuiltinToolProvider(Base):
# name of the tool provider
provider: Mapped[str] = mapped_column(String(256), nullable=False)
# credential of the tool provider
encrypted_credentials: Mapped[str] = mapped_column(db.Text, nullable=True)
encrypted_credentials: Mapped[str] = mapped_column(sa.Text, nullable=True)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")
sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")
sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
)
is_default: Mapped[bool] = mapped_column(db.Boolean, nullable=False, server_default=db.text("false"))
is_default: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
# credential type, e.g., "api-key", "oauth2"
credential_type: Mapped[str] = mapped_column(
String(32), nullable=False, server_default=db.text("'api-key'::character varying")
String(32), nullable=False, server_default=sa.text("'api-key'::character varying")
)
expires_at: Mapped[int] = mapped_column(db.BigInteger, nullable=False, server_default=db.text("-1"))
expires_at: Mapped[int] = mapped_column(sa.BigInteger, nullable=False, server_default=sa.text("-1"))
@property
def credentials(self) -> dict:
@@ -107,28 +107,28 @@ class ApiToolProvider(Base):
__tablename__ = "tool_api_providers"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="tool_api_provider_pkey"),
db.UniqueConstraint("name", "tenant_id", name="unique_api_tool_provider"),
sa.PrimaryKeyConstraint("id", name="tool_api_provider_pkey"),
sa.UniqueConstraint("name", "tenant_id", name="unique_api_tool_provider"),
)
id = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
# name of the api provider
name = mapped_column(String(255), nullable=False, server_default=db.text("'API KEY 1'::character varying"))
name = mapped_column(String(255), nullable=False, server_default=sa.text("'API KEY 1'::character varying"))
# icon
icon: Mapped[str] = mapped_column(String(255), nullable=False)
# original schema
schema = mapped_column(db.Text, nullable=False)
schema = mapped_column(sa.Text, nullable=False)
schema_type_str: Mapped[str] = mapped_column(String(40), nullable=False)
# who created this tool
user_id = mapped_column(StringUUID, nullable=False)
# tenant id
tenant_id = mapped_column(StringUUID, nullable=False)
# description of the provider
description = mapped_column(db.Text, nullable=False)
description = mapped_column(sa.Text, nullable=False)
# json format tools
tools_str = mapped_column(db.Text, nullable=False)
tools_str = mapped_column(sa.Text, nullable=False)
# json format credentials
credentials_str = mapped_column(db.Text, nullable=False)
credentials_str = mapped_column(sa.Text, nullable=False)
# privacy policy
privacy_policy = mapped_column(String(255), nullable=True)
# custom_disclaimer
@@ -167,11 +167,11 @@ class ToolLabelBinding(Base):
__tablename__ = "tool_label_bindings"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="tool_label_bind_pkey"),
db.UniqueConstraint("tool_id", "label_name", name="unique_tool_label_bind"),
sa.PrimaryKeyConstraint("id", name="tool_label_bind_pkey"),
sa.UniqueConstraint("tool_id", "label_name", name="unique_tool_label_bind"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
# tool id
tool_id: Mapped[str] = mapped_column(String(64), nullable=False)
# tool type
@@ -187,12 +187,12 @@ class WorkflowToolProvider(Base):
__tablename__ = "tool_workflow_providers"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="tool_workflow_provider_pkey"),
db.UniqueConstraint("name", "tenant_id", name="unique_workflow_tool_provider"),
db.UniqueConstraint("tenant_id", "app_id", name="unique_workflow_tool_provider_app_id"),
sa.PrimaryKeyConstraint("id", name="tool_workflow_provider_pkey"),
sa.UniqueConstraint("name", "tenant_id", name="unique_workflow_tool_provider"),
sa.UniqueConstraint("tenant_id", "app_id", name="unique_workflow_tool_provider_app_id"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
# name of the workflow provider
name: Mapped[str] = mapped_column(String(255), nullable=False)
# label of the workflow provider
@@ -208,17 +208,17 @@ class WorkflowToolProvider(Base):
# tenant id
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# description of the provider
description: Mapped[str] = mapped_column(db.Text, nullable=False)
description: Mapped[str] = mapped_column(sa.Text, nullable=False)
# parameter configuration
parameter_configuration: Mapped[str] = mapped_column(db.Text, nullable=False, server_default="[]")
parameter_configuration: Mapped[str] = mapped_column(sa.Text, nullable=False, server_default="[]")
# privacy policy
privacy_policy: Mapped[str] = mapped_column(String(255), nullable=True, server_default="")
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")
sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")
sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
)
@property
@@ -245,19 +245,19 @@ class MCPToolProvider(Base):
__tablename__ = "tool_mcp_providers"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="tool_mcp_provider_pkey"),
db.UniqueConstraint("tenant_id", "server_url_hash", name="unique_mcp_provider_server_url"),
db.UniqueConstraint("tenant_id", "name", name="unique_mcp_provider_name"),
db.UniqueConstraint("tenant_id", "server_identifier", name="unique_mcp_provider_server_identifier"),
sa.PrimaryKeyConstraint("id", name="tool_mcp_provider_pkey"),
sa.UniqueConstraint("tenant_id", "server_url_hash", name="unique_mcp_provider_server_url"),
sa.UniqueConstraint("tenant_id", "name", name="unique_mcp_provider_name"),
sa.UniqueConstraint("tenant_id", "server_identifier", name="unique_mcp_provider_server_identifier"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
# name of the mcp provider
name: Mapped[str] = mapped_column(String(40), nullable=False)
# server identifier of the mcp provider
server_identifier: Mapped[str] = mapped_column(String(64), nullable=False)
# encrypted url of the mcp provider
server_url: Mapped[str] = mapped_column(db.Text, nullable=False)
server_url: Mapped[str] = mapped_column(sa.Text, nullable=False)
# hash of server_url for uniqueness check
server_url_hash: Mapped[str] = mapped_column(String(64), nullable=False)
# icon of the mcp provider
@@ -267,16 +267,16 @@ class MCPToolProvider(Base):
# who created this tool
user_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# encrypted credentials
encrypted_credentials: Mapped[str] = mapped_column(db.Text, nullable=True)
encrypted_credentials: Mapped[str] = mapped_column(sa.Text, nullable=True)
# authed
authed: Mapped[bool] = mapped_column(db.Boolean, nullable=False, default=False)
authed: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=False)
# tools
tools: Mapped[str] = mapped_column(db.Text, nullable=False, default="[]")
tools: Mapped[str] = mapped_column(sa.Text, nullable=False, default="[]")
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")
sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")
sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
)
def load_user(self) -> Account | None:
@@ -347,9 +347,9 @@ class ToolModelInvoke(Base):
"""
__tablename__ = "tool_model_invokes"
__table_args__ = (db.PrimaryKeyConstraint("id", name="tool_model_invoke_pkey"),)
__table_args__ = (sa.PrimaryKeyConstraint("id", name="tool_model_invoke_pkey"),)
id = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
# who invoke this tool
user_id = mapped_column(StringUUID, nullable=False)
# tenant id
@@ -361,18 +361,18 @@ class ToolModelInvoke(Base):
# tool name
tool_name = mapped_column(String(128), nullable=False)
# invoke parameters
model_parameters = mapped_column(db.Text, nullable=False)
model_parameters = mapped_column(sa.Text, nullable=False)
# prompt messages
prompt_messages = mapped_column(db.Text, nullable=False)
prompt_messages = mapped_column(sa.Text, nullable=False)
# invoke response
model_response = mapped_column(db.Text, nullable=False)
model_response = mapped_column(sa.Text, nullable=False)
prompt_tokens: Mapped[int] = mapped_column(db.Integer, nullable=False, server_default=db.text("0"))
answer_tokens: Mapped[int] = mapped_column(db.Integer, nullable=False, server_default=db.text("0"))
answer_unit_price = mapped_column(db.Numeric(10, 4), nullable=False)
answer_price_unit = mapped_column(db.Numeric(10, 7), nullable=False, server_default=db.text("0.001"))
provider_response_latency = mapped_column(db.Float, nullable=False, server_default=db.text("0"))
total_price = mapped_column(db.Numeric(10, 7))
prompt_tokens: Mapped[int] = mapped_column(sa.Integer, nullable=False, server_default=sa.text("0"))
answer_tokens: Mapped[int] = mapped_column(sa.Integer, nullable=False, server_default=sa.text("0"))
answer_unit_price = mapped_column(sa.Numeric(10, 4), nullable=False)
answer_price_unit = mapped_column(sa.Numeric(10, 7), nullable=False, server_default=sa.text("0.001"))
provider_response_latency = mapped_column(sa.Float, nullable=False, server_default=sa.text("0"))
total_price = mapped_column(sa.Numeric(10, 7))
currency: Mapped[str] = mapped_column(String(255), nullable=False)
created_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
updated_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
@@ -386,13 +386,13 @@ class ToolConversationVariables(Base):
__tablename__ = "tool_conversation_variables"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="tool_conversation_variables_pkey"),
sa.PrimaryKeyConstraint("id", name="tool_conversation_variables_pkey"),
# add index for user_id and conversation_id
db.Index("user_id_idx", "user_id"),
db.Index("conversation_id_idx", "conversation_id"),
sa.Index("user_id_idx", "user_id"),
sa.Index("conversation_id_idx", "conversation_id"),
)
id = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
# conversation user id
user_id = mapped_column(StringUUID, nullable=False)
# tenant id
@@ -400,7 +400,7 @@ class ToolConversationVariables(Base):
# conversation id
conversation_id = mapped_column(StringUUID, nullable=False)
# variables pool
variables_str = mapped_column(db.Text, nullable=False)
variables_str = mapped_column(sa.Text, nullable=False)
created_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
updated_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
@@ -417,11 +417,11 @@ class ToolFile(Base):
__tablename__ = "tool_files"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="tool_file_pkey"),
db.Index("tool_file_conversation_id_idx", "conversation_id"),
sa.PrimaryKeyConstraint("id", name="tool_file_pkey"),
sa.Index("tool_file_conversation_id_idx", "conversation_id"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
# conversation user id
user_id: Mapped[str] = mapped_column(StringUUID)
# tenant id
@@ -448,30 +448,30 @@ class DeprecatedPublishedAppTool(Base):
__tablename__ = "tool_published_apps"
__table_args__ = (
db.PrimaryKeyConstraint("id", name="published_app_tool_pkey"),
db.UniqueConstraint("app_id", "user_id", name="unique_published_app_tool"),
sa.PrimaryKeyConstraint("id", name="published_app_tool_pkey"),
sa.UniqueConstraint("app_id", "user_id", name="unique_published_app_tool"),
)
id = mapped_column(StringUUID, server_default=db.text("uuid_generate_v4()"))
id = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
# id of the app
app_id = mapped_column(StringUUID, ForeignKey("apps.id"), nullable=False)
user_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# who published this tool
description = mapped_column(db.Text, nullable=False)
description = mapped_column(sa.Text, nullable=False)
# llm_description of the tool, for LLM
llm_description = mapped_column(db.Text, nullable=False)
llm_description = mapped_column(sa.Text, nullable=False)
# query description, query will be seem as a parameter of the tool,
# to describe this parameter to llm, we need this field
query_description = mapped_column(db.Text, nullable=False)
query_description = mapped_column(sa.Text, nullable=False)
# query name, the name of the query parameter
query_name = mapped_column(String(40), nullable=False)
# name of the tool provider
tool_name = mapped_column(String(40), nullable=False)
# author
author = mapped_column(String(40), nullable=False)
created_at = mapped_column(sa.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
updated_at = mapped_column(sa.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
created_at = mapped_column(sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)"))
updated_at = mapped_column(sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)"))
@property
def description_i18n(self) -> I18nObject: