dify/api/models/tools.py
Novice 0ded6303c1
feat: implement MCP specification 2025-06-18 (#25766)
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2025-10-27 17:07:51 +08:00

526 lines
21 KiB
Python

import json
from datetime import datetime
from decimal import Decimal
from typing import TYPE_CHECKING, Any, cast
import sqlalchemy as sa
from deprecated import deprecated
from sqlalchemy import ForeignKey, String, func
from sqlalchemy.orm import Mapped, mapped_column
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.tool_bundle import ApiToolBundle
from core.tools.entities.tool_entities import ApiProviderSchemaType, WorkflowToolParameterConfiguration
from models.base import TypeBase
from .engine import db
from .model import Account, App, Tenant
from .types import StringUUID
if TYPE_CHECKING:
from core.entities.mcp_provider import MCPProviderEntity
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.tool_bundle import ApiToolBundle
from core.tools.entities.tool_entities import ApiProviderSchemaType, WorkflowToolParameterConfiguration
# system level tool oauth client params (client_id, client_secret, etc.)
class ToolOAuthSystemClient(TypeBase):
__tablename__ = "tool_oauth_system_clients"
__table_args__ = (
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=sa.text("uuid_generate_v4()"), init=False)
plugin_id: Mapped[str] = 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(sa.Text, nullable=False)
# tenant level tool oauth client params (client_id, client_secret, etc.)
class ToolOAuthTenantClient(TypeBase):
__tablename__ = "tool_oauth_tenant_clients"
__table_args__ = (
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=sa.text("uuid_generate_v4()"), init=False)
# 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(sa.Boolean, nullable=False, server_default=sa.text("true"), init=False)
# oauth params of the tool provider
encrypted_oauth_params: Mapped[str] = mapped_column(sa.Text, nullable=False, init=False)
@property
def oauth_params(self) -> dict[str, Any]:
return cast(dict[str, Any], json.loads(self.encrypted_oauth_params or "{}"))
class BuiltinToolProvider(TypeBase):
"""
This table stores the tool provider information for built-in tools for each tenant.
"""
__tablename__ = "tool_builtin_providers"
__table_args__ = (
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=sa.text("uuid_generate_v4()"), init=False)
name: Mapped[str] = mapped_column(
String(256),
nullable=False,
server_default=sa.text("'API KEY 1'::character varying"),
)
# id of the tenant
tenant_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
# who created this tool provider
user_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# name of the tool provider
provider: Mapped[str] = mapped_column(String(256), nullable=False)
# credential of the tool provider
encrypted_credentials: Mapped[str | None] = mapped_column(sa.Text, nullable=True, default=None)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)"), init=False
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime,
nullable=False,
server_default=sa.text("CURRENT_TIMESTAMP(0)"),
onupdate=func.current_timestamp(),
init=False,
)
is_default: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"), default=False)
# credential type, e.g., "api-key", "oauth2"
credential_type: Mapped[str] = mapped_column(
String(32), nullable=False, server_default=sa.text("'api-key'::character varying"), default="api-key"
)
expires_at: Mapped[int] = mapped_column(sa.BigInteger, nullable=False, server_default=sa.text("-1"), default=-1)
@property
def credentials(self) -> dict[str, Any]:
if not self.encrypted_credentials:
return {}
return cast(dict[str, Any], json.loads(self.encrypted_credentials))
class ApiToolProvider(TypeBase):
"""
The table stores the api providers.
"""
__tablename__ = "tool_api_providers"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="tool_api_provider_pkey"),
sa.UniqueConstraint("name", "tenant_id", name="unique_api_tool_provider"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"), init=False)
# name of the api provider
name: Mapped[str] = 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[str] = mapped_column(sa.Text, nullable=False)
schema_type_str: Mapped[str] = mapped_column(String(40), nullable=False)
# who created this tool
user_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# tenant id
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# description of the provider
description: Mapped[str] = mapped_column(sa.Text, nullable=False)
# json format tools
tools_str: Mapped[str] = mapped_column(sa.Text, nullable=False)
# json format credentials
credentials_str: Mapped[str] = mapped_column(sa.Text, nullable=False)
# privacy policy
privacy_policy: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
# custom_disclaimer
custom_disclaimer: Mapped[str] = mapped_column(sa.TEXT, default="")
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime,
nullable=False,
server_default=func.current_timestamp(),
onupdate=func.current_timestamp(),
init=False,
)
@property
def schema_type(self) -> "ApiProviderSchemaType":
from core.tools.entities.tool_entities import ApiProviderSchemaType
return ApiProviderSchemaType.value_of(self.schema_type_str)
@property
def tools(self) -> list["ApiToolBundle"]:
from core.tools.entities.tool_bundle import ApiToolBundle
return [ApiToolBundle.model_validate(tool) for tool in json.loads(self.tools_str)]
@property
def credentials(self) -> dict[str, Any]:
return dict[str, Any](json.loads(self.credentials_str))
@property
def user(self) -> Account | None:
if not self.user_id:
return None
return db.session.query(Account).where(Account.id == self.user_id).first()
@property
def tenant(self) -> Tenant | None:
return db.session.query(Tenant).where(Tenant.id == self.tenant_id).first()
class ToolLabelBinding(TypeBase):
"""
The table stores the labels for tools.
"""
__tablename__ = "tool_label_bindings"
__table_args__ = (
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=sa.text("uuid_generate_v4()"), init=False)
# tool id
tool_id: Mapped[str] = mapped_column(String(64), nullable=False)
# tool type
tool_type: Mapped[str] = mapped_column(String(40), nullable=False)
# label name
label_name: Mapped[str] = mapped_column(String(40), nullable=False)
class WorkflowToolProvider(TypeBase):
"""
The table stores the workflow providers.
"""
__tablename__ = "tool_workflow_providers"
__table_args__ = (
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=sa.text("uuid_generate_v4()"))
# name of the workflow provider
name: Mapped[str] = mapped_column(String(255), nullable=False)
# label of the workflow provider
label: Mapped[str] = mapped_column(String(255), nullable=False, server_default="")
# icon
icon: Mapped[str] = mapped_column(String(255), nullable=False)
# app id of the workflow provider
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# version of the workflow provider
version: Mapped[str] = mapped_column(String(255), nullable=False, server_default="")
# who created this tool
user_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# tenant id
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# description of the provider
description: Mapped[str] = mapped_column(sa.Text, nullable=False)
# parameter configuration
parameter_configuration: Mapped[str] = mapped_column(sa.Text, nullable=False, server_default="[]", default="[]")
# privacy policy
privacy_policy: Mapped[str | None] = mapped_column(String(255), nullable=True, server_default="", default=None)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)"), init=False
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime,
nullable=False,
server_default=sa.text("CURRENT_TIMESTAMP(0)"),
onupdate=func.current_timestamp(),
init=False,
)
@property
def user(self) -> Account | None:
return db.session.query(Account).where(Account.id == self.user_id).first()
@property
def tenant(self) -> Tenant | None:
return db.session.query(Tenant).where(Tenant.id == self.tenant_id).first()
@property
def parameter_configurations(self) -> list["WorkflowToolParameterConfiguration"]:
from core.tools.entities.tool_entities import WorkflowToolParameterConfiguration
return [
WorkflowToolParameterConfiguration.model_validate(config)
for config in json.loads(self.parameter_configuration)
]
@property
def app(self) -> App | None:
return db.session.query(App).where(App.id == self.app_id).first()
class MCPToolProvider(TypeBase):
"""
The table stores the mcp providers.
"""
__tablename__ = "tool_mcp_providers"
__table_args__ = (
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=sa.text("uuid_generate_v4()"), init=False)
# 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(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
icon: Mapped[str | None] = mapped_column(String(255), nullable=True)
# tenant id
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# who created this tool
user_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# encrypted credentials
encrypted_credentials: Mapped[str | None] = mapped_column(sa.Text, nullable=True, default=None)
# authed
authed: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=False)
# tools
tools: Mapped[str] = mapped_column(sa.Text, nullable=False, default="[]")
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)"), init=False
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime,
nullable=False,
server_default=sa.text("CURRENT_TIMESTAMP(0)"),
onupdate=func.current_timestamp(),
init=False,
)
timeout: Mapped[float] = mapped_column(sa.Float, nullable=False, server_default=sa.text("30"), default=30.0)
sse_read_timeout: Mapped[float] = mapped_column(
sa.Float, nullable=False, server_default=sa.text("300"), default=300.0
)
# encrypted headers for MCP server requests
encrypted_headers: Mapped[str | None] = mapped_column(sa.Text, nullable=True, default=None)
def load_user(self) -> Account | None:
return db.session.query(Account).where(Account.id == self.user_id).first()
@property
def credentials(self) -> dict[str, Any]:
if not self.encrypted_credentials:
return {}
try:
return json.loads(self.encrypted_credentials)
except Exception:
return {}
@property
def headers(self) -> dict[str, Any]:
if self.encrypted_headers is None:
return {}
try:
return json.loads(self.encrypted_headers)
except Exception:
return {}
@property
def tool_dict(self) -> list[dict[str, Any]]:
try:
return json.loads(self.tools) if self.tools else []
except (json.JSONDecodeError, TypeError):
return []
def to_entity(self) -> "MCPProviderEntity":
"""Convert to domain entity"""
from core.entities.mcp_provider import MCPProviderEntity
return MCPProviderEntity.from_db_model(self)
class ToolModelInvoke(TypeBase):
"""
store the invoke logs from tool invoke
"""
__tablename__ = "tool_model_invokes"
__table_args__ = (sa.PrimaryKeyConstraint("id", name="tool_model_invoke_pkey"),)
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"), init=False)
# who invoke this tool
user_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# tenant id
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# provider
provider: Mapped[str] = mapped_column(String(255), nullable=False)
# type
tool_type: Mapped[str] = mapped_column(String(40), nullable=False)
# tool name
tool_name: Mapped[str] = mapped_column(String(128), nullable=False)
# invoke parameters
model_parameters: Mapped[str] = mapped_column(sa.Text, nullable=False)
# prompt messages
prompt_messages: Mapped[str] = mapped_column(sa.Text, nullable=False)
# invoke response
model_response: Mapped[str] = mapped_column(sa.Text, nullable=False)
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[Decimal] = mapped_column(sa.Numeric(10, 4), nullable=False)
answer_price_unit: Mapped[Decimal] = mapped_column(
sa.Numeric(10, 7), nullable=False, server_default=sa.text("0.001")
)
provider_response_latency: Mapped[float] = mapped_column(sa.Float, nullable=False, server_default=sa.text("0"))
total_price: Mapped[Decimal | None] = mapped_column(sa.Numeric(10, 7))
currency: Mapped[str] = mapped_column(String(255), nullable=False)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime,
nullable=False,
server_default=func.current_timestamp(),
onupdate=func.current_timestamp(),
init=False,
)
@deprecated
class ToolConversationVariables(TypeBase):
"""
store the conversation variables from tool invoke
"""
__tablename__ = "tool_conversation_variables"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="tool_conversation_variables_pkey"),
# add index for user_id and conversation_id
sa.Index("user_id_idx", "user_id"),
sa.Index("conversation_id_idx", "conversation_id"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"), init=False)
# conversation user id
user_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# tenant id
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# conversation id
conversation_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# variables pool
variables_str: Mapped[str] = mapped_column(sa.Text, nullable=False)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime,
nullable=False,
server_default=func.current_timestamp(),
onupdate=func.current_timestamp(),
init=False,
)
@property
def variables(self):
return json.loads(self.variables_str)
class ToolFile(TypeBase):
"""This table stores file metadata generated in workflows,
not only files created by agent.
"""
__tablename__ = "tool_files"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="tool_file_pkey"),
sa.Index("tool_file_conversation_id_idx", "conversation_id"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"), init=False)
# conversation user id
user_id: Mapped[str] = mapped_column(StringUUID)
# tenant id
tenant_id: Mapped[str] = mapped_column(StringUUID)
# conversation id
conversation_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
# file key
file_key: Mapped[str] = mapped_column(String(255), nullable=False)
# mime type
mimetype: Mapped[str] = mapped_column(String(255), nullable=False)
# original url
original_url: Mapped[str | None] = mapped_column(String(2048), nullable=True, default=None)
# name
name: Mapped[str] = mapped_column(default="")
# size
size: Mapped[int] = mapped_column(default=-1)
@deprecated
class DeprecatedPublishedAppTool(TypeBase):
"""
The table stores the apps published as a tool for each person.
"""
__tablename__ = "tool_published_apps"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="published_app_tool_pkey"),
sa.UniqueConstraint("app_id", "user_id", name="unique_published_app_tool"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"), init=False)
# id of the app
app_id: Mapped[str] = mapped_column(StringUUID, ForeignKey("apps.id"), nullable=False)
user_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
# who published this tool
description: Mapped[str] = mapped_column(sa.Text, nullable=False)
# llm_description of the tool, for LLM
llm_description: Mapped[str] = 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[str] = mapped_column(sa.Text, nullable=False)
# query name, the name of the query parameter
query_name: Mapped[str] = mapped_column(String(40), nullable=False)
# name of the tool provider
tool_name: Mapped[str] = mapped_column(String(40), nullable=False)
# author
author: Mapped[str] = mapped_column(String(40), nullable=False)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)"), init=False
)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime,
nullable=False,
server_default=sa.text("CURRENT_TIMESTAMP(0)"),
onupdate=func.current_timestamp(),
init=False,
)
@property
def description_i18n(self) -> "I18nObject":
from core.tools.entities.common_entities import I18nObject
return I18nObject.model_validate(json.loads(self.description))