Underground-Digital/Workflow-Engine

1import json2from typing import Optional3 4import sqlalchemy as sa5from sqlalchemy import ForeignKey6from sqlalchemy.orm import Mapped, mapped_column7 8from core.tools.entities.common_entities import I18nObject9from core.tools.entities.tool_bundle import ApiToolBundle10from core.tools.entities.tool_entities import ApiProviderSchemaType, WorkflowToolParameterConfiguration11from extensions.ext_database import db12 13from .model import Account, App, Tenant14from .types import StringUUID15 16 17class BuiltinToolProvider(db.Model):18    """19    This table stores the tool provider information for built-in tools for each tenant.20    """21 22    __tablename__ = "tool_builtin_providers"23    __table_args__ = (24        db.PrimaryKeyConstraint("id", name="tool_builtin_provider_pkey"),25        # one tenant can only have one tool provider with the same name26        db.UniqueConstraint("tenant_id", "provider", name="unique_builtin_tool_provider"),27    )28 29    # id of the tool provider30    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))31    # id of the tenant32    tenant_id = db.Column(StringUUID, nullable=True)33    # who created this tool provider34    user_id = db.Column(StringUUID, nullable=False)35    # name of the tool provider36    provider = db.Column(db.String(40), nullable=False)37    # credential of the tool provider38    encrypted_credentials = db.Column(db.Text, nullable=True)39    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))40    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))41 42    @property43    def credentials(self) -> dict:44        return json.loads(self.encrypted_credentials)45 46 47class PublishedAppTool(db.Model):48    """49    The table stores the apps published as a tool for each person.50    """51 52    __tablename__ = "tool_published_apps"53    __table_args__ = (54        db.PrimaryKeyConstraint("id", name="published_app_tool_pkey"),55        db.UniqueConstraint("app_id", "user_id", name="unique_published_app_tool"),56    )57 58    # id of the tool provider59    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))60    # id of the app61    app_id = db.Column(StringUUID, ForeignKey("apps.id"), nullable=False)62    # who published this tool63    user_id = db.Column(StringUUID, nullable=False)64    # description of the tool, stored in i18n format, for human65    description = db.Column(db.Text, nullable=False)66    # llm_description of the tool, for LLM67    llm_description = db.Column(db.Text, nullable=False)68    # query description, query will be seem as a parameter of the tool,69    # to describe this parameter to llm, we need this field70    query_description = db.Column(db.Text, nullable=False)71    # query name, the name of the query parameter72    query_name = db.Column(db.String(40), nullable=False)73    # name of the tool provider74    tool_name = db.Column(db.String(40), nullable=False)75    # author76    author = db.Column(db.String(40), nullable=False)77    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))78    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))79 80    @property81    def description_i18n(self) -> I18nObject:82        return I18nObject(**json.loads(self.description))83 84    @property85    def app(self) -> App:86        return db.session.query(App).filter(App.id == self.app_id).first()87 88 89class ApiToolProvider(db.Model):90    """91    The table stores the api providers.92    """93 94    __tablename__ = "tool_api_providers"95    __table_args__ = (96        db.PrimaryKeyConstraint("id", name="tool_api_provider_pkey"),97        db.UniqueConstraint("name", "tenant_id", name="unique_api_tool_provider"),98    )99 100    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))101    # name of the api provider102    name = db.Column(db.String(40), nullable=False)103    # icon104    icon = db.Column(db.String(255), nullable=False)105    # original schema106    schema = db.Column(db.Text, nullable=False)107    schema_type_str: Mapped[str] = db.Column(db.String(40), nullable=False)108    # who created this tool109    user_id = db.Column(StringUUID, nullable=False)110    # tenant id111    tenant_id = db.Column(StringUUID, nullable=False)112    # description of the provider113    description = db.Column(db.Text, nullable=False)114    # json format tools115    tools_str = db.Column(db.Text, nullable=False)116    # json format credentials117    credentials_str = db.Column(db.Text, nullable=False)118    # privacy policy119    privacy_policy = db.Column(db.String(255), nullable=True)120    # custom_disclaimer121    custom_disclaimer: Mapped[str] = mapped_column(sa.TEXT, default="")122 123    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))124    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))125 126    @property127    def schema_type(self) -> ApiProviderSchemaType:128        return ApiProviderSchemaType.value_of(self.schema_type_str)129 130    @property131    def tools(self) -> list[ApiToolBundle]:132        return [ApiToolBundle(**tool) for tool in json.loads(self.tools_str)]133 134    @property135    def credentials(self) -> dict:136        return json.loads(self.credentials_str)137 138    @property139    def user(self) -> Account | None:140        return db.session.query(Account).filter(Account.id == self.user_id).first()141 142    @property143    def tenant(self) -> Tenant | None:144        return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first()145 146 147class ToolLabelBinding(db.Model):148    """149    The table stores the labels for tools.150    """151 152    __tablename__ = "tool_label_bindings"153    __table_args__ = (154        db.PrimaryKeyConstraint("id", name="tool_label_bind_pkey"),155        db.UniqueConstraint("tool_id", "label_name", name="unique_tool_label_bind"),156    )157 158    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))159    # tool id160    tool_id = db.Column(db.String(64), nullable=False)161    # tool type162    tool_type = db.Column(db.String(40), nullable=False)163    # label name164    label_name = db.Column(db.String(40), nullable=False)165 166 167class WorkflowToolProvider(db.Model):168    """169    The table stores the workflow providers.170    """171 172    __tablename__ = "tool_workflow_providers"173    __table_args__ = (174        db.PrimaryKeyConstraint("id", name="tool_workflow_provider_pkey"),175        db.UniqueConstraint("name", "tenant_id", name="unique_workflow_tool_provider"),176        db.UniqueConstraint("tenant_id", "app_id", name="unique_workflow_tool_provider_app_id"),177    )178 179    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))180    # name of the workflow provider181    name = db.Column(db.String(40), nullable=False)182    # label of the workflow provider183    label = db.Column(db.String(255), nullable=False, server_default="")184    # icon185    icon = db.Column(db.String(255), nullable=False)186    # app id of the workflow provider187    app_id = db.Column(StringUUID, nullable=False)188    # version of the workflow provider189    version = db.Column(db.String(255), nullable=False, server_default="")190    # who created this tool191    user_id = db.Column(StringUUID, nullable=False)192    # tenant id193    tenant_id = db.Column(StringUUID, nullable=False)194    # description of the provider195    description = db.Column(db.Text, nullable=False)196    # parameter configuration197    parameter_configuration = db.Column(db.Text, nullable=False, server_default="[]")198    # privacy policy199    privacy_policy = db.Column(db.String(255), nullable=True, server_default="")200 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 204    @property205    def schema_type(self) -> ApiProviderSchemaType:206        return ApiProviderSchemaType.value_of(self.schema_type_str)207 208    @property209    def user(self) -> Account | None:210        return db.session.query(Account).filter(Account.id == self.user_id).first()211 212    @property213    def tenant(self) -> Tenant | None:214        return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first()215 216    @property217    def parameter_configurations(self) -> list[WorkflowToolParameterConfiguration]:218        return [WorkflowToolParameterConfiguration(**config) for config in json.loads(self.parameter_configuration)]219 220    @property221    def app(self) -> App | None:222        return db.session.query(App).filter(App.id == self.app_id).first()223 224 225class ToolModelInvoke(db.Model):226    """227    store the invoke logs from tool invoke228    """229 230    __tablename__ = "tool_model_invokes"231    __table_args__ = (db.PrimaryKeyConstraint("id", name="tool_model_invoke_pkey"),)232 233    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))234    # who invoke this tool235    user_id = db.Column(StringUUID, nullable=False)236    # tenant id237    tenant_id = db.Column(StringUUID, nullable=False)238    # provider239    provider = db.Column(db.String(40), nullable=False)240    # type241    tool_type = db.Column(db.String(40), nullable=False)242    # tool name243    tool_name = db.Column(db.String(40), nullable=False)244    # invoke parameters245    model_parameters = db.Column(db.Text, nullable=False)246    # prompt messages247    prompt_messages = db.Column(db.Text, nullable=False)248    # invoke response249    model_response = db.Column(db.Text, nullable=False)250 251    prompt_tokens = db.Column(db.Integer, nullable=False, server_default=db.text("0"))252    answer_tokens = db.Column(db.Integer, nullable=False, server_default=db.text("0"))253    answer_unit_price = db.Column(db.Numeric(10, 4), nullable=False)254    answer_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text("0.001"))255    provider_response_latency = db.Column(db.Float, nullable=False, server_default=db.text("0"))256    total_price = db.Column(db.Numeric(10, 7))257    currency = db.Column(db.String(255), nullable=False)258    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))259    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))260 261 262class ToolConversationVariables(db.Model):263    """264    store the conversation variables from tool invoke265    """266 267    __tablename__ = "tool_conversation_variables"268    __table_args__ = (269        db.PrimaryKeyConstraint("id", name="tool_conversation_variables_pkey"),270        # add index for user_id and conversation_id271        db.Index("user_id_idx", "user_id"),272        db.Index("conversation_id_idx", "conversation_id"),273    )274 275    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))276    # conversation user id277    user_id = db.Column(StringUUID, nullable=False)278    # tenant id279    tenant_id = db.Column(StringUUID, nullable=False)280    # conversation id281    conversation_id = db.Column(StringUUID, nullable=False)282    # variables pool283    variables_str = db.Column(db.Text, nullable=False)284 285    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))286    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))287 288    @property289    def variables(self) -> dict:290        return json.loads(self.variables_str)291 292 293class ToolFile(db.Model):294    __tablename__ = "tool_files"295    __table_args__ = (296        db.PrimaryKeyConstraint("id", name="tool_file_pkey"),297        db.Index("tool_file_conversation_id_idx", "conversation_id"),298    )299 300    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))301    user_id: Mapped[str] = db.Column(StringUUID, nullable=False)302    tenant_id: Mapped[str] = db.Column(StringUUID, nullable=False)303    conversation_id: Mapped[Optional[str]] = db.Column(StringUUID, nullable=True)304    file_key: Mapped[str] = db.Column(db.String(255), nullable=False)305    mimetype: Mapped[str] = db.Column(db.String(255), nullable=False)306    original_url: Mapped[Optional[str]] = db.Column(db.String(2048), nullable=True)307    name: Mapped[str] = mapped_column(default="")308    size: Mapped[int] = mapped_column(default=-1)309 310    def __init__(311        self,312        *,313        user_id: str,314        tenant_id: str,315        conversation_id: Optional[str] = None,316        file_key: str,317        mimetype: str,318        original_url: Optional[str] = None,319        name: str,320        size: int,321    ):322        self.user_id = user_id323        self.tenant_id = tenant_id324        self.conversation_id = conversation_id325        self.file_key = file_key326        self.mimetype = mimetype327        self.original_url = original_url328        self.name = name329        self.size = size330