From d109881410e6630c582f4eeb812e53b144e6189b Mon Sep 17 00:00:00 2001 From: Bowen Liang Date: Tue, 10 Sep 2024 17:08:06 +0800 Subject: [PATCH] chore(api/models): apply ruff reformatting (#7600) Co-authored-by: -LAN- --- api/models/__init__.py | 10 +- api/models/account.py | 169 +++--- api/models/api_based_extension.py | 18 +- api/models/dataset.py | 581 ++++++++++---------- api/models/model.py | 885 ++++++++++++++++-------------- api/models/provider.py | 122 ++-- api/models/source.py | 48 +- api/models/task.py | 21 +- api/models/tool.py | 20 +- api/models/tools.py | 144 ++--- api/models/types.py | 6 +- api/models/web.py | 21 +- api/models/workflow.py | 357 ++++++------ api/pyproject.toml | 1 - 14 files changed, 1265 insertions(+), 1138 deletions(-) diff --git a/api/models/__init__.py b/api/models/__init__.py index 4012611471..30ceef057e 100644 --- a/api/models/__init__.py +++ b/api/models/__init__.py @@ -4,7 +4,7 @@ from .model import App, AppMode, Message from .types import StringUUID from .workflow import ConversationVariable, Workflow, WorkflowNodeExecutionStatus -__all__ = ['ConversationVariable', 'StringUUID', 'AppMode', 'WorkflowNodeExecutionStatus', 'Workflow', 'App', 'Message'] +__all__ = ["ConversationVariable", "StringUUID", "AppMode", "WorkflowNodeExecutionStatus", "Workflow", "App", "Message"] class CreatedByRole(Enum): @@ -12,11 +12,11 @@ class CreatedByRole(Enum): Enum class for createdByRole """ - ACCOUNT = 'account' - END_USER = 'end_user' + ACCOUNT = "account" + END_USER = "end_user" @classmethod - def value_of(cls, value: str) -> 'CreatedByRole': + def value_of(cls, value: str) -> "CreatedByRole": """ Get value of given mode. @@ -26,4 +26,4 @@ class CreatedByRole(Enum): for role in cls: if role.value == value: return role - raise ValueError(f'invalid createdByRole value {value}') + raise ValueError(f"invalid createdByRole value {value}") diff --git a/api/models/account.py b/api/models/account.py index 67d940b7b7..60b4f11aad 100644 --- a/api/models/account.py +++ b/api/models/account.py @@ -9,21 +9,18 @@ from .types import StringUUID class AccountStatus(str, enum.Enum): - PENDING = 'pending' - UNINITIALIZED = 'uninitialized' - ACTIVE = 'active' - BANNED = 'banned' - CLOSED = 'closed' + PENDING = "pending" + UNINITIALIZED = "uninitialized" + ACTIVE = "active" + BANNED = "banned" + CLOSED = "closed" class Account(UserMixin, db.Model): - __tablename__ = 'accounts' - __table_args__ = ( - db.PrimaryKeyConstraint('id', name='account_pkey'), - db.Index('account_email_idx', 'email') - ) + __tablename__ = "accounts" + __table_args__ = (db.PrimaryKeyConstraint("id", name="account_pkey"), db.Index("account_email_idx", "email")) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) name = db.Column(db.String(255), nullable=False) email = db.Column(db.String(255), nullable=False) password = db.Column(db.String(255), nullable=True) @@ -34,11 +31,11 @@ class Account(UserMixin, db.Model): timezone = db.Column(db.String(255)) last_login_at = db.Column(db.DateTime) last_login_ip = db.Column(db.String(255)) - last_active_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + last_active_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) status = db.Column(db.String(16), nullable=False, server_default=db.text("'active'::character varying")) initialized_at = db.Column(db.DateTime) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def is_password_set(self): @@ -65,11 +62,13 @@ class Account(UserMixin, db.Model): @current_tenant_id.setter def current_tenant_id(self, value: str): try: - tenant_account_join = db.session.query(Tenant, TenantAccountJoin) \ - .filter(Tenant.id == value) \ - .filter(TenantAccountJoin.tenant_id == Tenant.id) \ - .filter(TenantAccountJoin.account_id == self.id) \ + tenant_account_join = ( + db.session.query(Tenant, TenantAccountJoin) + .filter(Tenant.id == value) + .filter(TenantAccountJoin.tenant_id == Tenant.id) + .filter(TenantAccountJoin.account_id == self.id) .one_or_none() + ) if tenant_account_join: tenant, ta = tenant_account_join @@ -91,20 +90,18 @@ class Account(UserMixin, db.Model): @classmethod def get_by_openid(cls, provider: str, open_id: str) -> db.Model: - account_integrate = db.session.query(AccountIntegrate). \ - filter(AccountIntegrate.provider == provider, AccountIntegrate.open_id == open_id). \ - one_or_none() + account_integrate = ( + db.session.query(AccountIntegrate) + .filter(AccountIntegrate.provider == provider, AccountIntegrate.open_id == open_id) + .one_or_none() + ) if account_integrate: - return db.session.query(Account). \ - filter(Account.id == account_integrate.account_id). \ - one_or_none() + return db.session.query(Account).filter(Account.id == account_integrate.account_id).one_or_none() return None def get_integrates(self) -> list[db.Model]: ai = db.Model - return db.session.query(ai).filter( - ai.account_id == self.id - ).all() + return db.session.query(ai).filter(ai.account_id == self.id).all() # check current_user.current_tenant.current_role in ['admin', 'owner'] @property @@ -123,61 +120,75 @@ class Account(UserMixin, db.Model): def is_dataset_operator(self): return self._current_tenant.current_role == TenantAccountRole.DATASET_OPERATOR + class TenantStatus(str, enum.Enum): - NORMAL = 'normal' - ARCHIVE = 'archive' + NORMAL = "normal" + ARCHIVE = "archive" class TenantAccountRole(str, enum.Enum): - OWNER = 'owner' - ADMIN = 'admin' - EDITOR = 'editor' - NORMAL = 'normal' - DATASET_OPERATOR = 'dataset_operator' + OWNER = "owner" + ADMIN = "admin" + EDITOR = "editor" + NORMAL = "normal" + DATASET_OPERATOR = "dataset_operator" @staticmethod def is_valid_role(role: str) -> bool: - return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, - TenantAccountRole.NORMAL, TenantAccountRole.DATASET_OPERATOR} + return role and role in { + TenantAccountRole.OWNER, + TenantAccountRole.ADMIN, + TenantAccountRole.EDITOR, + TenantAccountRole.NORMAL, + TenantAccountRole.DATASET_OPERATOR, + } @staticmethod def is_privileged_role(role: str) -> bool: return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN} - + @staticmethod def is_non_owner_role(role: str) -> bool: - return role and role in {TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, TenantAccountRole.NORMAL, - TenantAccountRole.DATASET_OPERATOR} - + return role and role in { + TenantAccountRole.ADMIN, + TenantAccountRole.EDITOR, + TenantAccountRole.NORMAL, + TenantAccountRole.DATASET_OPERATOR, + } + @staticmethod def is_editing_role(role: str) -> bool: return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR} @staticmethod def is_dataset_edit_role(role: str) -> bool: - return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, - TenantAccountRole.DATASET_OPERATOR} + return role and role in { + TenantAccountRole.OWNER, + TenantAccountRole.ADMIN, + TenantAccountRole.EDITOR, + TenantAccountRole.DATASET_OPERATOR, + } + class Tenant(db.Model): - __tablename__ = 'tenants' - __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tenant_pkey'), - ) + __tablename__ = "tenants" + __table_args__ = (db.PrimaryKeyConstraint("id", name="tenant_pkey"),) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) name = db.Column(db.String(255), nullable=False) encrypt_public_key = db.Column(db.Text) plan = db.Column(db.String(255), nullable=False, server_default=db.text("'basic'::character varying")) status = db.Column(db.String(255), nullable=False, server_default=db.text("'normal'::character varying")) custom_config = db.Column(db.Text) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) def get_accounts(self) -> list[Account]: - return db.session.query(Account).filter( - Account.id == TenantAccountJoin.account_id, - TenantAccountJoin.tenant_id == self.id - ).all() + return ( + db.session.query(Account) + .filter(Account.id == TenantAccountJoin.account_id, TenantAccountJoin.tenant_id == self.id) + .all() + ) @property def custom_config_dict(self) -> dict: @@ -189,54 +200,54 @@ class Tenant(db.Model): class TenantAccountJoinRole(enum.Enum): - OWNER = 'owner' - ADMIN = 'admin' - NORMAL = 'normal' - DATASET_OPERATOR = 'dataset_operator' + OWNER = "owner" + ADMIN = "admin" + NORMAL = "normal" + DATASET_OPERATOR = "dataset_operator" class TenantAccountJoin(db.Model): - __tablename__ = 'tenant_account_joins' + __tablename__ = "tenant_account_joins" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tenant_account_join_pkey'), - db.Index('tenant_account_join_account_id_idx', 'account_id'), - db.Index('tenant_account_join_tenant_id_idx', 'tenant_id'), - db.UniqueConstraint('tenant_id', 'account_id', name='unique_tenant_account_join') + db.PrimaryKeyConstraint("id", name="tenant_account_join_pkey"), + db.Index("tenant_account_join_account_id_idx", "account_id"), + db.Index("tenant_account_join_tenant_id_idx", "tenant_id"), + db.UniqueConstraint("tenant_id", "account_id", name="unique_tenant_account_join"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) account_id = db.Column(StringUUID, nullable=False) - current = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) - role = db.Column(db.String(16), nullable=False, server_default='normal') + current = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) + role = db.Column(db.String(16), nullable=False, server_default="normal") invited_by = db.Column(StringUUID, nullable=True) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class AccountIntegrate(db.Model): - __tablename__ = 'account_integrates' + __tablename__ = "account_integrates" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='account_integrate_pkey'), - db.UniqueConstraint('account_id', 'provider', name='unique_account_provider'), - db.UniqueConstraint('provider', 'open_id', name='unique_provider_open_id') + db.PrimaryKeyConstraint("id", name="account_integrate_pkey"), + db.UniqueConstraint("account_id", "provider", name="unique_account_provider"), + db.UniqueConstraint("provider", "open_id", name="unique_provider_open_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) account_id = db.Column(StringUUID, nullable=False) provider = db.Column(db.String(16), nullable=False) open_id = db.Column(db.String(255), nullable=False) encrypted_token = db.Column(db.String(255), nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class InvitationCode(db.Model): - __tablename__ = 'invitation_codes' + __tablename__ = "invitation_codes" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='invitation_code_pkey'), - db.Index('invitation_codes_batch_idx', 'batch'), - db.Index('invitation_codes_code_idx', 'code', 'status') + db.PrimaryKeyConstraint("id", name="invitation_code_pkey"), + db.Index("invitation_codes_batch_idx", "batch"), + db.Index("invitation_codes_code_idx", "code", "status"), ) id = db.Column(db.Integer, nullable=False) @@ -247,4 +258,4 @@ class InvitationCode(db.Model): used_by_tenant_id = db.Column(StringUUID) used_by_account_id = db.Column(StringUUID) deprecated_at = db.Column(db.DateTime) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) diff --git a/api/models/api_based_extension.py b/api/models/api_based_extension.py index 7f69323628..97173747af 100644 --- a/api/models/api_based_extension.py +++ b/api/models/api_based_extension.py @@ -6,22 +6,22 @@ from .types import StringUUID class APIBasedExtensionPoint(enum.Enum): - APP_EXTERNAL_DATA_TOOL_QUERY = 'app.external_data_tool.query' - PING = 'ping' - APP_MODERATION_INPUT = 'app.moderation.input' - APP_MODERATION_OUTPUT = 'app.moderation.output' + APP_EXTERNAL_DATA_TOOL_QUERY = "app.external_data_tool.query" + PING = "ping" + APP_MODERATION_INPUT = "app.moderation.input" + APP_MODERATION_OUTPUT = "app.moderation.output" class APIBasedExtension(db.Model): - __tablename__ = 'api_based_extensions' + __tablename__ = "api_based_extensions" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='api_based_extension_pkey'), - db.Index('api_based_extension_tenant_idx', 'tenant_id'), + db.PrimaryKeyConstraint("id", name="api_based_extension_pkey"), + db.Index("api_based_extension_tenant_idx", "tenant_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) name = db.Column(db.String(255), nullable=False) api_endpoint = db.Column(db.String(255), nullable=False) api_key = db.Column(db.Text, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) diff --git a/api/models/dataset.py b/api/models/dataset.py index bf3f12a2c5..55f6ed3180 100644 --- a/api/models/dataset.py +++ b/api/models/dataset.py @@ -24,37 +24,34 @@ from .types import StringUUID class DatasetPermissionEnum(str, enum.Enum): - ONLY_ME = 'only_me' - ALL_TEAM = 'all_team_members' - PARTIAL_TEAM = 'partial_members' + ONLY_ME = "only_me" + ALL_TEAM = "all_team_members" + PARTIAL_TEAM = "partial_members" + class Dataset(db.Model): - __tablename__ = 'datasets' + __tablename__ = "datasets" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='dataset_pkey'), - db.Index('dataset_tenant_idx', 'tenant_id'), - db.Index('retrieval_model_idx', "retrieval_model", postgresql_using='gin') + db.PrimaryKeyConstraint("id", name="dataset_pkey"), + db.Index("dataset_tenant_idx", "tenant_id"), + db.Index("retrieval_model_idx", "retrieval_model", postgresql_using="gin"), ) - INDEXING_TECHNIQUE_LIST = ['high_quality', 'economy', None] + INDEXING_TECHNIQUE_LIST = ["high_quality", "economy", None] - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) name = db.Column(db.String(255), nullable=False) description = db.Column(db.Text, nullable=True) - provider = db.Column(db.String(255), nullable=False, - server_default=db.text("'vendor'::character varying")) - permission = db.Column(db.String(255), nullable=False, - server_default=db.text("'only_me'::character varying")) + provider = db.Column(db.String(255), nullable=False, server_default=db.text("'vendor'::character varying")) + permission = db.Column(db.String(255), nullable=False, server_default=db.text("'only_me'::character varying")) data_source_type = db.Column(db.String(255)) indexing_technique = db.Column(db.String(255), nullable=True) index_struct = db.Column(db.Text, nullable=True) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, - server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) updated_by = db.Column(StringUUID, nullable=True) - updated_at = db.Column(db.DateTime, nullable=False, - server_default=db.text('CURRENT_TIMESTAMP(0)')) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) embedding_model = db.Column(db.String(255), nullable=True) embedding_model_provider = db.Column(db.String(255), nullable=True) collection_binding_id = db.Column(StringUUID, nullable=True) @@ -62,8 +59,9 @@ class Dataset(db.Model): @property def dataset_keyword_table(self): - dataset_keyword_table = db.session.query(DatasetKeywordTable).filter( - DatasetKeywordTable.dataset_id == self.id).first() + dataset_keyword_table = ( + db.session.query(DatasetKeywordTable).filter(DatasetKeywordTable.dataset_id == self.id).first() + ) if dataset_keyword_table: return dataset_keyword_table @@ -79,13 +77,19 @@ class Dataset(db.Model): @property def latest_process_rule(self): - return DatasetProcessRule.query.filter(DatasetProcessRule.dataset_id == self.id) \ - .order_by(DatasetProcessRule.created_at.desc()).first() + return ( + DatasetProcessRule.query.filter(DatasetProcessRule.dataset_id == self.id) + .order_by(DatasetProcessRule.created_at.desc()) + .first() + ) @property def app_count(self): - return db.session.query(func.count(AppDatasetJoin.id)).filter(AppDatasetJoin.dataset_id == self.id, - App.id == AppDatasetJoin.app_id).scalar() + return ( + db.session.query(func.count(AppDatasetJoin.id)) + .filter(AppDatasetJoin.dataset_id == self.id, App.id == AppDatasetJoin.app_id) + .scalar() + ) @property def document_count(self): @@ -93,30 +97,40 @@ class Dataset(db.Model): @property def available_document_count(self): - return db.session.query(func.count(Document.id)).filter( - Document.dataset_id == self.id, - Document.indexing_status == 'completed', - Document.enabled == True, - Document.archived == False - ).scalar() + return ( + db.session.query(func.count(Document.id)) + .filter( + Document.dataset_id == self.id, + Document.indexing_status == "completed", + Document.enabled == True, + Document.archived == False, + ) + .scalar() + ) @property def available_segment_count(self): - return db.session.query(func.count(DocumentSegment.id)).filter( - DocumentSegment.dataset_id == self.id, - DocumentSegment.status == 'completed', - DocumentSegment.enabled == True - ).scalar() + return ( + db.session.query(func.count(DocumentSegment.id)) + .filter( + DocumentSegment.dataset_id == self.id, + DocumentSegment.status == "completed", + DocumentSegment.enabled == True, + ) + .scalar() + ) @property def word_count(self): - return Document.query.with_entities(func.coalesce(func.sum(Document.word_count))) \ - .filter(Document.dataset_id == self.id).scalar() + return ( + Document.query.with_entities(func.coalesce(func.sum(Document.word_count))) + .filter(Document.dataset_id == self.id) + .scalar() + ) @property def doc_form(self): - document = db.session.query(Document).filter( - Document.dataset_id == self.id).first() + document = db.session.query(Document).filter(Document.dataset_id == self.id).first() if document: return document.doc_form return None @@ -124,76 +138,68 @@ class Dataset(db.Model): @property def retrieval_model_dict(self): default_retrieval_model = { - 'search_method': RetrievalMethod.SEMANTIC_SEARCH.value, - 'reranking_enable': False, - 'reranking_model': { - 'reranking_provider_name': '', - 'reranking_model_name': '' - }, - 'top_k': 2, - 'score_threshold_enabled': False + "search_method": RetrievalMethod.SEMANTIC_SEARCH.value, + "reranking_enable": False, + "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, + "top_k": 2, + "score_threshold_enabled": False, } return self.retrieval_model if self.retrieval_model else default_retrieval_model @property def tags(self): - tags = db.session.query(Tag).join( - TagBinding, - Tag.id == TagBinding.tag_id - ).filter( - TagBinding.target_id == self.id, - TagBinding.tenant_id == self.tenant_id, - Tag.tenant_id == self.tenant_id, - Tag.type == 'knowledge' - ).all() + tags = ( + db.session.query(Tag) + .join(TagBinding, Tag.id == TagBinding.tag_id) + .filter( + TagBinding.target_id == self.id, + TagBinding.tenant_id == self.tenant_id, + Tag.tenant_id == self.tenant_id, + Tag.type == "knowledge", + ) + .all() + ) return tags if tags else [] @staticmethod def gen_collection_name_by_id(dataset_id: str) -> str: normalized_dataset_id = dataset_id.replace("-", "_") - return f'Vector_index_{normalized_dataset_id}_Node' + return f"Vector_index_{normalized_dataset_id}_Node" class DatasetProcessRule(db.Model): - __tablename__ = 'dataset_process_rules' + __tablename__ = "dataset_process_rules" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='dataset_process_rule_pkey'), - db.Index('dataset_process_rule_dataset_id_idx', 'dataset_id'), + db.PrimaryKeyConstraint("id", name="dataset_process_rule_pkey"), + db.Index("dataset_process_rule_dataset_id_idx", "dataset_id"), ) - id = db.Column(StringUUID, nullable=False, - server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()")) dataset_id = db.Column(StringUUID, nullable=False) - mode = db.Column(db.String(255), nullable=False, - server_default=db.text("'automatic'::character varying")) + mode = db.Column(db.String(255), nullable=False, server_default=db.text("'automatic'::character varying")) rules = db.Column(db.Text, nullable=True) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, - server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) - MODES = ['automatic', 'custom'] - PRE_PROCESSING_RULES = ['remove_stopwords', 'remove_extra_spaces', 'remove_urls_emails'] + MODES = ["automatic", "custom"] + PRE_PROCESSING_RULES = ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"] AUTOMATIC_RULES = { - 'pre_processing_rules': [ - {'id': 'remove_extra_spaces', 'enabled': True}, - {'id': 'remove_urls_emails', 'enabled': False} + "pre_processing_rules": [ + {"id": "remove_extra_spaces", "enabled": True}, + {"id": "remove_urls_emails", "enabled": False}, ], - 'segmentation': { - 'delimiter': '\n', - 'max_tokens': 500, - 'chunk_overlap': 50 - } + "segmentation": {"delimiter": "\n", "max_tokens": 500, "chunk_overlap": 50}, } def to_dict(self): return { - 'id': self.id, - 'dataset_id': self.dataset_id, - 'mode': self.mode, - 'rules': self.rules_dict, - 'created_by': self.created_by, - 'created_at': self.created_at, + "id": self.id, + "dataset_id": self.dataset_id, + "mode": self.mode, + "rules": self.rules_dict, + "created_by": self.created_by, + "created_at": self.created_at, } @property @@ -205,17 +211,16 @@ class DatasetProcessRule(db.Model): class Document(db.Model): - __tablename__ = 'documents' + __tablename__ = "documents" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='document_pkey'), - db.Index('document_dataset_id_idx', 'dataset_id'), - db.Index('document_is_paused_idx', 'is_paused'), - db.Index('document_tenant_idx', 'tenant_id'), + db.PrimaryKeyConstraint("id", name="document_pkey"), + db.Index("document_dataset_id_idx", "dataset_id"), + db.Index("document_is_paused_idx", "is_paused"), + db.Index("document_tenant_idx", "tenant_id"), ) # initial fields - id = db.Column(StringUUID, nullable=False, - server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) dataset_id = db.Column(StringUUID, nullable=False) position = db.Column(db.Integer, nullable=False) @@ -227,8 +232,7 @@ class Document(db.Model): created_from = db.Column(db.String(255), nullable=False) created_by = db.Column(StringUUID, nullable=False) created_api_request_id = db.Column(StringUUID, nullable=True) - created_at = db.Column(db.DateTime, nullable=False, - server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) # start processing processing_started_at = db.Column(db.DateTime, nullable=True) @@ -250,7 +254,7 @@ class Document(db.Model): completed_at = db.Column(db.DateTime, nullable=True) # pause - is_paused = db.Column(db.Boolean, nullable=True, server_default=db.text('false')) + is_paused = db.Column(db.Boolean, nullable=True, server_default=db.text("false")) paused_by = db.Column(StringUUID, nullable=True) paused_at = db.Column(db.DateTime, nullable=True) @@ -259,44 +263,39 @@ class Document(db.Model): stopped_at = db.Column(db.DateTime, nullable=True) # basic fields - indexing_status = db.Column(db.String( - 255), nullable=False, server_default=db.text("'waiting'::character varying")) - enabled = db.Column(db.Boolean, nullable=False, - server_default=db.text('true')) + indexing_status = db.Column(db.String(255), nullable=False, server_default=db.text("'waiting'::character varying")) + enabled = db.Column(db.Boolean, nullable=False, server_default=db.text("true")) disabled_at = db.Column(db.DateTime, nullable=True) disabled_by = db.Column(StringUUID, nullable=True) - archived = db.Column(db.Boolean, nullable=False, - server_default=db.text('false')) + archived = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) archived_reason = db.Column(db.String(255), nullable=True) archived_by = db.Column(StringUUID, nullable=True) archived_at = db.Column(db.DateTime, nullable=True) - updated_at = db.Column(db.DateTime, nullable=False, - server_default=db.text('CURRENT_TIMESTAMP(0)')) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) doc_type = db.Column(db.String(40), nullable=True) doc_metadata = db.Column(db.JSON, nullable=True) - doc_form = db.Column(db.String( - 255), nullable=False, server_default=db.text("'text_model'::character varying")) + doc_form = db.Column(db.String(255), nullable=False, server_default=db.text("'text_model'::character varying")) doc_language = db.Column(db.String(255), nullable=True) - DATA_SOURCES = ['upload_file', 'notion_import', 'website_crawl'] + DATA_SOURCES = ["upload_file", "notion_import", "website_crawl"] @property def display_status(self): status = None - if self.indexing_status == 'waiting': - status = 'queuing' - elif self.indexing_status not in ['completed', 'error', 'waiting'] and self.is_paused: - status = 'paused' - elif self.indexing_status in ['parsing', 'cleaning', 'splitting', 'indexing']: - status = 'indexing' - elif self.indexing_status == 'error': - status = 'error' - elif self.indexing_status == 'completed' and not self.archived and self.enabled: - status = 'available' - elif self.indexing_status == 'completed' and not self.archived and not self.enabled: - status = 'disabled' - elif self.indexing_status == 'completed' and self.archived: - status = 'archived' + if self.indexing_status == "waiting": + status = "queuing" + elif self.indexing_status not in ["completed", "error", "waiting"] and self.is_paused: + status = "paused" + elif self.indexing_status in ["parsing", "cleaning", "splitting", "indexing"]: + status = "indexing" + elif self.indexing_status == "error": + status = "error" + elif self.indexing_status == "completed" and not self.archived and self.enabled: + status = "available" + elif self.indexing_status == "completed" and not self.archived and not self.enabled: + status = "disabled" + elif self.indexing_status == "completed" and self.archived: + status = "archived" return status @property @@ -313,24 +312,26 @@ class Document(db.Model): @property def data_source_detail_dict(self): if self.data_source_info: - if self.data_source_type == 'upload_file': + if self.data_source_type == "upload_file": data_source_info_dict = json.loads(self.data_source_info) - file_detail = db.session.query(UploadFile). \ - filter(UploadFile.id == data_source_info_dict['upload_file_id']). \ - one_or_none() + file_detail = ( + db.session.query(UploadFile) + .filter(UploadFile.id == data_source_info_dict["upload_file_id"]) + .one_or_none() + ) if file_detail: return { - 'upload_file': { - 'id': file_detail.id, - 'name': file_detail.name, - 'size': file_detail.size, - 'extension': file_detail.extension, - 'mime_type': file_detail.mime_type, - 'created_by': file_detail.created_by, - 'created_at': file_detail.created_at.timestamp() + "upload_file": { + "id": file_detail.id, + "name": file_detail.name, + "size": file_detail.size, + "extension": file_detail.extension, + "mime_type": file_detail.mime_type, + "created_by": file_detail.created_by, + "created_at": file_detail.created_at.timestamp(), } } - elif self.data_source_type == 'notion_import' or self.data_source_type == 'website_crawl': + elif self.data_source_type == "notion_import" or self.data_source_type == "website_crawl": return json.loads(self.data_source_info) return {} @@ -356,120 +357,123 @@ class Document(db.Model): @property def hit_count(self): - return DocumentSegment.query.with_entities(func.coalesce(func.sum(DocumentSegment.hit_count))) \ - .filter(DocumentSegment.document_id == self.id).scalar() + return ( + DocumentSegment.query.with_entities(func.coalesce(func.sum(DocumentSegment.hit_count))) + .filter(DocumentSegment.document_id == self.id) + .scalar() + ) def to_dict(self): return { - 'id': self.id, - 'tenant_id': self.tenant_id, - 'dataset_id': self.dataset_id, - 'position': self.position, - 'data_source_type': self.data_source_type, - 'data_source_info': self.data_source_info, - 'dataset_process_rule_id': self.dataset_process_rule_id, - 'batch': self.batch, - 'name': self.name, - 'created_from': self.created_from, - 'created_by': self.created_by, - 'created_api_request_id': self.created_api_request_id, - 'created_at': self.created_at, - 'processing_started_at': self.processing_started_at, - 'file_id': self.file_id, - 'word_count': self.word_count, - 'parsing_completed_at': self.parsing_completed_at, - 'cleaning_completed_at': self.cleaning_completed_at, - 'splitting_completed_at': self.splitting_completed_at, - 'tokens': self.tokens, - 'indexing_latency': self.indexing_latency, - 'completed_at': self.completed_at, - 'is_paused': self.is_paused, - 'paused_by': self.paused_by, - 'paused_at': self.paused_at, - 'error': self.error, - 'stopped_at': self.stopped_at, - 'indexing_status': self.indexing_status, - 'enabled': self.enabled, - 'disabled_at': self.disabled_at, - 'disabled_by': self.disabled_by, - 'archived': self.archived, - 'archived_reason': self.archived_reason, - 'archived_by': self.archived_by, - 'archived_at': self.archived_at, - 'updated_at': self.updated_at, - 'doc_type': self.doc_type, - 'doc_metadata': self.doc_metadata, - 'doc_form': self.doc_form, - 'doc_language': self.doc_language, - 'display_status': self.display_status, - 'data_source_info_dict': self.data_source_info_dict, - 'average_segment_length': self.average_segment_length, - 'dataset_process_rule': self.dataset_process_rule.to_dict() if self.dataset_process_rule else None, - 'dataset': self.dataset.to_dict() if self.dataset else None, - 'segment_count': self.segment_count, - 'hit_count': self.hit_count + "id": self.id, + "tenant_id": self.tenant_id, + "dataset_id": self.dataset_id, + "position": self.position, + "data_source_type": self.data_source_type, + "data_source_info": self.data_source_info, + "dataset_process_rule_id": self.dataset_process_rule_id, + "batch": self.batch, + "name": self.name, + "created_from": self.created_from, + "created_by": self.created_by, + "created_api_request_id": self.created_api_request_id, + "created_at": self.created_at, + "processing_started_at": self.processing_started_at, + "file_id": self.file_id, + "word_count": self.word_count, + "parsing_completed_at": self.parsing_completed_at, + "cleaning_completed_at": self.cleaning_completed_at, + "splitting_completed_at": self.splitting_completed_at, + "tokens": self.tokens, + "indexing_latency": self.indexing_latency, + "completed_at": self.completed_at, + "is_paused": self.is_paused, + "paused_by": self.paused_by, + "paused_at": self.paused_at, + "error": self.error, + "stopped_at": self.stopped_at, + "indexing_status": self.indexing_status, + "enabled": self.enabled, + "disabled_at": self.disabled_at, + "disabled_by": self.disabled_by, + "archived": self.archived, + "archived_reason": self.archived_reason, + "archived_by": self.archived_by, + "archived_at": self.archived_at, + "updated_at": self.updated_at, + "doc_type": self.doc_type, + "doc_metadata": self.doc_metadata, + "doc_form": self.doc_form, + "doc_language": self.doc_language, + "display_status": self.display_status, + "data_source_info_dict": self.data_source_info_dict, + "average_segment_length": self.average_segment_length, + "dataset_process_rule": self.dataset_process_rule.to_dict() if self.dataset_process_rule else None, + "dataset": self.dataset.to_dict() if self.dataset else None, + "segment_count": self.segment_count, + "hit_count": self.hit_count, } @classmethod def from_dict(cls, data: dict): return cls( - id=data.get('id'), - tenant_id=data.get('tenant_id'), - dataset_id=data.get('dataset_id'), - position=data.get('position'), - data_source_type=data.get('data_source_type'), - data_source_info=data.get('data_source_info'), - dataset_process_rule_id=data.get('dataset_process_rule_id'), - batch=data.get('batch'), - name=data.get('name'), - created_from=data.get('created_from'), - created_by=data.get('created_by'), - created_api_request_id=data.get('created_api_request_id'), - created_at=data.get('created_at'), - processing_started_at=data.get('processing_started_at'), - file_id=data.get('file_id'), - word_count=data.get('word_count'), - parsing_completed_at=data.get('parsing_completed_at'), - cleaning_completed_at=data.get('cleaning_completed_at'), - splitting_completed_at=data.get('splitting_completed_at'), - tokens=data.get('tokens'), - indexing_latency=data.get('indexing_latency'), - completed_at=data.get('completed_at'), - is_paused=data.get('is_paused'), - paused_by=data.get('paused_by'), - paused_at=data.get('paused_at'), - error=data.get('error'), - stopped_at=data.get('stopped_at'), - indexing_status=data.get('indexing_status'), - enabled=data.get('enabled'), - disabled_at=data.get('disabled_at'), - disabled_by=data.get('disabled_by'), - archived=data.get('archived'), - archived_reason=data.get('archived_reason'), - archived_by=data.get('archived_by'), - archived_at=data.get('archived_at'), - updated_at=data.get('updated_at'), - doc_type=data.get('doc_type'), - doc_metadata=data.get('doc_metadata'), - doc_form=data.get('doc_form'), - doc_language=data.get('doc_language') + id=data.get("id"), + tenant_id=data.get("tenant_id"), + dataset_id=data.get("dataset_id"), + position=data.get("position"), + data_source_type=data.get("data_source_type"), + data_source_info=data.get("data_source_info"), + dataset_process_rule_id=data.get("dataset_process_rule_id"), + batch=data.get("batch"), + name=data.get("name"), + created_from=data.get("created_from"), + created_by=data.get("created_by"), + created_api_request_id=data.get("created_api_request_id"), + created_at=data.get("created_at"), + processing_started_at=data.get("processing_started_at"), + file_id=data.get("file_id"), + word_count=data.get("word_count"), + parsing_completed_at=data.get("parsing_completed_at"), + cleaning_completed_at=data.get("cleaning_completed_at"), + splitting_completed_at=data.get("splitting_completed_at"), + tokens=data.get("tokens"), + indexing_latency=data.get("indexing_latency"), + completed_at=data.get("completed_at"), + is_paused=data.get("is_paused"), + paused_by=data.get("paused_by"), + paused_at=data.get("paused_at"), + error=data.get("error"), + stopped_at=data.get("stopped_at"), + indexing_status=data.get("indexing_status"), + enabled=data.get("enabled"), + disabled_at=data.get("disabled_at"), + disabled_by=data.get("disabled_by"), + archived=data.get("archived"), + archived_reason=data.get("archived_reason"), + archived_by=data.get("archived_by"), + archived_at=data.get("archived_at"), + updated_at=data.get("updated_at"), + doc_type=data.get("doc_type"), + doc_metadata=data.get("doc_metadata"), + doc_form=data.get("doc_form"), + doc_language=data.get("doc_language"), ) + class DocumentSegment(db.Model): - __tablename__ = 'document_segments' + __tablename__ = "document_segments" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='document_segment_pkey'), - db.Index('document_segment_dataset_id_idx', 'dataset_id'), - db.Index('document_segment_document_id_idx', 'document_id'), - db.Index('document_segment_tenant_dataset_idx', 'dataset_id', 'tenant_id'), - db.Index('document_segment_tenant_document_idx', 'document_id', 'tenant_id'), - db.Index('document_segment_dataset_node_idx', 'dataset_id', 'index_node_id'), - db.Index('document_segment_tenant_idx', 'tenant_id'), + db.PrimaryKeyConstraint("id", name="document_segment_pkey"), + db.Index("document_segment_dataset_id_idx", "dataset_id"), + db.Index("document_segment_document_id_idx", "document_id"), + db.Index("document_segment_tenant_dataset_idx", "dataset_id", "tenant_id"), + db.Index("document_segment_tenant_document_idx", "document_id", "tenant_id"), + db.Index("document_segment_dataset_node_idx", "dataset_id", "index_node_id"), + db.Index("document_segment_tenant_idx", "tenant_id"), ) # initial fields - id = db.Column(StringUUID, nullable=False, - server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) dataset_id = db.Column(StringUUID, nullable=False) document_id = db.Column(StringUUID, nullable=False) @@ -486,18 +490,14 @@ class DocumentSegment(db.Model): # basic fields hit_count = db.Column(db.Integer, nullable=False, default=0) - enabled = db.Column(db.Boolean, nullable=False, - server_default=db.text('true')) + enabled = db.Column(db.Boolean, nullable=False, server_default=db.text("true")) disabled_at = db.Column(db.DateTime, nullable=True) disabled_by = db.Column(StringUUID, nullable=True) - status = db.Column(db.String(255), nullable=False, - server_default=db.text("'waiting'::character varying")) + status = db.Column(db.String(255), nullable=False, server_default=db.text("'waiting'::character varying")) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, - server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) updated_by = db.Column(StringUUID, nullable=True) - updated_at = db.Column(db.DateTime, nullable=False, - server_default=db.text('CURRENT_TIMESTAMP(0)')) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) indexing_at = db.Column(db.DateTime, nullable=True) completed_at = db.Column(db.DateTime, nullable=True) error = db.Column(db.Text, nullable=True) @@ -513,17 +513,19 @@ class DocumentSegment(db.Model): @property def previous_segment(self): - return db.session.query(DocumentSegment).filter( - DocumentSegment.document_id == self.document_id, - DocumentSegment.position == self.position - 1 - ).first() + return ( + db.session.query(DocumentSegment) + .filter(DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position - 1) + .first() + ) @property def next_segment(self): - return db.session.query(DocumentSegment).filter( - DocumentSegment.document_id == self.document_id, - DocumentSegment.position == self.position + 1 - ).first() + return ( + db.session.query(DocumentSegment) + .filter(DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position + 1) + .first() + ) def get_sign_content(self): pattern = r"/files/([a-f0-9\-]+)/image-preview" @@ -535,7 +537,7 @@ class DocumentSegment(db.Model): nonce = os.urandom(16).hex() timestamp = str(int(time.time())) data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}" - secret_key = dify_config.SECRET_KEY.encode() if dify_config.SECRET_KEY else b'' + secret_key = dify_config.SECRET_KEY.encode() if dify_config.SECRET_KEY else b"" sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() encoded_sign = base64.urlsafe_b64encode(sign).decode() @@ -546,21 +548,20 @@ class DocumentSegment(db.Model): # Reconstruct the text with signed URLs offset = 0 for start, end, signed_url in signed_urls: - text = text[:start + offset] + signed_url + text[end + offset:] + text = text[: start + offset] + signed_url + text[end + offset :] offset += len(signed_url) - (end - start) return text - class AppDatasetJoin(db.Model): - __tablename__ = 'app_dataset_joins' + __tablename__ = "app_dataset_joins" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='app_dataset_join_pkey'), - db.Index('app_dataset_join_app_dataset_idx', 'dataset_id', 'app_id'), + db.PrimaryKeyConstraint("id", name="app_dataset_join_pkey"), + db.Index("app_dataset_join_app_dataset_idx", "dataset_id", "app_id"), ) - id = db.Column(StringUUID, primary_key=True, nullable=False, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, primary_key=True, nullable=False, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) dataset_id = db.Column(StringUUID, nullable=False) created_at = db.Column(db.DateTime, nullable=False, server_default=db.func.current_timestamp()) @@ -571,13 +572,13 @@ class AppDatasetJoin(db.Model): class DatasetQuery(db.Model): - __tablename__ = 'dataset_queries' + __tablename__ = "dataset_queries" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='dataset_query_pkey'), - db.Index('dataset_query_dataset_id_idx', 'dataset_id'), + db.PrimaryKeyConstraint("id", name="dataset_query_pkey"), + db.Index("dataset_query_dataset_id_idx", "dataset_id"), ) - id = db.Column(StringUUID, primary_key=True, nullable=False, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, primary_key=True, nullable=False, server_default=db.text("uuid_generate_v4()")) dataset_id = db.Column(StringUUID, nullable=False) content = db.Column(db.Text, nullable=False) source = db.Column(db.String(255), nullable=False) @@ -588,17 +589,18 @@ class DatasetQuery(db.Model): class DatasetKeywordTable(db.Model): - __tablename__ = 'dataset_keyword_tables' + __tablename__ = "dataset_keyword_tables" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='dataset_keyword_table_pkey'), - db.Index('dataset_keyword_table_dataset_id_idx', 'dataset_id'), + db.PrimaryKeyConstraint("id", name="dataset_keyword_table_pkey"), + db.Index("dataset_keyword_table_dataset_id_idx", "dataset_id"), ) - id = db.Column(StringUUID, primary_key=True, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, primary_key=True, server_default=db.text("uuid_generate_v4()")) dataset_id = db.Column(StringUUID, nullable=False, unique=True) keyword_table = db.Column(db.Text, nullable=False) - data_source_type = db.Column(db.String(255), nullable=False, - server_default=db.text("'database'::character varying")) + data_source_type = db.Column( + db.String(255), nullable=False, server_default=db.text("'database'::character varying") + ) @property def keyword_table_dict(self): @@ -614,19 +616,17 @@ class DatasetKeywordTable(db.Model): return dct # get dataset - dataset = Dataset.query.filter_by( - id=self.dataset_id - ).first() + dataset = Dataset.query.filter_by(id=self.dataset_id).first() if not dataset: return None - if self.data_source_type == 'database': + if self.data_source_type == "database": return json.loads(self.keyword_table, cls=SetDecoder) if self.keyword_table else None else: - file_key = 'keyword_files/' + dataset.tenant_id + '/' + self.dataset_id + '.txt' + file_key = "keyword_files/" + dataset.tenant_id + "/" + self.dataset_id + ".txt" try: keyword_table_text = storage.load_once(file_key) if keyword_table_text: - return json.loads(keyword_table_text.decode('utf-8'), cls=SetDecoder) + return json.loads(keyword_table_text.decode("utf-8"), cls=SetDecoder) return None except Exception as e: logging.exception(str(e)) @@ -634,21 +634,21 @@ class DatasetKeywordTable(db.Model): class Embedding(db.Model): - __tablename__ = 'embeddings' + __tablename__ = "embeddings" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='embedding_pkey'), - db.UniqueConstraint('model_name', 'hash', 'provider_name', name='embedding_hash_idx'), - db.Index('created_at_idx', 'created_at') + db.PrimaryKeyConstraint("id", name="embedding_pkey"), + db.UniqueConstraint("model_name", "hash", "provider_name", name="embedding_hash_idx"), + db.Index("created_at_idx", "created_at"), ) - id = db.Column(StringUUID, primary_key=True, server_default=db.text('uuid_generate_v4()')) - model_name = db.Column(db.String(255), nullable=False, - server_default=db.text("'text-embedding-ada-002'::character varying")) + id = db.Column(StringUUID, primary_key=True, server_default=db.text("uuid_generate_v4()")) + model_name = db.Column( + db.String(255), nullable=False, server_default=db.text("'text-embedding-ada-002'::character varying") + ) hash = db.Column(db.String(64), nullable=False) embedding = db.Column(db.LargeBinary, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - provider_name = db.Column(db.String(255), nullable=False, - server_default=db.text("''::character varying")) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + provider_name = db.Column(db.String(255), nullable=False, server_default=db.text("''::character varying")) def set_embedding(self, embedding_data: list[float]): self.embedding = pickle.dumps(embedding_data, protocol=pickle.HIGHEST_PROTOCOL) @@ -658,33 +658,32 @@ class Embedding(db.Model): class DatasetCollectionBinding(db.Model): - __tablename__ = 'dataset_collection_bindings' + __tablename__ = "dataset_collection_bindings" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='dataset_collection_bindings_pkey'), - db.Index('provider_model_name_idx', 'provider_name', 'model_name') - + db.PrimaryKeyConstraint("id", name="dataset_collection_bindings_pkey"), + db.Index("provider_model_name_idx", "provider_name", "model_name"), ) - id = db.Column(StringUUID, primary_key=True, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, primary_key=True, server_default=db.text("uuid_generate_v4()")) provider_name = db.Column(db.String(40), nullable=False) model_name = db.Column(db.String(255), nullable=False) type = db.Column(db.String(40), server_default=db.text("'dataset'::character varying"), nullable=False) collection_name = db.Column(db.String(64), nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class DatasetPermission(db.Model): - __tablename__ = 'dataset_permissions' + __tablename__ = "dataset_permissions" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='dataset_permission_pkey'), - db.Index('idx_dataset_permissions_dataset_id', 'dataset_id'), - db.Index('idx_dataset_permissions_account_id', 'account_id'), - db.Index('idx_dataset_permissions_tenant_id', 'tenant_id') + db.PrimaryKeyConstraint("id", name="dataset_permission_pkey"), + db.Index("idx_dataset_permissions_dataset_id", "dataset_id"), + db.Index("idx_dataset_permissions_account_id", "account_id"), + db.Index("idx_dataset_permissions_tenant_id", "tenant_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'), primary_key=True) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"), primary_key=True) dataset_id = db.Column(StringUUID, nullable=False) account_id = db.Column(StringUUID, nullable=False) tenant_id = db.Column(StringUUID, nullable=False) - has_permission = db.Column(db.Boolean, nullable=False, server_default=db.text('true')) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + has_permission = db.Column(db.Boolean, nullable=False, server_default=db.text("true")) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) diff --git a/api/models/model.py b/api/models/model.py index e81d25fbc9..8ab3026522 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -20,25 +20,23 @@ from .types import StringUUID class DifySetup(db.Model): - __tablename__ = 'dify_setups' - __table_args__ = ( - db.PrimaryKeyConstraint('version', name='dify_setup_pkey'), - ) + __tablename__ = "dify_setups" + __table_args__ = (db.PrimaryKeyConstraint("version", name="dify_setup_pkey"),) version = db.Column(db.String(255), nullable=False) - setup_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + setup_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class AppMode(Enum): - COMPLETION = 'completion' - WORKFLOW = 'workflow' - CHAT = 'chat' - ADVANCED_CHAT = 'advanced-chat' - AGENT_CHAT = 'agent-chat' - CHANNEL = 'channel' + COMPLETION = "completion" + WORKFLOW = "workflow" + CHAT = "chat" + ADVANCED_CHAT = "advanced-chat" + AGENT_CHAT = "agent-chat" + CHANNEL = "channel" @classmethod - def value_of(cls, value: str) -> 'AppMode': + def value_of(cls, value: str) -> "AppMode": """ Get value of given mode. @@ -48,21 +46,19 @@ class AppMode(Enum): for mode in cls: if mode.value == value: return mode - raise ValueError(f'invalid mode value {value}') + raise ValueError(f"invalid mode value {value}") class IconType(Enum): IMAGE = "image" EMOJI = "emoji" -class App(db.Model): - __tablename__ = 'apps' - __table_args__ = ( - db.PrimaryKeyConstraint('id', name='app_pkey'), - db.Index('app_tenant_id_idx', 'tenant_id') - ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) +class App(db.Model): + __tablename__ = "apps" + __table_args__ = (db.PrimaryKeyConstraint("id", name="app_pkey"), db.Index("app_tenant_id_idx", "tenant_id")) + + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) name = db.Column(db.String(255), nullable=False) description = db.Column(db.Text, nullable=False, server_default=db.text("''::character varying")) @@ -75,17 +71,17 @@ class App(db.Model): status = db.Column(db.String(255), nullable=False, server_default=db.text("'normal'::character varying")) enable_site = db.Column(db.Boolean, nullable=False) enable_api = db.Column(db.Boolean, nullable=False) - api_rpm = db.Column(db.Integer, nullable=False, server_default=db.text('0')) - api_rph = db.Column(db.Integer, nullable=False, server_default=db.text('0')) - is_demo = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) - is_public = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) - is_universal = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) + api_rpm = db.Column(db.Integer, nullable=False, server_default=db.text("0")) + api_rph = db.Column(db.Integer, nullable=False, server_default=db.text("0")) + is_demo = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) + is_public = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) + is_universal = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) tracing = db.Column(db.Text, nullable=True) max_active_requests = db.Column(db.Integer, nullable=True) created_by = db.Column(StringUUID, nullable=True) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) updated_by = db.Column(StringUUID, nullable=True) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) use_icon_as_answer_icon = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) @property @@ -97,7 +93,7 @@ class App(db.Model): if app_model_config: return app_model_config.pre_prompt else: - return '' + return "" @property def site(self): @@ -105,24 +101,24 @@ class App(db.Model): return site @property - def app_model_config(self) -> Optional['AppModelConfig']: + def app_model_config(self) -> Optional["AppModelConfig"]: if self.app_model_config_id: return db.session.query(AppModelConfig).filter(AppModelConfig.id == self.app_model_config_id).first() return None @property - def workflow(self) -> Optional['Workflow']: + def workflow(self) -> Optional["Workflow"]: if self.workflow_id: from .workflow import Workflow + return db.session.query(Workflow).filter(Workflow.id == self.workflow_id).first() return None @property def api_base_url(self): - return (dify_config.SERVICE_API_URL if dify_config.SERVICE_API_URL - else request.host_url.rstrip('/')) + '/v1' + return (dify_config.SERVICE_API_URL if dify_config.SERVICE_API_URL else request.host_url.rstrip("/")) + "/v1" @property def tenant(self): @@ -136,8 +132,9 @@ class App(db.Model): return False if not app_model_config.agent_mode: return False - if self.app_model_config.agent_mode_dict.get('enabled', False) \ - and self.app_model_config.agent_mode_dict.get('strategy', '') in ['function_call', 'react']: + if self.app_model_config.agent_mode_dict.get("enabled", False) and self.app_model_config.agent_mode_dict.get( + "strategy", "" + ) in ["function_call", "react"]: self.mode = AppMode.AGENT_CHAT.value db.session.commit() return True @@ -159,16 +156,16 @@ class App(db.Model): if not app_model_config.agent_mode: return [] agent_mode = app_model_config.agent_mode_dict - tools = agent_mode.get('tools', []) + tools = agent_mode.get("tools", []) provider_ids = [] for tool in tools: keys = list(tool.keys()) if len(keys) >= 4: - provider_type = tool.get('provider_type', '') - provider_id = tool.get('provider_id', '') - if provider_type == 'api': + provider_type = tool.get("provider_type", "") + provider_id = tool.get("provider_id", "") + if provider_type == "api": # check if provider id is a uuid string, if not, skip try: uuid.UUID(provider_id) @@ -180,8 +177,7 @@ class App(db.Model): return [] api_providers = db.session.execute( - text('SELECT id FROM tool_api_providers WHERE id IN :provider_ids'), - {'provider_ids': tuple(provider_ids)} + text("SELECT id FROM tool_api_providers WHERE id IN :provider_ids"), {"provider_ids": tuple(provider_ids)} ).fetchall() deleted_tools = [] @@ -190,44 +186,43 @@ class App(db.Model): for tool in tools: keys = list(tool.keys()) if len(keys) >= 4: - provider_type = tool.get('provider_type', '') - provider_id = tool.get('provider_id', '') - if provider_type == 'api' and provider_id not in current_api_provider_ids: - deleted_tools.append(tool['tool_name']) + provider_type = tool.get("provider_type", "") + provider_id = tool.get("provider_id", "") + if provider_type == "api" and provider_id not in current_api_provider_ids: + deleted_tools.append(tool["tool_name"]) return deleted_tools @property def tags(self): - tags = db.session.query(Tag).join( - TagBinding, - Tag.id == TagBinding.tag_id - ).filter( - TagBinding.target_id == self.id, - TagBinding.tenant_id == self.tenant_id, - Tag.tenant_id == self.tenant_id, - Tag.type == 'app' - ).all() + tags = ( + db.session.query(Tag) + .join(TagBinding, Tag.id == TagBinding.tag_id) + .filter( + TagBinding.target_id == self.id, + TagBinding.tenant_id == self.tenant_id, + Tag.tenant_id == self.tenant_id, + Tag.type == "app", + ) + .all() + ) return tags if tags else [] class AppModelConfig(db.Model): - __tablename__ = 'app_model_configs' - __table_args__ = ( - db.PrimaryKeyConstraint('id', name='app_model_config_pkey'), - db.Index('app_app_id_idx', 'app_id') - ) + __tablename__ = "app_model_configs" + __table_args__ = (db.PrimaryKeyConstraint("id", name="app_model_config_pkey"), db.Index("app_app_id_idx", "app_id")) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) provider = db.Column(db.String(255), nullable=True) model_id = db.Column(db.String(255), nullable=True) configs = db.Column(db.JSON, nullable=True) created_by = db.Column(StringUUID, nullable=True) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) updated_by = db.Column(StringUUID, nullable=True) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) opening_statement = db.Column(db.Text) suggested_questions = db.Column(db.Text) suggested_questions_after_answer = db.Column(db.Text) @@ -263,28 +258,29 @@ class AppModelConfig(db.Model): @property def suggested_questions_after_answer_dict(self) -> dict: - return json.loads(self.suggested_questions_after_answer) if self.suggested_questions_after_answer \ + return ( + json.loads(self.suggested_questions_after_answer) + if self.suggested_questions_after_answer else {"enabled": False} + ) @property def speech_to_text_dict(self) -> dict: - return json.loads(self.speech_to_text) if self.speech_to_text \ - else {"enabled": False} + return json.loads(self.speech_to_text) if self.speech_to_text else {"enabled": False} @property def text_to_speech_dict(self) -> dict: - return json.loads(self.text_to_speech) if self.text_to_speech \ - else {"enabled": False} + return json.loads(self.text_to_speech) if self.text_to_speech else {"enabled": False} @property def retriever_resource_dict(self) -> dict: - return json.loads(self.retriever_resource) if self.retriever_resource \ - else {"enabled": True} + return json.loads(self.retriever_resource) if self.retriever_resource else {"enabled": True} @property def annotation_reply_dict(self) -> dict: - annotation_setting = db.session.query(AppAnnotationSetting).filter( - AppAnnotationSetting.app_id == self.app_id).first() + annotation_setting = ( + db.session.query(AppAnnotationSetting).filter(AppAnnotationSetting.app_id == self.app_id).first() + ) if annotation_setting: collection_binding_detail = annotation_setting.collection_binding_detail return { @@ -293,8 +289,8 @@ class AppModelConfig(db.Model): "score_threshold": annotation_setting.score_threshold, "embedding_model": { "embedding_provider_name": collection_binding_detail.provider_name, - "embedding_model_name": collection_binding_detail.model_name - } + "embedding_model_name": collection_binding_detail.model_name, + }, } else: @@ -306,13 +302,15 @@ class AppModelConfig(db.Model): @property def sensitive_word_avoidance_dict(self) -> dict: - return json.loads(self.sensitive_word_avoidance) if self.sensitive_word_avoidance \ + return ( + json.loads(self.sensitive_word_avoidance) + if self.sensitive_word_avoidance else {"enabled": False, "type": "", "configs": []} + ) @property def external_data_tools_list(self) -> list[dict]: - return json.loads(self.external_data_tools) if self.external_data_tools \ - else [] + return json.loads(self.external_data_tools) if self.external_data_tools else [] @property def user_input_form_list(self) -> dict: @@ -320,8 +318,11 @@ class AppModelConfig(db.Model): @property def agent_mode_dict(self) -> dict: - return json.loads(self.agent_mode) if self.agent_mode else {"enabled": False, "strategy": None, "tools": [], - "prompt": None} + return ( + json.loads(self.agent_mode) + if self.agent_mode + else {"enabled": False, "strategy": None, "tools": [], "prompt": None} + ) @property def chat_prompt_config_dict(self) -> dict: @@ -335,19 +336,28 @@ class AppModelConfig(db.Model): def dataset_configs_dict(self) -> dict: if self.dataset_configs: dataset_configs = json.loads(self.dataset_configs) - if 'retrieval_model' not in dataset_configs: - return {'retrieval_model': 'single'} + if "retrieval_model" not in dataset_configs: + return {"retrieval_model": "single"} else: return dataset_configs return { - 'retrieval_model': 'multiple', - } + "retrieval_model": "multiple", + } @property def file_upload_dict(self) -> dict: - return json.loads(self.file_upload) if self.file_upload else { - "image": {"enabled": False, "number_limits": 3, "detail": "high", - "transfer_methods": ["remote_url", "local_file"]}} + return ( + json.loads(self.file_upload) + if self.file_upload + else { + "image": { + "enabled": False, + "number_limits": 3, + "detail": "high", + "transfer_methods": ["remote_url", "local_file"], + } + } + ) def to_dict(self) -> dict: return { @@ -370,44 +380,53 @@ class AppModelConfig(db.Model): "chat_prompt_config": self.chat_prompt_config_dict, "completion_prompt_config": self.completion_prompt_config_dict, "dataset_configs": self.dataset_configs_dict, - "file_upload": self.file_upload_dict + "file_upload": self.file_upload_dict, } def from_model_config_dict(self, model_config: dict): - self.opening_statement = model_config.get('opening_statement') - self.suggested_questions = json.dumps(model_config['suggested_questions']) \ - if model_config.get('suggested_questions') else None - self.suggested_questions_after_answer = json.dumps(model_config['suggested_questions_after_answer']) \ - if model_config.get('suggested_questions_after_answer') else None - self.speech_to_text = json.dumps(model_config['speech_to_text']) \ - if model_config.get('speech_to_text') else None - self.text_to_speech = json.dumps(model_config['text_to_speech']) \ - if model_config.get('text_to_speech') else None - self.more_like_this = json.dumps(model_config['more_like_this']) \ - if model_config.get('more_like_this') else None - self.sensitive_word_avoidance = json.dumps(model_config['sensitive_word_avoidance']) \ - if model_config.get('sensitive_word_avoidance') else None - self.external_data_tools = json.dumps(model_config['external_data_tools']) \ - if model_config.get('external_data_tools') else None - self.model = json.dumps(model_config['model']) \ - if model_config.get('model') else None - self.user_input_form = json.dumps(model_config['user_input_form']) \ - if model_config.get('user_input_form') else None - self.dataset_query_variable = model_config.get('dataset_query_variable') - self.pre_prompt = model_config['pre_prompt'] - self.agent_mode = json.dumps(model_config['agent_mode']) \ - if model_config.get('agent_mode') else None - self.retriever_resource = json.dumps(model_config['retriever_resource']) \ - if model_config.get('retriever_resource') else None - self.prompt_type = model_config.get('prompt_type', 'simple') - self.chat_prompt_config = json.dumps(model_config.get('chat_prompt_config')) \ - if model_config.get('chat_prompt_config') else None - self.completion_prompt_config = json.dumps(model_config.get('completion_prompt_config')) \ - if model_config.get('completion_prompt_config') else None - self.dataset_configs = json.dumps(model_config.get('dataset_configs')) \ - if model_config.get('dataset_configs') else None - self.file_upload = json.dumps(model_config.get('file_upload')) \ - if model_config.get('file_upload') else None + self.opening_statement = model_config.get("opening_statement") + self.suggested_questions = ( + json.dumps(model_config["suggested_questions"]) if model_config.get("suggested_questions") else None + ) + self.suggested_questions_after_answer = ( + json.dumps(model_config["suggested_questions_after_answer"]) + if model_config.get("suggested_questions_after_answer") + else None + ) + self.speech_to_text = json.dumps(model_config["speech_to_text"]) if model_config.get("speech_to_text") else None + self.text_to_speech = json.dumps(model_config["text_to_speech"]) if model_config.get("text_to_speech") else None + self.more_like_this = json.dumps(model_config["more_like_this"]) if model_config.get("more_like_this") else None + self.sensitive_word_avoidance = ( + json.dumps(model_config["sensitive_word_avoidance"]) + if model_config.get("sensitive_word_avoidance") + else None + ) + self.external_data_tools = ( + json.dumps(model_config["external_data_tools"]) if model_config.get("external_data_tools") else None + ) + self.model = json.dumps(model_config["model"]) if model_config.get("model") else None + self.user_input_form = ( + json.dumps(model_config["user_input_form"]) if model_config.get("user_input_form") else None + ) + self.dataset_query_variable = model_config.get("dataset_query_variable") + self.pre_prompt = model_config["pre_prompt"] + self.agent_mode = json.dumps(model_config["agent_mode"]) if model_config.get("agent_mode") else None + self.retriever_resource = ( + json.dumps(model_config["retriever_resource"]) if model_config.get("retriever_resource") else None + ) + self.prompt_type = model_config.get("prompt_type", "simple") + self.chat_prompt_config = ( + json.dumps(model_config.get("chat_prompt_config")) if model_config.get("chat_prompt_config") else None + ) + self.completion_prompt_config = ( + json.dumps(model_config.get("completion_prompt_config")) + if model_config.get("completion_prompt_config") + else None + ) + self.dataset_configs = ( + json.dumps(model_config.get("dataset_configs")) if model_config.get("dataset_configs") else None + ) + self.file_upload = json.dumps(model_config.get("file_upload")) if model_config.get("file_upload") else None return self def copy(self): @@ -432,21 +451,21 @@ class AppModelConfig(db.Model): chat_prompt_config=self.chat_prompt_config, completion_prompt_config=self.completion_prompt_config, dataset_configs=self.dataset_configs, - file_upload=self.file_upload + file_upload=self.file_upload, ) return new_app_model_config class RecommendedApp(db.Model): - __tablename__ = 'recommended_apps' + __tablename__ = "recommended_apps" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='recommended_app_pkey'), - db.Index('recommended_app_app_id_idx', 'app_id'), - db.Index('recommended_app_is_listed_idx', 'is_listed', 'language') + db.PrimaryKeyConstraint("id", name="recommended_app_pkey"), + db.Index("recommended_app_app_id_idx", "app_id"), + db.Index("recommended_app_is_listed_idx", "is_listed", "language"), ) - id = db.Column(StringUUID, primary_key=True, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, primary_key=True, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) description = db.Column(db.JSON, nullable=False) copyright = db.Column(db.String(255), nullable=False) @@ -457,8 +476,8 @@ class RecommendedApp(db.Model): is_listed = db.Column(db.Boolean, nullable=False, default=True) install_count = db.Column(db.Integer, nullable=False, default=0) language = db.Column(db.String(255), nullable=False, server_default=db.text("'en-US'::character varying")) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def app(self): @@ -467,22 +486,22 @@ class RecommendedApp(db.Model): class InstalledApp(db.Model): - __tablename__ = 'installed_apps' + __tablename__ = "installed_apps" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='installed_app_pkey'), - db.Index('installed_app_tenant_id_idx', 'tenant_id'), - db.Index('installed_app_app_id_idx', 'app_id'), - db.UniqueConstraint('tenant_id', 'app_id', name='unique_tenant_app') + db.PrimaryKeyConstraint("id", name="installed_app_pkey"), + db.Index("installed_app_tenant_id_idx", "tenant_id"), + db.Index("installed_app_app_id_idx", "app_id"), + db.UniqueConstraint("tenant_id", "app_id", name="unique_tenant_app"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) app_id = db.Column(StringUUID, nullable=False) app_owner_tenant_id = db.Column(StringUUID, nullable=False) position = db.Column(db.Integer, nullable=False, default=0) - is_pinned = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) + is_pinned = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) last_used_at = db.Column(db.DateTime, nullable=True) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def app(self): @@ -496,13 +515,13 @@ class InstalledApp(db.Model): class Conversation(db.Model): - __tablename__ = 'conversations' + __tablename__ = "conversations" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='conversation_pkey'), - db.Index('conversation_app_from_user_idx', 'app_id', 'from_source', 'from_end_user_id') + db.PrimaryKeyConstraint("id", name="conversation_pkey"), + db.Index("conversation_app_from_user_idx", "app_id", "from_source", "from_end_user_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) app_model_config_id = db.Column(StringUUID, nullable=True) model_provider = db.Column(db.String(255), nullable=True) @@ -514,7 +533,7 @@ class Conversation(db.Model): inputs = db.Column(db.JSON) introduction = db.Column(db.Text) system_instruction = db.Column(db.Text) - system_instruction_tokens = db.Column(db.Integer, nullable=False, server_default=db.text('0')) + system_instruction_tokens = db.Column(db.Integer, nullable=False, server_default=db.text("0")) status = db.Column(db.String(255), nullable=False) invoke_from = db.Column(db.String(255), nullable=True) from_source = db.Column(db.String(255), nullable=False) @@ -523,13 +542,15 @@ class Conversation(db.Model): read_at = db.Column(db.DateTime) read_account_id = db.Column(StringUUID) dialogue_count: Mapped[int] = mapped_column(default=0) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) - messages = db.relationship("Message", backref="conversation", lazy='select', passive_deletes="all") - message_annotations = db.relationship("MessageAnnotation", backref="conversation", lazy='select', passive_deletes="all") + messages = db.relationship("Message", backref="conversation", lazy="select", passive_deletes="all") + message_annotations = db.relationship( + "MessageAnnotation", backref="conversation", lazy="select", passive_deletes="all" + ) - is_deleted = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) + is_deleted = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) @property def model_config(self): @@ -542,20 +563,21 @@ class Conversation(db.Model): if self.override_model_configs: override_model_configs = json.loads(self.override_model_configs) - if 'model' in override_model_configs: + if "model" in override_model_configs: app_model_config = AppModelConfig() app_model_config = app_model_config.from_model_config_dict(override_model_configs) model_config = app_model_config.to_dict() else: - model_config['configs'] = override_model_configs + model_config["configs"] = override_model_configs else: - app_model_config = db.session.query(AppModelConfig).filter( - AppModelConfig.id == self.app_model_config_id).first() + app_model_config = ( + db.session.query(AppModelConfig).filter(AppModelConfig.id == self.app_model_config_id).first() + ) model_config = app_model_config.to_dict() - model_config['model_id'] = self.model_id - model_config['provider'] = self.model_provider + model_config["model_id"] = self.model_id + model_config["provider"] = self.model_provider return model_config @@ -568,7 +590,7 @@ class Conversation(db.Model): if first_message: return first_message.query else: - return '' + return "" @property def annotated(self): @@ -584,31 +606,51 @@ class Conversation(db.Model): @property def user_feedback_stats(self): - like = db.session.query(MessageFeedback) \ - .filter(MessageFeedback.conversation_id == self.id, - MessageFeedback.from_source == 'user', - MessageFeedback.rating == 'like').count() + like = ( + db.session.query(MessageFeedback) + .filter( + MessageFeedback.conversation_id == self.id, + MessageFeedback.from_source == "user", + MessageFeedback.rating == "like", + ) + .count() + ) - dislike = db.session.query(MessageFeedback) \ - .filter(MessageFeedback.conversation_id == self.id, - MessageFeedback.from_source == 'user', - MessageFeedback.rating == 'dislike').count() + dislike = ( + db.session.query(MessageFeedback) + .filter( + MessageFeedback.conversation_id == self.id, + MessageFeedback.from_source == "user", + MessageFeedback.rating == "dislike", + ) + .count() + ) - return {'like': like, 'dislike': dislike} + return {"like": like, "dislike": dislike} @property def admin_feedback_stats(self): - like = db.session.query(MessageFeedback) \ - .filter(MessageFeedback.conversation_id == self.id, - MessageFeedback.from_source == 'admin', - MessageFeedback.rating == 'like').count() + like = ( + db.session.query(MessageFeedback) + .filter( + MessageFeedback.conversation_id == self.id, + MessageFeedback.from_source == "admin", + MessageFeedback.rating == "like", + ) + .count() + ) - dislike = db.session.query(MessageFeedback) \ - .filter(MessageFeedback.conversation_id == self.id, - MessageFeedback.from_source == 'admin', - MessageFeedback.rating == 'dislike').count() + dislike = ( + db.session.query(MessageFeedback) + .filter( + MessageFeedback.conversation_id == self.id, + MessageFeedback.from_source == "admin", + MessageFeedback.rating == "dislike", + ) + .count() + ) - return {'like': like, 'dislike': dislike} + return {"like": like, "dislike": dislike} @property def first_message(self): @@ -642,33 +684,33 @@ class Conversation(db.Model): class Message(db.Model): - __tablename__ = 'messages' + __tablename__ = "messages" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='message_pkey'), - db.Index('message_app_id_idx', 'app_id', 'created_at'), - db.Index('message_conversation_id_idx', 'conversation_id'), - db.Index('message_end_user_idx', 'app_id', 'from_source', 'from_end_user_id'), - db.Index('message_account_idx', 'app_id', 'from_source', 'from_account_id'), - db.Index('message_workflow_run_id_idx', 'conversation_id', 'workflow_run_id') + db.PrimaryKeyConstraint("id", name="message_pkey"), + db.Index("message_app_id_idx", "app_id", "created_at"), + db.Index("message_conversation_id_idx", "conversation_id"), + db.Index("message_end_user_idx", "app_id", "from_source", "from_end_user_id"), + db.Index("message_account_idx", "app_id", "from_source", "from_account_id"), + db.Index("message_workflow_run_id_idx", "conversation_id", "workflow_run_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) model_provider = db.Column(db.String(255), nullable=True) model_id = db.Column(db.String(255), nullable=True) override_model_configs = db.Column(db.Text) - conversation_id = db.Column(StringUUID, db.ForeignKey('conversations.id'), nullable=False) + conversation_id = db.Column(StringUUID, db.ForeignKey("conversations.id"), nullable=False) inputs = db.Column(db.JSON) query = db.Column(db.Text, nullable=False) message = db.Column(db.JSON, nullable=False) - message_tokens = db.Column(db.Integer, nullable=False, server_default=db.text('0')) + message_tokens = db.Column(db.Integer, nullable=False, server_default=db.text("0")) message_unit_price = db.Column(db.Numeric(10, 4), nullable=False) - message_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text('0.001')) + message_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text("0.001")) answer = db.Column(db.Text, nullable=False) - answer_tokens = db.Column(db.Integer, nullable=False, server_default=db.text('0')) + answer_tokens = db.Column(db.Integer, nullable=False, server_default=db.text("0")) answer_unit_price = db.Column(db.Numeric(10, 4), nullable=False) - answer_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text('0.001')) - provider_response_latency = db.Column(db.Float, nullable=False, server_default=db.text('0')) + answer_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text("0.001")) + provider_response_latency = db.Column(db.Float, nullable=False, server_default=db.text("0")) total_price = db.Column(db.Numeric(10, 7)) currency = db.Column(db.String(255), nullable=False) status = db.Column(db.String(255), nullable=False, server_default=db.text("'normal'::character varying")) @@ -678,9 +720,9 @@ class Message(db.Model): from_source = db.Column(db.String(255), nullable=False) from_end_user_id = db.Column(StringUUID) from_account_id = db.Column(StringUUID) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - agent_based = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + agent_based = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) workflow_run_id = db.Column(StringUUID) @property @@ -688,7 +730,7 @@ class Message(db.Model): if not self.answer: return self.answer - pattern = r'\[!?.*?\]\((((http|https):\/\/.+)?\/files\/(tools\/)?[\w-]+.*?timestamp=.*&nonce=.*&sign=.*)\)' + pattern = r"\[!?.*?\]\((((http|https):\/\/.+)?\/files\/(tools\/)?[\w-]+.*?timestamp=.*&nonce=.*&sign=.*)\)" matches = re.findall(pattern, self.answer) if not matches: @@ -704,9 +746,9 @@ class Message(db.Model): re_sign_file_url_answer = self.answer for url in urls: - if 'files/tools' in url: + if "files/tools" in url: # get tool file id - tool_file_id_pattern = r'\/files\/tools\/([\.\w-]+)?\?timestamp=' + tool_file_id_pattern = r"\/files\/tools\/([\.\w-]+)?\?timestamp=" result = re.search(tool_file_id_pattern, url) if not result: continue @@ -714,25 +756,24 @@ class Message(db.Model): tool_file_id = result.group(1) # get extension - if '.' in tool_file_id: - split_result = tool_file_id.split('.') - extension = f'.{split_result[-1]}' + if "." in tool_file_id: + split_result = tool_file_id.split(".") + extension = f".{split_result[-1]}" if len(extension) > 10: - extension = '.bin' + extension = ".bin" tool_file_id = split_result[0] else: - extension = '.bin' + extension = ".bin" if not tool_file_id: continue sign_url = ToolFileParser.get_tool_file_manager().sign_file( - tool_file_id=tool_file_id, - extension=extension + tool_file_id=tool_file_id, extension=extension ) else: # get upload file id - upload_file_id_pattern = r'\/files\/([\w-]+)\/image-preview?\?timestamp=' + upload_file_id_pattern = r"\/files\/([\w-]+)\/image-preview?\?timestamp=" result = re.search(upload_file_id_pattern, url) if not result: continue @@ -750,14 +791,20 @@ class Message(db.Model): @property def user_feedback(self): - feedback = db.session.query(MessageFeedback).filter(MessageFeedback.message_id == self.id, - MessageFeedback.from_source == 'user').first() + feedback = ( + db.session.query(MessageFeedback) + .filter(MessageFeedback.message_id == self.id, MessageFeedback.from_source == "user") + .first() + ) return feedback @property def admin_feedback(self): - feedback = db.session.query(MessageFeedback).filter(MessageFeedback.message_id == self.id, - MessageFeedback.from_source == 'admin').first() + feedback = ( + db.session.query(MessageFeedback) + .filter(MessageFeedback.message_id == self.id, MessageFeedback.from_source == "admin") + .first() + ) return feedback @property @@ -772,11 +819,15 @@ class Message(db.Model): @property def annotation_hit_history(self): - annotation_history = (db.session.query(AppAnnotationHitHistory) - .filter(AppAnnotationHitHistory.message_id == self.id).first()) + annotation_history = ( + db.session.query(AppAnnotationHitHistory).filter(AppAnnotationHitHistory.message_id == self.id).first() + ) if annotation_history: - annotation = (db.session.query(MessageAnnotation). - filter(MessageAnnotation.id == annotation_history.annotation_id).first()) + annotation = ( + db.session.query(MessageAnnotation) + .filter(MessageAnnotation.id == annotation_history.annotation_id) + .first() + ) return annotation return None @@ -784,8 +835,9 @@ class Message(db.Model): def app_model_config(self): conversation = db.session.query(Conversation).filter(Conversation.id == self.conversation_id).first() if conversation: - return db.session.query(AppModelConfig).filter( - AppModelConfig.id == conversation.app_model_config_id).first() + return ( + db.session.query(AppModelConfig).filter(AppModelConfig.id == conversation.app_model_config_id).first() + ) return None @@ -799,13 +851,21 @@ class Message(db.Model): @property def agent_thoughts(self): - return db.session.query(MessageAgentThought).filter(MessageAgentThought.message_id == self.id) \ - .order_by(MessageAgentThought.position.asc()).all() + return ( + db.session.query(MessageAgentThought) + .filter(MessageAgentThought.message_id == self.id) + .order_by(MessageAgentThought.position.asc()) + .all() + ) @property def retriever_resources(self): - return db.session.query(DatasetRetrieverResource).filter(DatasetRetrieverResource.message_id == self.id) \ - .order_by(DatasetRetrieverResource.position.asc()).all() + return ( + db.session.query(DatasetRetrieverResource) + .filter(DatasetRetrieverResource.message_id == self.id) + .order_by(DatasetRetrieverResource.position.asc()) + .all() + ) @property def message_files(self): @@ -818,39 +878,39 @@ class Message(db.Model): files = [] for message_file in message_files: url = message_file.url - if message_file.type == 'image': - if message_file.transfer_method == 'local_file': - upload_file = (db.session.query(UploadFile) - .filter( - UploadFile.id == message_file.upload_file_id - ).first()) - - url = UploadFileParser.get_image_data( - upload_file=upload_file, - force_url=True + if message_file.type == "image": + if message_file.transfer_method == "local_file": + upload_file = ( + db.session.query(UploadFile).filter(UploadFile.id == message_file.upload_file_id).first() ) - if message_file.transfer_method == 'tool_file': + + url = UploadFileParser.get_image_data(upload_file=upload_file, force_url=True) + if message_file.transfer_method == "tool_file": # get tool file id - tool_file_id = message_file.url.split('/')[-1] + tool_file_id = message_file.url.split("/")[-1] # trim extension - tool_file_id = tool_file_id.split('.')[0] + tool_file_id = tool_file_id.split(".")[0] # get extension - if '.' in message_file.url: + if "." in message_file.url: extension = f'.{message_file.url.split(".")[-1]}' if len(extension) > 10: - extension = '.bin' + extension = ".bin" else: - extension = '.bin' + extension = ".bin" # add sign url - url = ToolFileParser.get_tool_file_manager().sign_file(tool_file_id=tool_file_id, extension=extension) + url = ToolFileParser.get_tool_file_manager().sign_file( + tool_file_id=tool_file_id, extension=extension + ) - files.append({ - 'id': message_file.id, - 'type': message_file.type, - 'url': url, - 'belongs_to': message_file.belongs_to if message_file.belongs_to else 'user' - }) + files.append( + { + "id": message_file.id, + "type": message_file.type, + "url": url, + "belongs_to": message_file.belongs_to if message_file.belongs_to else "user", + } + ) return files @@ -858,64 +918,65 @@ class Message(db.Model): def workflow_run(self): if self.workflow_run_id: from .workflow import WorkflowRun + return db.session.query(WorkflowRun).filter(WorkflowRun.id == self.workflow_run_id).first() return None def to_dict(self) -> dict: return { - 'id': self.id, - 'app_id': self.app_id, - 'conversation_id': self.conversation_id, - 'inputs': self.inputs, - 'query': self.query, - 'message': self.message, - 'answer': self.answer, - 'status': self.status, - 'error': self.error, - 'message_metadata': self.message_metadata_dict, - 'from_source': self.from_source, - 'from_end_user_id': self.from_end_user_id, - 'from_account_id': self.from_account_id, - 'created_at': self.created_at.isoformat(), - 'updated_at': self.updated_at.isoformat(), - 'agent_based': self.agent_based, - 'workflow_run_id': self.workflow_run_id + "id": self.id, + "app_id": self.app_id, + "conversation_id": self.conversation_id, + "inputs": self.inputs, + "query": self.query, + "message": self.message, + "answer": self.answer, + "status": self.status, + "error": self.error, + "message_metadata": self.message_metadata_dict, + "from_source": self.from_source, + "from_end_user_id": self.from_end_user_id, + "from_account_id": self.from_account_id, + "created_at": self.created_at.isoformat(), + "updated_at": self.updated_at.isoformat(), + "agent_based": self.agent_based, + "workflow_run_id": self.workflow_run_id, } @classmethod def from_dict(cls, data: dict): return cls( - id=data['id'], - app_id=data['app_id'], - conversation_id=data['conversation_id'], - inputs=data['inputs'], - query=data['query'], - message=data['message'], - answer=data['answer'], - status=data['status'], - error=data['error'], - message_metadata=json.dumps(data['message_metadata']), - from_source=data['from_source'], - from_end_user_id=data['from_end_user_id'], - from_account_id=data['from_account_id'], - created_at=data['created_at'], - updated_at=data['updated_at'], - agent_based=data['agent_based'], - workflow_run_id=data['workflow_run_id'] + id=data["id"], + app_id=data["app_id"], + conversation_id=data["conversation_id"], + inputs=data["inputs"], + query=data["query"], + message=data["message"], + answer=data["answer"], + status=data["status"], + error=data["error"], + message_metadata=json.dumps(data["message_metadata"]), + from_source=data["from_source"], + from_end_user_id=data["from_end_user_id"], + from_account_id=data["from_account_id"], + created_at=data["created_at"], + updated_at=data["updated_at"], + agent_based=data["agent_based"], + workflow_run_id=data["workflow_run_id"], ) class MessageFeedback(db.Model): - __tablename__ = 'message_feedbacks' + __tablename__ = "message_feedbacks" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='message_feedback_pkey'), - db.Index('message_feedback_app_idx', 'app_id'), - db.Index('message_feedback_message_idx', 'message_id', 'from_source'), - db.Index('message_feedback_conversation_idx', 'conversation_id', 'from_source', 'rating') + db.PrimaryKeyConstraint("id", name="message_feedback_pkey"), + db.Index("message_feedback_app_idx", "app_id"), + db.Index("message_feedback_message_idx", "message_id", "from_source"), + db.Index("message_feedback_conversation_idx", "conversation_id", "from_source", "rating"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) conversation_id = db.Column(StringUUID, nullable=False) message_id = db.Column(StringUUID, nullable=False) @@ -924,8 +985,8 @@ class MessageFeedback(db.Model): from_source = db.Column(db.String(255), nullable=False) from_end_user_id = db.Column(StringUUID) from_account_id = db.Column(StringUUID) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def from_account(self): @@ -934,14 +995,14 @@ class MessageFeedback(db.Model): class MessageFile(db.Model): - __tablename__ = 'message_files' + __tablename__ = "message_files" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='message_file_pkey'), - db.Index('message_file_message_idx', 'message_id'), - db.Index('message_file_created_by_idx', 'created_by') + db.PrimaryKeyConstraint("id", name="message_file_pkey"), + db.Index("message_file_message_idx", "message_id"), + db.Index("message_file_created_by_idx", "created_by"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) message_id = db.Column(StringUUID, nullable=False) type = db.Column(db.String(255), nullable=False) transfer_method = db.Column(db.String(255), nullable=False) @@ -950,28 +1011,28 @@ class MessageFile(db.Model): upload_file_id = db.Column(StringUUID, nullable=True) created_by_role = db.Column(db.String(255), nullable=False) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class MessageAnnotation(db.Model): - __tablename__ = 'message_annotations' + __tablename__ = "message_annotations" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='message_annotation_pkey'), - db.Index('message_annotation_app_idx', 'app_id'), - db.Index('message_annotation_conversation_idx', 'conversation_id'), - db.Index('message_annotation_message_idx', 'message_id') + db.PrimaryKeyConstraint("id", name="message_annotation_pkey"), + db.Index("message_annotation_app_idx", "app_id"), + db.Index("message_annotation_conversation_idx", "conversation_id"), + db.Index("message_annotation_message_idx", "message_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) - conversation_id = db.Column(StringUUID, db.ForeignKey('conversations.id'), nullable=True) + conversation_id = db.Column(StringUUID, db.ForeignKey("conversations.id"), nullable=True) message_id = db.Column(StringUUID, nullable=True) question = db.Column(db.Text, nullable=True) content = db.Column(db.Text, nullable=False) - hit_count = db.Column(db.Integer, nullable=False, server_default=db.text('0')) + hit_count = db.Column(db.Integer, nullable=False, server_default=db.text("0")) account_id = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def account(self): @@ -985,32 +1046,35 @@ class MessageAnnotation(db.Model): class AppAnnotationHitHistory(db.Model): - __tablename__ = 'app_annotation_hit_histories' + __tablename__ = "app_annotation_hit_histories" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='app_annotation_hit_histories_pkey'), - db.Index('app_annotation_hit_histories_app_idx', 'app_id'), - db.Index('app_annotation_hit_histories_account_idx', 'account_id'), - db.Index('app_annotation_hit_histories_annotation_idx', 'annotation_id'), - db.Index('app_annotation_hit_histories_message_idx', 'message_id'), + db.PrimaryKeyConstraint("id", name="app_annotation_hit_histories_pkey"), + db.Index("app_annotation_hit_histories_app_idx", "app_id"), + db.Index("app_annotation_hit_histories_account_idx", "account_id"), + db.Index("app_annotation_hit_histories_annotation_idx", "annotation_id"), + db.Index("app_annotation_hit_histories_message_idx", "message_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) annotation_id = db.Column(StringUUID, nullable=False) source = db.Column(db.Text, nullable=False) question = db.Column(db.Text, nullable=False) account_id = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - score = db.Column(Float, nullable=False, server_default=db.text('0')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + score = db.Column(Float, nullable=False, server_default=db.text("0")) message_id = db.Column(StringUUID, nullable=False) annotation_question = db.Column(db.Text, nullable=False) annotation_content = db.Column(db.Text, nullable=False) @property def account(self): - account = (db.session.query(Account) - .join(MessageAnnotation, MessageAnnotation.account_id == Account.id) - .filter(MessageAnnotation.id == self.annotation_id).first()) + account = ( + db.session.query(Account) + .join(MessageAnnotation, MessageAnnotation.account_id == Account.id) + .filter(MessageAnnotation.id == self.annotation_id) + .first() + ) return account @property @@ -1020,89 +1084,99 @@ class AppAnnotationHitHistory(db.Model): class AppAnnotationSetting(db.Model): - __tablename__ = 'app_annotation_settings' + __tablename__ = "app_annotation_settings" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='app_annotation_settings_pkey'), - db.Index('app_annotation_settings_app_idx', 'app_id') + db.PrimaryKeyConstraint("id", name="app_annotation_settings_pkey"), + db.Index("app_annotation_settings_app_idx", "app_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) - score_threshold = db.Column(Float, nullable=False, server_default=db.text('0')) + score_threshold = db.Column(Float, nullable=False, server_default=db.text("0")) collection_binding_id = db.Column(StringUUID, nullable=False) created_user_id = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) updated_user_id = db.Column(StringUUID, nullable=False) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def created_account(self): - account = (db.session.query(Account) - .join(AppAnnotationSetting, AppAnnotationSetting.created_user_id == Account.id) - .filter(AppAnnotationSetting.id == self.annotation_id).first()) + account = ( + db.session.query(Account) + .join(AppAnnotationSetting, AppAnnotationSetting.created_user_id == Account.id) + .filter(AppAnnotationSetting.id == self.annotation_id) + .first() + ) return account @property def updated_account(self): - account = (db.session.query(Account) - .join(AppAnnotationSetting, AppAnnotationSetting.updated_user_id == Account.id) - .filter(AppAnnotationSetting.id == self.annotation_id).first()) + account = ( + db.session.query(Account) + .join(AppAnnotationSetting, AppAnnotationSetting.updated_user_id == Account.id) + .filter(AppAnnotationSetting.id == self.annotation_id) + .first() + ) return account @property def collection_binding_detail(self): from .dataset import DatasetCollectionBinding - collection_binding_detail = (db.session.query(DatasetCollectionBinding) - .filter(DatasetCollectionBinding.id == self.collection_binding_id).first()) + + collection_binding_detail = ( + db.session.query(DatasetCollectionBinding) + .filter(DatasetCollectionBinding.id == self.collection_binding_id) + .first() + ) return collection_binding_detail class OperationLog(db.Model): - __tablename__ = 'operation_logs' + __tablename__ = "operation_logs" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='operation_log_pkey'), - db.Index('operation_log_account_action_idx', 'tenant_id', 'account_id', 'action') + db.PrimaryKeyConstraint("id", name="operation_log_pkey"), + db.Index("operation_log_account_action_idx", "tenant_id", "account_id", "action"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) account_id = db.Column(StringUUID, nullable=False) action = db.Column(db.String(255), nullable=False) content = db.Column(db.JSON) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) created_ip = db.Column(db.String(255), nullable=False) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class EndUser(UserMixin, db.Model): - __tablename__ = 'end_users' + __tablename__ = "end_users" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='end_user_pkey'), - db.Index('end_user_session_id_idx', 'session_id', 'type'), - db.Index('end_user_tenant_session_id_idx', 'tenant_id', 'session_id', 'type'), + db.PrimaryKeyConstraint("id", name="end_user_pkey"), + db.Index("end_user_session_id_idx", "session_id", "type"), + db.Index("end_user_tenant_session_id_idx", "tenant_id", "session_id", "type"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) app_id = db.Column(StringUUID, nullable=True) type = db.Column(db.String(255), nullable=False) external_user_id = db.Column(db.String(255), nullable=True) name = db.Column(db.String(255)) - is_anonymous = db.Column(db.Boolean, nullable=False, server_default=db.text('true')) + is_anonymous = db.Column(db.Boolean, nullable=False, server_default=db.text("true")) session_id = db.Column(db.String(255), nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class Site(db.Model): - __tablename__ = 'sites' + __tablename__ = "sites" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='site_pkey'), - db.Index('site_app_id_idx', 'app_id'), - db.Index('site_code_idx', 'code', 'status') + db.PrimaryKeyConstraint("id", name="site_pkey"), + db.Index("site_app_id_idx", "app_id"), + db.Index("site_code_idx", "code", "status"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) title = db.Column(db.String(255), nullable=False) icon_type = db.Column(db.String(255), nullable=True) @@ -1111,20 +1185,20 @@ class Site(db.Model): description = db.Column(db.Text) default_language = db.Column(db.String(255), nullable=False) chat_color_theme = db.Column(db.String(255)) - chat_color_theme_inverted = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) + chat_color_theme_inverted = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) copyright = db.Column(db.String(255)) privacy_policy = db.Column(db.String(255)) - show_workflow_steps = db.Column(db.Boolean, nullable=False, server_default=db.text('true')) + show_workflow_steps = db.Column(db.Boolean, nullable=False, server_default=db.text("true")) use_icon_as_answer_icon = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) custom_disclaimer = db.Column(db.String(255), nullable=True) customize_domain = db.Column(db.String(255)) customize_token_strategy = db.Column(db.String(255), nullable=False) - prompt_public = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) + prompt_public = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) status = db.Column(db.String(255), nullable=False, server_default=db.text("'normal'::character varying")) created_by = db.Column(StringUUID, nullable=True) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) updated_by = db.Column(StringUUID, nullable=True) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) code = db.Column(db.String(255)) @staticmethod @@ -1138,26 +1212,25 @@ class Site(db.Model): @property def app_base_url(self): - return ( - dify_config.APP_WEB_URL if dify_config.APP_WEB_URL else request.url_root.rstrip('/')) + return dify_config.APP_WEB_URL if dify_config.APP_WEB_URL else request.url_root.rstrip("/") class ApiToken(db.Model): - __tablename__ = 'api_tokens' + __tablename__ = "api_tokens" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='api_token_pkey'), - db.Index('api_token_app_id_type_idx', 'app_id', 'type'), - db.Index('api_token_token_idx', 'token', 'type'), - db.Index('api_token_tenant_idx', 'tenant_id', 'type') + db.PrimaryKeyConstraint("id", name="api_token_pkey"), + db.Index("api_token_app_id_type_idx", "app_id", "type"), + db.Index("api_token_token_idx", "token", "type"), + db.Index("api_token_tenant_idx", "tenant_id", "type"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=True) tenant_id = db.Column(StringUUID, nullable=True) type = db.Column(db.String(16), nullable=False) token = db.Column(db.String(255), nullable=False) last_used_at = db.Column(db.DateTime, nullable=True) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @staticmethod def generate_api_key(prefix, n): @@ -1170,13 +1243,13 @@ class ApiToken(db.Model): class UploadFile(db.Model): - __tablename__ = 'upload_files' + __tablename__ = "upload_files" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='upload_file_pkey'), - db.Index('upload_file_tenant_idx', 'tenant_id') + db.PrimaryKeyConstraint("id", name="upload_file_pkey"), + db.Index("upload_file_tenant_idx", "tenant_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) storage_type = db.Column(db.String(255), nullable=False) key = db.Column(db.String(255), nullable=False) @@ -1186,38 +1259,38 @@ class UploadFile(db.Model): mime_type = db.Column(db.String(255), nullable=True) created_by_role = db.Column(db.String(255), nullable=False, server_default=db.text("'account'::character varying")) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - used = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + used = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) used_by = db.Column(StringUUID, nullable=True) used_at = db.Column(db.DateTime, nullable=True) hash = db.Column(db.String(255), nullable=True) class ApiRequest(db.Model): - __tablename__ = 'api_requests' + __tablename__ = "api_requests" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='api_request_pkey'), - db.Index('api_request_token_idx', 'tenant_id', 'api_token_id') + db.PrimaryKeyConstraint("id", name="api_request_pkey"), + db.Index("api_request_token_idx", "tenant_id", "api_token_id"), ) - id = db.Column(StringUUID, nullable=False, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) api_token_id = db.Column(StringUUID, nullable=False) path = db.Column(db.String(255), nullable=False) request = db.Column(db.Text, nullable=True) response = db.Column(db.Text, nullable=True) ip = db.Column(db.String(255), nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class MessageChain(db.Model): - __tablename__ = 'message_chains' + __tablename__ = "message_chains" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='message_chain_pkey'), - db.Index('message_chain_message_id_idx', 'message_id') + db.PrimaryKeyConstraint("id", name="message_chain_pkey"), + db.Index("message_chain_message_id_idx", "message_id"), ) - id = db.Column(StringUUID, nullable=False, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()")) message_id = db.Column(StringUUID, nullable=False) type = db.Column(db.String(255), nullable=False) input = db.Column(db.Text, nullable=True) @@ -1226,14 +1299,14 @@ class MessageChain(db.Model): class MessageAgentThought(db.Model): - __tablename__ = 'message_agent_thoughts' + __tablename__ = "message_agent_thoughts" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='message_agent_thought_pkey'), - db.Index('message_agent_thought_message_id_idx', 'message_id'), - db.Index('message_agent_thought_message_chain_id_idx', 'message_chain_id'), + db.PrimaryKeyConstraint("id", name="message_agent_thought_pkey"), + db.Index("message_agent_thought_message_id_idx", "message_id"), + db.Index("message_agent_thought_message_chain_id_idx", "message_chain_id"), ) - id = db.Column(StringUUID, nullable=False, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()")) message_id = db.Column(StringUUID, nullable=False) message_chain_id = db.Column(StringUUID, nullable=True) position = db.Column(db.Integer, nullable=False) @@ -1248,12 +1321,12 @@ class MessageAgentThought(db.Model): message = db.Column(db.Text, nullable=True) message_token = db.Column(db.Integer, nullable=True) message_unit_price = db.Column(db.Numeric, nullable=True) - message_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text('0.001')) + message_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text("0.001")) message_files = db.Column(db.Text, nullable=True) answer = db.Column(db.Text, nullable=True) answer_token = db.Column(db.Integer, nullable=True) answer_unit_price = db.Column(db.Numeric, nullable=True) - answer_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text('0.001')) + answer_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text("0.001")) tokens = db.Column(db.Integer, nullable=True) total_price = db.Column(db.Numeric, nullable=True) currency = db.Column(db.String, nullable=True) @@ -1310,9 +1383,7 @@ class MessageAgentThought(db.Model): result[tool] = {} return result else: - return { - tool: {} for tool in tools - } + return {tool: {} for tool in tools} except Exception as e: return {} @@ -1333,22 +1404,20 @@ class MessageAgentThought(db.Model): result[tool] = {} return result else: - return { - tool: {} for tool in tools - } + return {tool: {} for tool in tools} except Exception as e: if self.observation: return dict.fromkeys(tools, self.observation) class DatasetRetrieverResource(db.Model): - __tablename__ = 'dataset_retriever_resources' + __tablename__ = "dataset_retriever_resources" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='dataset_retriever_resource_pkey'), - db.Index('dataset_retriever_resource_message_id_idx', 'message_id'), + db.PrimaryKeyConstraint("id", name="dataset_retriever_resource_pkey"), + db.Index("dataset_retriever_resource_message_id_idx", "message_id"), ) - id = db.Column(StringUUID, nullable=False, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()")) message_id = db.Column(StringUUID, nullable=False) position = db.Column(db.Integer, nullable=False) dataset_id = db.Column(StringUUID, nullable=False) @@ -1369,53 +1438,53 @@ class DatasetRetrieverResource(db.Model): class Tag(db.Model): - __tablename__ = 'tags' + __tablename__ = "tags" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tag_pkey'), - db.Index('tag_type_idx', 'type'), - db.Index('tag_name_idx', 'name'), + db.PrimaryKeyConstraint("id", name="tag_pkey"), + db.Index("tag_type_idx", "type"), + db.Index("tag_name_idx", "name"), ) - TAG_TYPE_LIST = ['knowledge', 'app'] + TAG_TYPE_LIST = ["knowledge", "app"] - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=True) type = db.Column(db.String(16), nullable=False) name = db.Column(db.String(255), nullable=False) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class TagBinding(db.Model): - __tablename__ = 'tag_bindings' + __tablename__ = "tag_bindings" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tag_binding_pkey'), - db.Index('tag_bind_target_id_idx', 'target_id'), - db.Index('tag_bind_tag_id_idx', 'tag_id'), + db.PrimaryKeyConstraint("id", name="tag_binding_pkey"), + db.Index("tag_bind_target_id_idx", "target_id"), + db.Index("tag_bind_tag_id_idx", "tag_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=True) tag_id = db.Column(StringUUID, nullable=True) target_id = db.Column(StringUUID, nullable=True) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class TraceAppConfig(db.Model): - __tablename__ = 'trace_app_config' + __tablename__ = "trace_app_config" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tracing_app_config_pkey'), - db.Index('trace_app_config_app_id_idx', 'app_id'), + db.PrimaryKeyConstraint("id", name="tracing_app_config_pkey"), + db.Index("trace_app_config_app_id_idx", "app_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) tracing_provider = db.Column(db.String(255), nullable=True) tracing_config = db.Column(db.JSON, nullable=True) created_at = db.Column(db.DateTime, nullable=False, server_default=func.now()) updated_at = db.Column(db.DateTime, nullable=False, server_default=func.now(), onupdate=func.now()) - is_active = db.Column(db.Boolean, nullable=False, server_default=db.text('true')) + is_active = db.Column(db.Boolean, nullable=False, server_default=db.text("true")) @property def tracing_config_dict(self): @@ -1427,11 +1496,11 @@ class TraceAppConfig(db.Model): def to_dict(self): return { - 'id': self.id, - 'app_id': self.app_id, - 'tracing_provider': self.tracing_provider, - 'tracing_config': self.tracing_config_dict, + "id": self.id, + "app_id": self.app_id, + "tracing_provider": self.tracing_provider, + "tracing_config": self.tracing_config_dict, "is_active": self.is_active, "created_at": self.created_at.__str__() if self.created_at else None, - 'updated_at': self.updated_at.__str__() if self.updated_at else None, + "updated_at": self.updated_at.__str__() if self.updated_at else None, } diff --git a/api/models/provider.py b/api/models/provider.py index 5d92ee6eb6..ff63e03a92 100644 --- a/api/models/provider.py +++ b/api/models/provider.py @@ -6,8 +6,8 @@ from .types import StringUUID class ProviderType(Enum): - CUSTOM = 'custom' - SYSTEM = 'system' + CUSTOM = "custom" + SYSTEM = "system" @staticmethod def value_of(value): @@ -18,13 +18,13 @@ class ProviderType(Enum): class ProviderQuotaType(Enum): - PAID = 'paid' + PAID = "paid" """hosted paid quota""" - FREE = 'free' + FREE = "free" """third-party free quota""" - TRIAL = 'trial' + TRIAL = "trial" """hosted trial quota""" @staticmethod @@ -39,27 +39,30 @@ class Provider(db.Model): """ Provider model representing the API providers and their configurations. """ - __tablename__ = 'providers' + + __tablename__ = "providers" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='provider_pkey'), - db.Index('provider_tenant_id_provider_idx', 'tenant_id', 'provider_name'), - db.UniqueConstraint('tenant_id', 'provider_name', 'provider_type', 'quota_type', name='unique_provider_name_type_quota') + db.PrimaryKeyConstraint("id", name="provider_pkey"), + db.Index("provider_tenant_id_provider_idx", "tenant_id", "provider_name"), + db.UniqueConstraint( + "tenant_id", "provider_name", "provider_type", "quota_type", name="unique_provider_name_type_quota" + ), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) provider_name = db.Column(db.String(255), nullable=False) provider_type = db.Column(db.String(40), nullable=False, server_default=db.text("'custom'::character varying")) encrypted_config = db.Column(db.Text, nullable=True) - is_valid = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) + is_valid = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) last_used = db.Column(db.DateTime, nullable=True) quota_type = db.Column(db.String(40), nullable=True, server_default=db.text("''::character varying")) quota_limit = db.Column(db.BigInteger, nullable=True) quota_used = db.Column(db.BigInteger, default=0) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) def __repr__(self): return f"" @@ -67,8 +70,8 @@ class Provider(db.Model): @property def token_is_set(self): """ - Returns True if the encrypted_config is not None, indicating that the token is set. - """ + Returns True if the encrypted_config is not None, indicating that the token is set. + """ return self.encrypted_config is not None @property @@ -86,118 +89,123 @@ class ProviderModel(db.Model): """ Provider model representing the API provider_models and their configurations. """ - __tablename__ = 'provider_models' + + __tablename__ = "provider_models" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='provider_model_pkey'), - db.Index('provider_model_tenant_id_provider_idx', 'tenant_id', 'provider_name'), - db.UniqueConstraint('tenant_id', 'provider_name', 'model_name', 'model_type', name='unique_provider_model_name') + db.PrimaryKeyConstraint("id", name="provider_model_pkey"), + db.Index("provider_model_tenant_id_provider_idx", "tenant_id", "provider_name"), + db.UniqueConstraint( + "tenant_id", "provider_name", "model_name", "model_type", name="unique_provider_model_name" + ), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) provider_name = db.Column(db.String(255), nullable=False) model_name = db.Column(db.String(255), nullable=False) model_type = db.Column(db.String(40), nullable=False) encrypted_config = db.Column(db.Text, nullable=True) - is_valid = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + is_valid = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class TenantDefaultModel(db.Model): - __tablename__ = 'tenant_default_models' + __tablename__ = "tenant_default_models" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tenant_default_model_pkey'), - db.Index('tenant_default_model_tenant_id_provider_type_idx', 'tenant_id', 'provider_name', 'model_type'), + db.PrimaryKeyConstraint("id", name="tenant_default_model_pkey"), + db.Index("tenant_default_model_tenant_id_provider_type_idx", "tenant_id", "provider_name", "model_type"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) provider_name = db.Column(db.String(255), nullable=False) model_name = db.Column(db.String(255), nullable=False) model_type = db.Column(db.String(40), nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class TenantPreferredModelProvider(db.Model): - __tablename__ = 'tenant_preferred_model_providers' + __tablename__ = "tenant_preferred_model_providers" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tenant_preferred_model_provider_pkey'), - db.Index('tenant_preferred_model_provider_tenant_provider_idx', 'tenant_id', 'provider_name'), + db.PrimaryKeyConstraint("id", name="tenant_preferred_model_provider_pkey"), + db.Index("tenant_preferred_model_provider_tenant_provider_idx", "tenant_id", "provider_name"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) provider_name = db.Column(db.String(255), nullable=False) preferred_provider_type = db.Column(db.String(40), nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class ProviderOrder(db.Model): - __tablename__ = 'provider_orders' + __tablename__ = "provider_orders" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='provider_order_pkey'), - db.Index('provider_order_tenant_provider_idx', 'tenant_id', 'provider_name'), + db.PrimaryKeyConstraint("id", name="provider_order_pkey"), + db.Index("provider_order_tenant_provider_idx", "tenant_id", "provider_name"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) provider_name = db.Column(db.String(255), nullable=False) account_id = db.Column(StringUUID, nullable=False) payment_product_id = db.Column(db.String(191), nullable=False) payment_id = db.Column(db.String(191)) transaction_id = db.Column(db.String(191)) - quantity = db.Column(db.Integer, nullable=False, server_default=db.text('1')) + quantity = db.Column(db.Integer, nullable=False, server_default=db.text("1")) currency = db.Column(db.String(40)) total_amount = db.Column(db.Integer) payment_status = db.Column(db.String(40), nullable=False, server_default=db.text("'wait_pay'::character varying")) paid_at = db.Column(db.DateTime) pay_failed_at = db.Column(db.DateTime) refunded_at = db.Column(db.DateTime) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class ProviderModelSetting(db.Model): """ Provider model settings for record the model enabled status and load balancing status. """ - __tablename__ = 'provider_model_settings' + + __tablename__ = "provider_model_settings" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='provider_model_setting_pkey'), - db.Index('provider_model_setting_tenant_provider_model_idx', 'tenant_id', 'provider_name', 'model_type'), + db.PrimaryKeyConstraint("id", name="provider_model_setting_pkey"), + db.Index("provider_model_setting_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) provider_name = db.Column(db.String(255), nullable=False) model_name = db.Column(db.String(255), nullable=False) model_type = db.Column(db.String(40), nullable=False) - enabled = db.Column(db.Boolean, nullable=False, server_default=db.text('true')) - load_balancing_enabled = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + enabled = db.Column(db.Boolean, nullable=False, server_default=db.text("true")) + load_balancing_enabled = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) class LoadBalancingModelConfig(db.Model): """ Configurations for load balancing models. """ - __tablename__ = 'load_balancing_model_configs' + + __tablename__ = "load_balancing_model_configs" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='load_balancing_model_config_pkey'), - db.Index('load_balancing_model_config_tenant_provider_model_idx', 'tenant_id', 'provider_name', 'model_type'), + db.PrimaryKeyConstraint("id", name="load_balancing_model_config_pkey"), + db.Index("load_balancing_model_config_tenant_provider_model_idx", "tenant_id", "provider_name", "model_type"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) provider_name = db.Column(db.String(255), nullable=False) model_name = db.Column(db.String(255), nullable=False) model_type = db.Column(db.String(40), nullable=False) name = db.Column(db.String(255), nullable=False) encrypted_config = db.Column(db.Text, nullable=True) - enabled = db.Column(db.Boolean, nullable=False, server_default=db.text('true')) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + enabled = db.Column(db.Boolean, nullable=False, server_default=db.text("true")) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) diff --git a/api/models/source.py b/api/models/source.py index adc00028be..07695f06e6 100644 --- a/api/models/source.py +++ b/api/models/source.py @@ -8,48 +8,48 @@ from .types import StringUUID class DataSourceOauthBinding(db.Model): - __tablename__ = 'data_source_oauth_bindings' + __tablename__ = "data_source_oauth_bindings" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='source_binding_pkey'), - db.Index('source_binding_tenant_id_idx', 'tenant_id'), - db.Index('source_info_idx', "source_info", postgresql_using='gin') + db.PrimaryKeyConstraint("id", name="source_binding_pkey"), + db.Index("source_binding_tenant_id_idx", "tenant_id"), + db.Index("source_info_idx", "source_info", postgresql_using="gin"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) access_token = db.Column(db.String(255), nullable=False) provider = db.Column(db.String(255), nullable=False) source_info = db.Column(JSONB, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - disabled = db.Column(db.Boolean, nullable=True, server_default=db.text('false')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + disabled = db.Column(db.Boolean, nullable=True, server_default=db.text("false")) class DataSourceApiKeyAuthBinding(db.Model): - __tablename__ = 'data_source_api_key_auth_bindings' + __tablename__ = "data_source_api_key_auth_bindings" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='data_source_api_key_auth_binding_pkey'), - db.Index('data_source_api_key_auth_binding_tenant_id_idx', 'tenant_id'), - db.Index('data_source_api_key_auth_binding_provider_idx', 'provider'), + db.PrimaryKeyConstraint("id", name="data_source_api_key_auth_binding_pkey"), + db.Index("data_source_api_key_auth_binding_tenant_id_idx", "tenant_id"), + db.Index("data_source_api_key_auth_binding_provider_idx", "provider"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) category = db.Column(db.String(255), nullable=False) provider = db.Column(db.String(255), nullable=False) credentials = db.Column(db.Text, nullable=True) # JSON - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - disabled = db.Column(db.Boolean, nullable=True, server_default=db.text('false')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + disabled = db.Column(db.Boolean, nullable=True, server_default=db.text("false")) def to_dict(self): return { - 'id': self.id, - 'tenant_id': self.tenant_id, - 'category': self.category, - 'provider': self.provider, - 'credentials': json.loads(self.credentials), - 'created_at': self.created_at.timestamp(), - 'updated_at': self.updated_at.timestamp(), - 'disabled': self.disabled + "id": self.id, + "tenant_id": self.tenant_id, + "category": self.category, + "provider": self.provider, + "credentials": json.loads(self.credentials), + "created_at": self.created_at.timestamp(), + "updated_at": self.updated_at.timestamp(), + "disabled": self.disabled, } diff --git a/api/models/task.py b/api/models/task.py index 618d831d8e..57b147c78d 100644 --- a/api/models/task.py +++ b/api/models/task.py @@ -8,15 +8,18 @@ from extensions.ext_database import db class CeleryTask(db.Model): """Task result/status.""" - __tablename__ = 'celery_taskmeta' + __tablename__ = "celery_taskmeta" - id = db.Column(db.Integer, db.Sequence('task_id_sequence'), - primary_key=True, autoincrement=True) + id = db.Column(db.Integer, db.Sequence("task_id_sequence"), primary_key=True, autoincrement=True) task_id = db.Column(db.String(155), unique=True) status = db.Column(db.String(50), default=states.PENDING) result = db.Column(db.PickleType, nullable=True) - date_done = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc).replace(tzinfo=None), - onupdate=lambda: datetime.now(timezone.utc).replace(tzinfo=None), nullable=True) + date_done = db.Column( + db.DateTime, + default=lambda: datetime.now(timezone.utc).replace(tzinfo=None), + onupdate=lambda: datetime.now(timezone.utc).replace(tzinfo=None), + nullable=True, + ) traceback = db.Column(db.Text, nullable=True) name = db.Column(db.String(155), nullable=True) args = db.Column(db.LargeBinary, nullable=True) @@ -29,11 +32,9 @@ class CeleryTask(db.Model): class CeleryTaskSet(db.Model): """TaskSet result.""" - __tablename__ = 'celery_tasksetmeta' + __tablename__ = "celery_tasksetmeta" - id = db.Column(db.Integer, db.Sequence('taskset_id_sequence'), - autoincrement=True, primary_key=True) + id = db.Column(db.Integer, db.Sequence("taskset_id_sequence"), autoincrement=True, primary_key=True) taskset_id = db.Column(db.String(155), unique=True) result = db.Column(db.PickleType, nullable=True) - date_done = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc).replace(tzinfo=None), - nullable=True) + date_done = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc).replace(tzinfo=None), nullable=True) diff --git a/api/models/tool.py b/api/models/tool.py index 79a70c6b1f..a81bb65174 100644 --- a/api/models/tool.py +++ b/api/models/tool.py @@ -7,7 +7,7 @@ from .types import StringUUID class ToolProviderName(Enum): - SERPAPI = 'serpapi' + SERPAPI = "serpapi" @staticmethod def value_of(value): @@ -18,25 +18,25 @@ class ToolProviderName(Enum): class ToolProvider(db.Model): - __tablename__ = 'tool_providers' + __tablename__ = "tool_providers" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tool_provider_pkey'), - db.UniqueConstraint('tenant_id', 'tool_name', name='unique_tool_provider_tool_name') + db.PrimaryKeyConstraint("id", name="tool_provider_pkey"), + db.UniqueConstraint("tenant_id", "tool_name", name="unique_tool_provider_tool_name"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) tool_name = db.Column(db.String(40), nullable=False) encrypted_credentials = db.Column(db.Text, nullable=True) - is_enabled = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + is_enabled = db.Column(db.Boolean, nullable=False, server_default=db.text("false")) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def credentials_is_set(self): """ - Returns True if the encrypted_config is not None, indicating that the token is set. - """ + Returns True if the encrypted_config is not None, indicating that the token is set. + """ return self.encrypted_credentials is not None @property diff --git a/api/models/tools.py b/api/models/tools.py index 069dc5bad0..6b69a219b1 100644 --- a/api/models/tools.py +++ b/api/models/tools.py @@ -15,15 +15,16 @@ class BuiltinToolProvider(db.Model): """ This table stores the tool provider information for built-in tools for each tenant. """ - __tablename__ = 'tool_builtin_providers' + + __tablename__ = "tool_builtin_providers" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tool_builtin_provider_pkey'), + db.PrimaryKeyConstraint("id", name="tool_builtin_provider_pkey"), # one tenant can only have one tool provider with the same name - db.UniqueConstraint('tenant_id', 'provider', name='unique_builtin_tool_provider') + db.UniqueConstraint("tenant_id", "provider", name="unique_builtin_tool_provider"), ) # id of the tool provider - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) # id of the tenant tenant_id = db.Column(StringUUID, nullable=True) # who created this tool provider @@ -32,27 +33,29 @@ class BuiltinToolProvider(db.Model): provider = db.Column(db.String(40), nullable=False) # credential of the tool provider encrypted_credentials = db.Column(db.Text, nullable=True) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def credentials(self) -> dict: return json.loads(self.encrypted_credentials) + class PublishedAppTool(db.Model): """ The table stores the apps published as a tool for each person. """ - __tablename__ = 'tool_published_apps' + + __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') + db.PrimaryKeyConstraint("id", name="published_app_tool_pkey"), + db.UniqueConstraint("app_id", "user_id", name="unique_published_app_tool"), ) # id of the tool provider - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) # id of the app - app_id = db.Column(StringUUID, ForeignKey('apps.id'), nullable=False) + app_id = db.Column(StringUUID, ForeignKey("apps.id"), nullable=False) # who published this tool user_id = db.Column(StringUUID, nullable=False) # description of the tool, stored in i18n format, for human @@ -67,28 +70,30 @@ class PublishedAppTool(db.Model): tool_name = db.Column(db.String(40), nullable=False) # author author = db.Column(db.String(40), nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def description_i18n(self) -> I18nObject: return I18nObject(**json.loads(self.description)) - + @property def app(self) -> App: return db.session.query(App).filter(App.id == self.app_id).first() + class ApiToolProvider(db.Model): """ The table stores the api providers. """ - __tablename__ = 'tool_api_providers' + + __tablename__ = "tool_api_providers" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tool_api_provider_pkey'), - db.UniqueConstraint('name', 'tenant_id', name='unique_api_tool_provider') + db.PrimaryKeyConstraint("id", name="tool_api_provider_pkey"), + db.UniqueConstraint("name", "tenant_id", name="unique_api_tool_provider"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) # name of the api provider name = db.Column(db.String(40), nullable=False) # icon @@ -111,21 +116,21 @@ class ApiToolProvider(db.Model): # custom_disclaimer custom_disclaimer = db.Column(db.String(255), nullable=True) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def schema_type(self) -> ApiProviderSchemaType: return ApiProviderSchemaType.value_of(self.schema_type_str) - + @property def tools(self) -> list[ApiToolBundle]: return [ApiToolBundle(**tool) for tool in json.loads(self.tools_str)] - + @property def credentials(self) -> dict: return json.loads(self.credentials_str) - + @property def user(self) -> Account: return db.session.query(Account).filter(Account.id == self.user_id).first() @@ -134,17 +139,19 @@ class ApiToolProvider(db.Model): def tenant(self) -> Tenant: return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first() + class ToolLabelBinding(db.Model): """ The table stores the labels for tools. """ - __tablename__ = 'tool_label_bindings' + + __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'), + db.PrimaryKeyConstraint("id", name="tool_label_bind_pkey"), + db.UniqueConstraint("tool_id", "label_name", name="unique_tool_label_bind"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) # tool id tool_id = db.Column(db.String(64), nullable=False) # tool type @@ -152,28 +159,30 @@ class ToolLabelBinding(db.Model): # label name label_name = db.Column(db.String(40), nullable=False) + class WorkflowToolProvider(db.Model): """ The table stores the workflow providers. """ - __tablename__ = 'tool_workflow_providers' + + __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'), + 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"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) # name of the workflow provider name = db.Column(db.String(40), nullable=False) # label of the workflow provider - label = db.Column(db.String(255), nullable=False, server_default='') + label = db.Column(db.String(255), nullable=False, server_default="") # icon icon = db.Column(db.String(255), nullable=False) # app id of the workflow provider app_id = db.Column(StringUUID, nullable=False) # version of the workflow provider - version = db.Column(db.String(255), nullable=False, server_default='') + version = db.Column(db.String(255), nullable=False, server_default="") # who created this tool user_id = db.Column(StringUUID, nullable=False) # tenant id @@ -181,17 +190,17 @@ class WorkflowToolProvider(db.Model): # description of the provider description = db.Column(db.Text, nullable=False) # parameter configuration - parameter_configuration = db.Column(db.Text, nullable=False, server_default='[]') + parameter_configuration = db.Column(db.Text, nullable=False, server_default="[]") # privacy policy - privacy_policy = db.Column(db.String(255), nullable=True, server_default='') + privacy_policy = db.Column(db.String(255), nullable=True, server_default="") - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def schema_type(self) -> ApiProviderSchemaType: return ApiProviderSchemaType.value_of(self.schema_type_str) - + @property def user(self) -> Account: return db.session.query(Account).filter(Account.id == self.user_id).first() @@ -199,28 +208,25 @@ class WorkflowToolProvider(db.Model): @property def tenant(self) -> Tenant: return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first() - + @property def parameter_configurations(self) -> list[WorkflowToolParameterConfiguration]: - return [ - WorkflowToolParameterConfiguration(**config) - for config in json.loads(self.parameter_configuration) - ] - + return [WorkflowToolParameterConfiguration(**config) for config in json.loads(self.parameter_configuration)] + @property def app(self) -> App: return db.session.query(App).filter(App.id == self.app_id).first() + class ToolModelInvoke(db.Model): """ store the invoke logs from tool invoke """ - __tablename__ = "tool_model_invokes" - __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tool_model_invoke_pkey'), - ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + __tablename__ = "tool_model_invokes" + __table_args__ = (db.PrimaryKeyConstraint("id", name="tool_model_invoke_pkey"),) + + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) # who invoke this tool user_id = db.Column(StringUUID, nullable=False) # tenant id @@ -238,29 +244,31 @@ class ToolModelInvoke(db.Model): # invoke response model_response = db.Column(db.Text, nullable=False) - prompt_tokens = db.Column(db.Integer, nullable=False, server_default=db.text('0')) - answer_tokens = db.Column(db.Integer, nullable=False, server_default=db.text('0')) + prompt_tokens = db.Column(db.Integer, nullable=False, server_default=db.text("0")) + answer_tokens = db.Column(db.Integer, nullable=False, server_default=db.text("0")) answer_unit_price = db.Column(db.Numeric(10, 4), nullable=False) - answer_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text('0.001')) - provider_response_latency = db.Column(db.Float, nullable=False, server_default=db.text('0')) + answer_price_unit = db.Column(db.Numeric(10, 7), nullable=False, server_default=db.text("0.001")) + provider_response_latency = db.Column(db.Float, nullable=False, server_default=db.text("0")) total_price = db.Column(db.Numeric(10, 7)) currency = db.Column(db.String(255), nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + class ToolConversationVariables(db.Model): """ store the conversation variables from tool invoke """ + __tablename__ = "tool_conversation_variables" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tool_conversation_variables_pkey'), + db.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'), + db.Index("user_id_idx", "user_id"), + db.Index("conversation_id_idx", "conversation_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) # conversation user id user_id = db.Column(StringUUID, nullable=False) # tenant id @@ -270,25 +278,27 @@ class ToolConversationVariables(db.Model): # variables pool variables_str = db.Column(db.Text, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def variables(self) -> dict: return json.loads(self.variables_str) - + + class ToolFile(db.Model): """ store the file created by agent """ + __tablename__ = "tool_files" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='tool_file_pkey'), + db.PrimaryKeyConstraint("id", name="tool_file_pkey"), # add index for conversation_id - db.Index('tool_file_conversation_id_idx', 'conversation_id'), + db.Index("tool_file_conversation_id_idx", "conversation_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) # conversation user id user_id = db.Column(StringUUID, nullable=False) # tenant id @@ -300,4 +310,4 @@ class ToolFile(db.Model): # mime type mimetype = db.Column(db.String(255), nullable=False) # original url - original_url = db.Column(db.String(2048), nullable=True) \ No newline at end of file + original_url = db.Column(db.String(2048), nullable=True) diff --git a/api/models/types.py b/api/models/types.py index 1614ec2018..cb6773e70c 100644 --- a/api/models/types.py +++ b/api/models/types.py @@ -9,13 +9,13 @@ class StringUUID(TypeDecorator): def process_bind_param(self, value, dialect): if value is None: return value - elif dialect.name == 'postgresql': + elif dialect.name == "postgresql": return str(value) else: return value.hex def load_dialect_impl(self, dialect): - if dialect.name == 'postgresql': + if dialect.name == "postgresql": return dialect.type_descriptor(UUID()) else: return dialect.type_descriptor(CHAR(36)) @@ -23,4 +23,4 @@ class StringUUID(TypeDecorator): def process_result_value(self, value, dialect): if value is None: return value - return str(value) \ No newline at end of file + return str(value) diff --git a/api/models/web.py b/api/models/web.py index 0e901d5f84..bc088c185d 100644 --- a/api/models/web.py +++ b/api/models/web.py @@ -1,4 +1,3 @@ - from extensions.ext_database import db from .model import Message @@ -6,18 +5,18 @@ from .types import StringUUID class SavedMessage(db.Model): - __tablename__ = 'saved_messages' + __tablename__ = "saved_messages" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='saved_message_pkey'), - db.Index('saved_message_message_idx', 'app_id', 'message_id', 'created_by_role', 'created_by'), + db.PrimaryKeyConstraint("id", name="saved_message_pkey"), + db.Index("saved_message_message_idx", "app_id", "message_id", "created_by_role", "created_by"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) message_id = db.Column(StringUUID, nullable=False) created_by_role = db.Column(db.String(255), nullable=False, server_default=db.text("'end_user'::character varying")) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def message(self): @@ -25,15 +24,15 @@ class SavedMessage(db.Model): class PinnedConversation(db.Model): - __tablename__ = 'pinned_conversations' + __tablename__ = "pinned_conversations" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='pinned_conversation_pkey'), - db.Index('pinned_conversation_conversation_idx', 'app_id', 'conversation_id', 'created_by_role', 'created_by'), + db.PrimaryKeyConstraint("id", name="pinned_conversation_pkey"), + db.Index("pinned_conversation_conversation_idx", "app_id", "conversation_id", "created_by_role", "created_by"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) app_id = db.Column(StringUUID, nullable=False) conversation_id = db.Column(StringUUID, nullable=False) created_by_role = db.Column(db.String(255), nullable=False, server_default=db.text("'end_user'::character varying")) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) diff --git a/api/models/workflow.py b/api/models/workflow.py index e78b5666bc..d52749f0ff 100644 --- a/api/models/workflow.py +++ b/api/models/workflow.py @@ -22,11 +22,12 @@ class CreatedByRole(Enum): """ Created By Role Enum """ - ACCOUNT = 'account' - END_USER = 'end_user' + + ACCOUNT = "account" + END_USER = "end_user" @classmethod - def value_of(cls, value: str) -> 'CreatedByRole': + def value_of(cls, value: str) -> "CreatedByRole": """ Get value of given mode. @@ -36,18 +37,19 @@ class CreatedByRole(Enum): for mode in cls: if mode.value == value: return mode - raise ValueError(f'invalid created by role value {value}') + raise ValueError(f"invalid created by role value {value}") class WorkflowType(Enum): """ Workflow Type Enum """ - WORKFLOW = 'workflow' - CHAT = 'chat' + + WORKFLOW = "workflow" + CHAT = "chat" @classmethod - def value_of(cls, value: str) -> 'WorkflowType': + def value_of(cls, value: str) -> "WorkflowType": """ Get value of given mode. @@ -57,10 +59,10 @@ class WorkflowType(Enum): for mode in cls: if mode.value == value: return mode - raise ValueError(f'invalid workflow type value {value}') + raise ValueError(f"invalid workflow type value {value}") @classmethod - def from_app_mode(cls, app_mode: Union[str, 'AppMode']) -> 'WorkflowType': + def from_app_mode(cls, app_mode: Union[str, "AppMode"]) -> "WorkflowType": """ Get workflow type from app mode. @@ -68,6 +70,7 @@ class WorkflowType(Enum): :return: workflow type """ from models.model import AppMode + app_mode = app_mode if isinstance(app_mode, AppMode) else AppMode.value_of(app_mode) return cls.WORKFLOW if app_mode == AppMode.WORKFLOW else cls.CHAT @@ -105,13 +108,13 @@ class Workflow(db.Model): - updated_at (timestamp) `optional` Last update time """ - __tablename__ = 'workflows' + __tablename__ = "workflows" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='workflow_pkey'), - db.Index('workflow_version_idx', 'tenant_id', 'app_id', 'version'), + db.PrimaryKeyConstraint("id", name="workflow_pkey"), + db.Index("workflow_version_idx", "tenant_id", "app_id", "version"), ) - id: Mapped[str] = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id: Mapped[str] = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id: Mapped[str] = db.Column(StringUUID, nullable=False) app_id: Mapped[str] = db.Column(StringUUID, nullable=False) type: Mapped[str] = db.Column(db.String(255), nullable=False) @@ -119,15 +122,31 @@ class Workflow(db.Model): graph: Mapped[str] = db.Column(db.Text) features: Mapped[str] = db.Column(db.Text) created_by: Mapped[str] = db.Column(StringUUID, nullable=False) - created_at: Mapped[datetime] = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at: Mapped[datetime] = db.Column( + db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)") + ) updated_by: Mapped[str] = db.Column(StringUUID) updated_at: Mapped[datetime] = db.Column(db.DateTime) - _environment_variables: Mapped[str] = db.Column('environment_variables', db.Text, nullable=False, server_default='{}') - _conversation_variables: Mapped[str] = db.Column('conversation_variables', db.Text, nullable=False, server_default='{}') + _environment_variables: Mapped[str] = db.Column( + "environment_variables", db.Text, nullable=False, server_default="{}" + ) + _conversation_variables: Mapped[str] = db.Column( + "conversation_variables", db.Text, nullable=False, server_default="{}" + ) - def __init__(self, *, tenant_id: str, app_id: str, type: str, version: str, graph: str, - features: str, created_by: str, environment_variables: Sequence[Variable], - conversation_variables: Sequence[Variable]): + def __init__( + self, + *, + tenant_id: str, + app_id: str, + type: str, + version: str, + graph: str, + features: str, + created_by: str, + environment_variables: Sequence[Variable], + conversation_variables: Sequence[Variable], + ): self.tenant_id = tenant_id self.app_id = app_id self.type = type @@ -160,22 +179,20 @@ class Workflow(db.Model): return [] graph_dict = self.graph_dict - if 'nodes' not in graph_dict: + if "nodes" not in graph_dict: return [] - start_node = next((node for node in graph_dict['nodes'] if node['data']['type'] == 'start'), None) + start_node = next((node for node in graph_dict["nodes"] if node["data"]["type"] == "start"), None) if not start_node: return [] # get user_input_form from start node - variables = start_node.get('data', {}).get('variables', []) + variables = start_node.get("data", {}).get("variables", []) if to_old_structure: old_structure_variables = [] for variable in variables: - old_structure_variables.append({ - variable['type']: variable - }) + old_structure_variables.append({variable["type"]: variable}) return old_structure_variables @@ -188,25 +205,24 @@ class Workflow(db.Model): :return: hash """ - entity = { - 'graph': self.graph_dict, - 'features': self.features_dict - } + entity = {"graph": self.graph_dict, "features": self.features_dict} return helper.generate_text_hash(json.dumps(entity, sort_keys=True)) @property def tool_published(self) -> bool: from models.tools import WorkflowToolProvider - return db.session.query(WorkflowToolProvider).filter( - WorkflowToolProvider.app_id == self.app_id - ).first() is not None + + return ( + db.session.query(WorkflowToolProvider).filter(WorkflowToolProvider.app_id == self.app_id).first() + is not None + ) @property def environment_variables(self) -> Sequence[Variable]: # TODO: find some way to init `self._environment_variables` when instance created. if self._environment_variables is None: - self._environment_variables = '{}' + self._environment_variables = "{}" tenant_id = contexts.tenant_id.get() @@ -215,9 +231,7 @@ class Workflow(db.Model): # decrypt secret variables value decrypt_func = ( - lambda var: var.model_copy( - update={'value': encrypter.decrypt_token(tenant_id=tenant_id, token=var.value)} - ) + lambda var: var.model_copy(update={"value": encrypter.decrypt_token(tenant_id=tenant_id, token=var.value)}) if isinstance(var, SecretVariable) else var ) @@ -230,19 +244,17 @@ class Workflow(db.Model): value = list(value) if any(var for var in value if not var.id): - raise ValueError('environment variable require a unique id') + raise ValueError("environment variable require a unique id") # Compare inputs and origin variables, if the value is HIDDEN_VALUE, use the origin variable value (only update `name`). origin_variables_dictionary = {var.id: var for var in self.environment_variables} for i, variable in enumerate(value): if variable.id in origin_variables_dictionary and variable.value == HIDDEN_VALUE: - value[i] = origin_variables_dictionary[variable.id].model_copy(update={'name': variable.name}) + value[i] = origin_variables_dictionary[variable.id].model_copy(update={"name": variable.name}) # encrypt secret variables value encrypt_func = ( - lambda var: var.model_copy( - update={'value': encrypter.encrypt_token(tenant_id=tenant_id, token=var.value)} - ) + lambda var: var.model_copy(update={"value": encrypter.encrypt_token(tenant_id=tenant_id, token=var.value)}) if isinstance(var, SecretVariable) else var ) @@ -256,15 +268,15 @@ class Workflow(db.Model): def to_dict(self, *, include_secret: bool = False) -> Mapping[str, Any]: environment_variables = list(self.environment_variables) environment_variables = [ - v if not isinstance(v, SecretVariable) or include_secret else v.model_copy(update={'value': ''}) + v if not isinstance(v, SecretVariable) or include_secret else v.model_copy(update={"value": ""}) for v in environment_variables ] result = { - 'graph': self.graph_dict, - 'features': self.features_dict, - 'environment_variables': [var.model_dump(mode='json') for var in environment_variables], - 'conversation_variables': [var.model_dump(mode='json') for var in self.conversation_variables], + "graph": self.graph_dict, + "features": self.features_dict, + "environment_variables": [var.model_dump(mode="json") for var in environment_variables], + "conversation_variables": [var.model_dump(mode="json") for var in self.conversation_variables], } return result @@ -272,7 +284,7 @@ class Workflow(db.Model): def conversation_variables(self) -> Sequence[Variable]: # TODO: find some way to init `self._conversation_variables` when instance created. if self._conversation_variables is None: - self._conversation_variables = '{}' + self._conversation_variables = "{}" variables_dict: dict[str, Any] = json.loads(self._conversation_variables) results = [factory.build_variable_from_mapping(v) for v in variables_dict.values()] @@ -290,11 +302,12 @@ class WorkflowRunTriggeredFrom(Enum): """ Workflow Run Triggered From Enum """ - DEBUGGING = 'debugging' - APP_RUN = 'app-run' + + DEBUGGING = "debugging" + APP_RUN = "app-run" @classmethod - def value_of(cls, value: str) -> 'WorkflowRunTriggeredFrom': + def value_of(cls, value: str) -> "WorkflowRunTriggeredFrom": """ Get value of given mode. @@ -304,20 +317,21 @@ class WorkflowRunTriggeredFrom(Enum): for mode in cls: if mode.value == value: return mode - raise ValueError(f'invalid workflow run triggered from value {value}') + raise ValueError(f"invalid workflow run triggered from value {value}") class WorkflowRunStatus(Enum): """ Workflow Run Status Enum """ - RUNNING = 'running' - SUCCEEDED = 'succeeded' - FAILED = 'failed' - STOPPED = 'stopped' + + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + STOPPED = "stopped" @classmethod - def value_of(cls, value: str) -> 'WorkflowRunStatus': + def value_of(cls, value: str) -> "WorkflowRunStatus": """ Get value of given mode. @@ -327,7 +341,7 @@ class WorkflowRunStatus(Enum): for mode in cls: if mode.value == value: return mode - raise ValueError(f'invalid workflow run status value {value}') + raise ValueError(f"invalid workflow run status value {value}") class WorkflowRun(db.Model): @@ -368,14 +382,14 @@ class WorkflowRun(db.Model): - finished_at (timestamp) End time """ - __tablename__ = 'workflow_runs' + __tablename__ = "workflow_runs" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='workflow_run_pkey'), - db.Index('workflow_run_triggerd_from_idx', 'tenant_id', 'app_id', 'triggered_from'), - db.Index('workflow_run_tenant_app_sequence_idx', 'tenant_id', 'app_id', 'sequence_number'), + db.PrimaryKeyConstraint("id", name="workflow_run_pkey"), + db.Index("workflow_run_triggerd_from_idx", "tenant_id", "app_id", "triggered_from"), + db.Index("workflow_run_tenant_app_sequence_idx", "tenant_id", "app_id", "sequence_number"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) app_id = db.Column(StringUUID, nullable=False) sequence_number = db.Column(db.Integer, nullable=False) @@ -388,26 +402,25 @@ class WorkflowRun(db.Model): status = db.Column(db.String(255), nullable=False) outputs = db.Column(db.Text) error = db.Column(db.Text) - elapsed_time = db.Column(db.Float, nullable=False, server_default=db.text('0')) - total_tokens = db.Column(db.Integer, nullable=False, server_default=db.text('0')) - total_steps = db.Column(db.Integer, server_default=db.text('0')) + elapsed_time = db.Column(db.Float, nullable=False, server_default=db.text("0")) + total_tokens = db.Column(db.Integer, nullable=False, server_default=db.text("0")) + total_steps = db.Column(db.Integer, server_default=db.text("0")) created_by_role = db.Column(db.String(255), nullable=False) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) finished_at = db.Column(db.DateTime) @property def created_by_account(self): created_by_role = CreatedByRole.value_of(self.created_by_role) - return db.session.get(Account, self.created_by) \ - if created_by_role == CreatedByRole.ACCOUNT else None + return db.session.get(Account, self.created_by) if created_by_role == CreatedByRole.ACCOUNT else None @property def created_by_end_user(self): from models.model import EndUser + created_by_role = CreatedByRole.value_of(self.created_by_role) - return db.session.get(EndUser, self.created_by) \ - if created_by_role == CreatedByRole.END_USER else None + return db.session.get(EndUser, self.created_by) if created_by_role == CreatedByRole.END_USER else None @property def graph_dict(self): @@ -422,12 +435,12 @@ class WorkflowRun(db.Model): return json.loads(self.outputs) if self.outputs else None @property - def message(self) -> Optional['Message']: + def message(self) -> Optional["Message"]: from models.model import Message - return db.session.query(Message).filter( - Message.app_id == self.app_id, - Message.workflow_run_id == self.id - ).first() + + return ( + db.session.query(Message).filter(Message.app_id == self.app_id, Message.workflow_run_id == self.id).first() + ) @property def workflow(self): @@ -435,51 +448,51 @@ class WorkflowRun(db.Model): def to_dict(self): return { - 'id': self.id, - 'tenant_id': self.tenant_id, - 'app_id': self.app_id, - 'sequence_number': self.sequence_number, - 'workflow_id': self.workflow_id, - 'type': self.type, - 'triggered_from': self.triggered_from, - 'version': self.version, - 'graph': self.graph_dict, - 'inputs': self.inputs_dict, - 'status': self.status, - 'outputs': self.outputs_dict, - 'error': self.error, - 'elapsed_time': self.elapsed_time, - 'total_tokens': self.total_tokens, - 'total_steps': self.total_steps, - 'created_by_role': self.created_by_role, - 'created_by': self.created_by, - 'created_at': self.created_at, - 'finished_at': self.finished_at, + "id": self.id, + "tenant_id": self.tenant_id, + "app_id": self.app_id, + "sequence_number": self.sequence_number, + "workflow_id": self.workflow_id, + "type": self.type, + "triggered_from": self.triggered_from, + "version": self.version, + "graph": self.graph_dict, + "inputs": self.inputs_dict, + "status": self.status, + "outputs": self.outputs_dict, + "error": self.error, + "elapsed_time": self.elapsed_time, + "total_tokens": self.total_tokens, + "total_steps": self.total_steps, + "created_by_role": self.created_by_role, + "created_by": self.created_by, + "created_at": self.created_at, + "finished_at": self.finished_at, } @classmethod - def from_dict(cls, data: dict) -> 'WorkflowRun': + def from_dict(cls, data: dict) -> "WorkflowRun": return cls( - id=data.get('id'), - tenant_id=data.get('tenant_id'), - app_id=data.get('app_id'), - sequence_number=data.get('sequence_number'), - workflow_id=data.get('workflow_id'), - type=data.get('type'), - triggered_from=data.get('triggered_from'), - version=data.get('version'), - graph=json.dumps(data.get('graph')), - inputs=json.dumps(data.get('inputs')), - status=data.get('status'), - outputs=json.dumps(data.get('outputs')), - error=data.get('error'), - elapsed_time=data.get('elapsed_time'), - total_tokens=data.get('total_tokens'), - total_steps=data.get('total_steps'), - created_by_role=data.get('created_by_role'), - created_by=data.get('created_by'), - created_at=data.get('created_at'), - finished_at=data.get('finished_at'), + id=data.get("id"), + tenant_id=data.get("tenant_id"), + app_id=data.get("app_id"), + sequence_number=data.get("sequence_number"), + workflow_id=data.get("workflow_id"), + type=data.get("type"), + triggered_from=data.get("triggered_from"), + version=data.get("version"), + graph=json.dumps(data.get("graph")), + inputs=json.dumps(data.get("inputs")), + status=data.get("status"), + outputs=json.dumps(data.get("outputs")), + error=data.get("error"), + elapsed_time=data.get("elapsed_time"), + total_tokens=data.get("total_tokens"), + total_steps=data.get("total_steps"), + created_by_role=data.get("created_by_role"), + created_by=data.get("created_by"), + created_at=data.get("created_at"), + finished_at=data.get("finished_at"), ) @@ -487,11 +500,12 @@ class WorkflowNodeExecutionTriggeredFrom(Enum): """ Workflow Node Execution Triggered From Enum """ - SINGLE_STEP = 'single-step' - WORKFLOW_RUN = 'workflow-run' + + SINGLE_STEP = "single-step" + WORKFLOW_RUN = "workflow-run" @classmethod - def value_of(cls, value: str) -> 'WorkflowNodeExecutionTriggeredFrom': + def value_of(cls, value: str) -> "WorkflowNodeExecutionTriggeredFrom": """ Get value of given mode. @@ -501,19 +515,20 @@ class WorkflowNodeExecutionTriggeredFrom(Enum): for mode in cls: if mode.value == value: return mode - raise ValueError(f'invalid workflow node execution triggered from value {value}') + raise ValueError(f"invalid workflow node execution triggered from value {value}") class WorkflowNodeExecutionStatus(Enum): """ Workflow Node Execution Status Enum """ - RUNNING = 'running' - SUCCEEDED = 'succeeded' - FAILED = 'failed' + + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" @classmethod - def value_of(cls, value: str) -> 'WorkflowNodeExecutionStatus': + def value_of(cls, value: str) -> "WorkflowNodeExecutionStatus": """ Get value of given mode. @@ -523,7 +538,7 @@ class WorkflowNodeExecutionStatus(Enum): for mode in cls: if mode.value == value: return mode - raise ValueError(f'invalid workflow node execution status value {value}') + raise ValueError(f"invalid workflow node execution status value {value}") class WorkflowNodeExecution(db.Model): @@ -574,18 +589,31 @@ class WorkflowNodeExecution(db.Model): - finished_at (timestamp) End time """ - __tablename__ = 'workflow_node_executions' + __tablename__ = "workflow_node_executions" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='workflow_node_execution_pkey'), - db.Index('workflow_node_execution_workflow_run_idx', 'tenant_id', 'app_id', 'workflow_id', - 'triggered_from', 'workflow_run_id'), - db.Index('workflow_node_execution_node_run_idx', 'tenant_id', 'app_id', 'workflow_id', - 'triggered_from', 'node_id'), - db.Index('workflow_node_execution_id_idx', 'tenant_id', 'app_id', 'workflow_id', - 'triggered_from', 'node_execution_id'), + db.PrimaryKeyConstraint("id", name="workflow_node_execution_pkey"), + db.Index( + "workflow_node_execution_workflow_run_idx", + "tenant_id", + "app_id", + "workflow_id", + "triggered_from", + "workflow_run_id", + ), + db.Index( + "workflow_node_execution_node_run_idx", "tenant_id", "app_id", "workflow_id", "triggered_from", "node_id" + ), + db.Index( + "workflow_node_execution_id_idx", + "tenant_id", + "app_id", + "workflow_id", + "triggered_from", + "node_execution_id", + ), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) app_id = db.Column(StringUUID, nullable=False) workflow_id = db.Column(StringUUID, nullable=False) @@ -602,9 +630,9 @@ class WorkflowNodeExecution(db.Model): outputs = db.Column(db.Text) status = db.Column(db.String(255), nullable=False) error = db.Column(db.Text) - elapsed_time = db.Column(db.Float, nullable=False, server_default=db.text('0')) + elapsed_time = db.Column(db.Float, nullable=False, server_default=db.text("0")) execution_metadata = db.Column(db.Text) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) created_by_role = db.Column(db.String(255), nullable=False) created_by = db.Column(StringUUID, nullable=False) finished_at = db.Column(db.DateTime) @@ -612,15 +640,14 @@ class WorkflowNodeExecution(db.Model): @property def created_by_account(self): created_by_role = CreatedByRole.value_of(self.created_by_role) - return db.session.get(Account, self.created_by) \ - if created_by_role == CreatedByRole.ACCOUNT else None + return db.session.get(Account, self.created_by) if created_by_role == CreatedByRole.ACCOUNT else None @property def created_by_end_user(self): from models.model import EndUser + created_by_role = CreatedByRole.value_of(self.created_by_role) - return db.session.get(EndUser, self.created_by) \ - if created_by_role == CreatedByRole.END_USER else None + return db.session.get(EndUser, self.created_by) if created_by_role == CreatedByRole.END_USER else None @property def inputs_dict(self): @@ -641,15 +668,17 @@ class WorkflowNodeExecution(db.Model): @property def extras(self): from core.tools.tool_manager import ToolManager + extras = {} if self.execution_metadata_dict: from core.workflow.entities.node_entities import NodeType - if self.node_type == NodeType.TOOL.value and 'tool_info' in self.execution_metadata_dict: - tool_info = self.execution_metadata_dict['tool_info'] - extras['icon'] = ToolManager.get_tool_icon( + + if self.node_type == NodeType.TOOL.value and "tool_info" in self.execution_metadata_dict: + tool_info = self.execution_metadata_dict["tool_info"] + extras["icon"] = ToolManager.get_tool_icon( tenant_id=self.tenant_id, - provider_type=tool_info['provider_type'], - provider_id=tool_info['provider_id'] + provider_type=tool_info["provider_type"], + provider_id=tool_info["provider_id"], ) return extras @@ -659,12 +688,13 @@ class WorkflowAppLogCreatedFrom(Enum): """ Workflow App Log Created From Enum """ - SERVICE_API = 'service-api' - WEB_APP = 'web-app' - INSTALLED_APP = 'installed-app' + + SERVICE_API = "service-api" + WEB_APP = "web-app" + INSTALLED_APP = "installed-app" @classmethod - def value_of(cls, value: str) -> 'WorkflowAppLogCreatedFrom': + def value_of(cls, value: str) -> "WorkflowAppLogCreatedFrom": """ Get value of given mode. @@ -674,7 +704,7 @@ class WorkflowAppLogCreatedFrom(Enum): for mode in cls: if mode.value == value: return mode - raise ValueError(f'invalid workflow app log created from value {value}') + raise ValueError(f"invalid workflow app log created from value {value}") class WorkflowAppLog(db.Model): @@ -706,13 +736,13 @@ class WorkflowAppLog(db.Model): - created_at (timestamp) Creation time """ - __tablename__ = 'workflow_app_logs' + __tablename__ = "workflow_app_logs" __table_args__ = ( - db.PrimaryKeyConstraint('id', name='workflow_app_log_pkey'), - db.Index('workflow_app_log_app_idx', 'tenant_id', 'app_id'), + db.PrimaryKeyConstraint("id", name="workflow_app_log_pkey"), + db.Index("workflow_app_log_app_idx", "tenant_id", "app_id"), ) - id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()')) + id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) tenant_id = db.Column(StringUUID, nullable=False) app_id = db.Column(StringUUID, nullable=False) workflow_id = db.Column(StringUUID, nullable=False) @@ -720,7 +750,7 @@ class WorkflowAppLog(db.Model): created_from = db.Column(db.String(255), nullable=False) created_by_role = db.Column(db.String(255), nullable=False) created_by = db.Column(StringUUID, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, server_default=db.text('CURRENT_TIMESTAMP(0)')) + created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")) @property def workflow_run(self): @@ -729,26 +759,27 @@ class WorkflowAppLog(db.Model): @property def created_by_account(self): created_by_role = CreatedByRole.value_of(self.created_by_role) - return db.session.get(Account, self.created_by) \ - if created_by_role == CreatedByRole.ACCOUNT else None + return db.session.get(Account, self.created_by) if created_by_role == CreatedByRole.ACCOUNT else None @property def created_by_end_user(self): from models.model import EndUser + created_by_role = CreatedByRole.value_of(self.created_by_role) - return db.session.get(EndUser, self.created_by) \ - if created_by_role == CreatedByRole.END_USER else None + return db.session.get(EndUser, self.created_by) if created_by_role == CreatedByRole.END_USER else None class ConversationVariable(db.Model): - __tablename__ = 'workflow_conversation_variables' + __tablename__ = "workflow_conversation_variables" id: Mapped[str] = db.Column(StringUUID, primary_key=True) conversation_id: Mapped[str] = db.Column(StringUUID, nullable=False, primary_key=True) app_id: Mapped[str] = db.Column(StringUUID, nullable=False, index=True) data = db.Column(db.Text, nullable=False) - created_at = db.Column(db.DateTime, nullable=False, index=True, server_default=db.text('CURRENT_TIMESTAMP(0)')) - updated_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp()) + created_at = db.Column(db.DateTime, nullable=False, index=True, server_default=db.text("CURRENT_TIMESTAMP(0)")) + updated_at = db.Column( + db.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp() + ) def __init__(self, *, id: str, app_id: str, conversation_id: str, data: str) -> None: self.id = id @@ -757,7 +788,7 @@ class ConversationVariable(db.Model): self.data = data @classmethod - def from_variable(cls, *, app_id: str, conversation_id: str, variable: Variable) -> 'ConversationVariable': + def from_variable(cls, *, app_id: str, conversation_id: str, variable: Variable) -> "ConversationVariable": obj = cls( id=variable.id, app_id=app_id, diff --git a/api/pyproject.toml b/api/pyproject.toml index 16b8c32c2a..dc7e271ccf 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -68,7 +68,6 @@ ignore = [ [tool.ruff.format] exclude = [ - "models/**/*.py", "migrations/**/*", ]