AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-12 17:39:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore(api/models): apply ruff reformatting (#7600)

Browse Source 
Co-authored-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
Bowen Liang 2024-09-10 17:08:06 +08:00 committed by GitHub
parent d1605952b0
commit d109881410
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 1265 additions and 1138 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
api/models/__init__.py
View File

@ -4,7 +4,7 @@ from .model import App, AppMode, Message
from .types import StringUUID from .types import StringUUID
from .workflow import ConversationVariable, Workflow, WorkflowNodeExecutionStatus 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): class CreatedByRole(Enum):
@ -12,11 +12,11 @@ class CreatedByRole(Enum):
Enum class for createdByRole Enum class for createdByRole
""" """
ACCOUNT = 'account' ACCOUNT = "account"
END_USER = 'end_user' END_USER = "end_user"
@classmethod @classmethod
def value_of(cls, value: str) -> 'CreatedByRole': def value_of(cls, value: str) -> "CreatedByRole":
""" """
Get value of given mode. Get value of given mode.
@ -26,4 +26,4 @@ class CreatedByRole(Enum):
for role in cls: for role in cls:
if role.value == value: if role.value == value:
return role return role
raise ValueError(f'invalid createdByRole value {value}') raise ValueError(f"invalid createdByRole value {value}")

165
api/models/account.py
View File

@ -9,21 +9,18 @@ from .types import StringUUID
class AccountStatus(str, enum.Enum): class AccountStatus(str, enum.Enum):
PENDING = 'pending' PENDING = "pending"
UNINITIALIZED = 'uninitialized' UNINITIALIZED = "uninitialized"
ACTIVE = 'active' ACTIVE = "active"
BANNED = 'banned' BANNED = "banned"
CLOSED = 'closed' CLOSED = "closed"
class Account(UserMixin, db.Model): class Account(UserMixin, db.Model):
__tablename__ = 'accounts' __tablename__ = "accounts"
__table_args__ = ( __table_args__ = (db.PrimaryKeyConstraint("id", name="account_pkey"), db.Index("account_email_idx", "email"))
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) name = db.Column(db.String(255), nullable=False)
email = db.Column(db.String(255), nullable=False) email = db.Column(db.String(255), nullable=False)
password = db.Column(db.String(255), nullable=True) password = db.Column(db.String(255), nullable=True)
@ -34,11 +31,11 @@ class Account(UserMixin, db.Model):
timezone = db.Column(db.String(255)) timezone = db.Column(db.String(255))
last_login_at = db.Column(db.DateTime) last_login_at = db.Column(db.DateTime)
last_login_ip = db.Column(db.String(255)) 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")) status = db.Column(db.String(16), nullable=False, server_default=db.text("'active'::character varying"))
initialized_at = db.Column(db.DateTime) initialized_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)"))
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 @property
def is_password_set(self): def is_password_set(self):
@ -65,11 +62,13 @@ class Account(UserMixin, db.Model):
@current_tenant_id.setter @current_tenant_id.setter
def current_tenant_id(self, value: str): def current_tenant_id(self, value: str):
try: try:
tenant_account_join = db.session.query(Tenant, TenantAccountJoin) \ tenant_account_join = (
.filter(Tenant.id == value) \ db.session.query(Tenant, TenantAccountJoin)
.filter(TenantAccountJoin.tenant_id == Tenant.id) \ .filter(Tenant.id == value)
.filter(TenantAccountJoin.account_id == self.id) \ .filter(TenantAccountJoin.tenant_id == Tenant.id)
.filter(TenantAccountJoin.account_id == self.id)
.one_or_none() .one_or_none()
)
if tenant_account_join: if tenant_account_join:
tenant, ta = tenant_account_join tenant, ta = tenant_account_join
@ -91,20 +90,18 @@ class Account(UserMixin, db.Model):
@classmethod @classmethod
def get_by_openid(cls, provider: str, open_id: str) -> db.Model: def get_by_openid(cls, provider: str, open_id: str) -> db.Model:
account_integrate = db.session.query(AccountIntegrate). \ account_integrate = (
filter(AccountIntegrate.provider == provider, AccountIntegrate.open_id == open_id). \ db.session.query(AccountIntegrate)
one_or_none() .filter(AccountIntegrate.provider == provider, AccountIntegrate.open_id == open_id)
.one_or_none()
)
if account_integrate: if account_integrate:
return db.session.query(Account). \ return db.session.query(Account).filter(Account.id == account_integrate.account_id).one_or_none()
filter(Account.id == account_integrate.account_id). \
one_or_none()
return None return None
def get_integrates(self) -> list[db.Model]: def get_integrates(self) -> list[db.Model]:
ai = db.Model ai = db.Model
return db.session.query(ai).filter( return db.session.query(ai).filter(ai.account_id == self.id).all()
ai.account_id == self.id
).all()
# check current_user.current_tenant.current_role in ['admin', 'owner'] # check current_user.current_tenant.current_role in ['admin', 'owner']
@property @property
@ -123,22 +120,28 @@ class Account(UserMixin, db.Model):
def is_dataset_operator(self): def is_dataset_operator(self):
return self._current_tenant.current_role == TenantAccountRole.DATASET_OPERATOR return self._current_tenant.current_role == TenantAccountRole.DATASET_OPERATOR
class TenantStatus(str, enum.Enum): class TenantStatus(str, enum.Enum):
NORMAL = 'normal' NORMAL = "normal"
ARCHIVE = 'archive' ARCHIVE = "archive"
class TenantAccountRole(str, enum.Enum): class TenantAccountRole(str, enum.Enum):
OWNER = 'owner' OWNER = "owner"
ADMIN = 'admin' ADMIN = "admin"
EDITOR = 'editor' EDITOR = "editor"
NORMAL = 'normal' NORMAL = "normal"
DATASET_OPERATOR = 'dataset_operator' DATASET_OPERATOR = "dataset_operator"
@staticmethod @staticmethod
def is_valid_role(role: str) -> bool: def is_valid_role(role: str) -> bool:
return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, return role and role in {
TenantAccountRole.NORMAL, TenantAccountRole.DATASET_OPERATOR} TenantAccountRole.OWNER,
TenantAccountRole.ADMIN,
TenantAccountRole.EDITOR,
TenantAccountRole.NORMAL,
TenantAccountRole.DATASET_OPERATOR,
}
@staticmethod @staticmethod
def is_privileged_role(role: str) -> bool: def is_privileged_role(role: str) -> bool:
@ -146,8 +149,12 @@ class TenantAccountRole(str, enum.Enum):
@staticmethod @staticmethod
def is_non_owner_role(role: str) -> bool: def is_non_owner_role(role: str) -> bool:
return role and role in {TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, TenantAccountRole.NORMAL, return role and role in {
TenantAccountRole.DATASET_OPERATOR} TenantAccountRole.ADMIN,
TenantAccountRole.EDITOR,
TenantAccountRole.NORMAL,
TenantAccountRole.DATASET_OPERATOR,
}
@staticmethod @staticmethod
def is_editing_role(role: str) -> bool: def is_editing_role(role: str) -> bool:
@ -155,29 +162,33 @@ class TenantAccountRole(str, enum.Enum):
@staticmethod @staticmethod
def is_dataset_edit_role(role: str) -> bool: def is_dataset_edit_role(role: str) -> bool:
return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR, return role and role in {
TenantAccountRole.DATASET_OPERATOR} TenantAccountRole.OWNER,
TenantAccountRole.ADMIN,
TenantAccountRole.EDITOR,
TenantAccountRole.DATASET_OPERATOR,
}
class Tenant(db.Model): class Tenant(db.Model):
__tablename__ = 'tenants' __tablename__ = "tenants"
__table_args__ = ( __table_args__ = (db.PrimaryKeyConstraint("id", name="tenant_pkey"),)
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) name = db.Column(db.String(255), nullable=False)
encrypt_public_key = db.Column(db.Text) encrypt_public_key = db.Column(db.Text)
plan = db.Column(db.String(255), nullable=False, server_default=db.text("'basic'::character varying")) 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")) status = db.Column(db.String(255), nullable=False, server_default=db.text("'normal'::character varying"))
custom_config = db.Column(db.Text) custom_config = 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)"))
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)"))
def get_accounts(self) -> list[Account]: def get_accounts(self) -> list[Account]:
return db.session.query(Account).filter( return (
Account.id == TenantAccountJoin.account_id, db.session.query(Account)
TenantAccountJoin.tenant_id == self.id .filter(Account.id == TenantAccountJoin.account_id, TenantAccountJoin.tenant_id == self.id)
).all() .all()
)
@property @property
def custom_config_dict(self) -> dict: def custom_config_dict(self) -> dict:
@ -189,54 +200,54 @@ class Tenant(db.Model):
class TenantAccountJoinRole(enum.Enum): class TenantAccountJoinRole(enum.Enum):
OWNER = 'owner' OWNER = "owner"
ADMIN = 'admin' ADMIN = "admin"
NORMAL = 'normal' NORMAL = "normal"
DATASET_OPERATOR = 'dataset_operator' DATASET_OPERATOR = "dataset_operator"
class TenantAccountJoin(db.Model): class TenantAccountJoin(db.Model):
__tablename__ = 'tenant_account_joins' __tablename__ = "tenant_account_joins"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='tenant_account_join_pkey'), db.PrimaryKeyConstraint("id", name="tenant_account_join_pkey"),
db.Index('tenant_account_join_account_id_idx', 'account_id'), db.Index("tenant_account_join_account_id_idx", "account_id"),
db.Index('tenant_account_join_tenant_id_idx', 'tenant_id'), db.Index("tenant_account_join_tenant_id_idx", "tenant_id"),
db.UniqueConstraint('tenant_id', 'account_id', name='unique_tenant_account_join') 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) tenant_id = db.Column(StringUUID, nullable=False)
account_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')) current = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))
role = db.Column(db.String(16), nullable=False, server_default='normal') role = db.Column(db.String(16), nullable=False, server_default="normal")
invited_by = db.Column(StringUUID, nullable=True) invited_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_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): class AccountIntegrate(db.Model):
__tablename__ = 'account_integrates' __tablename__ = "account_integrates"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='account_integrate_pkey'), db.PrimaryKeyConstraint("id", name="account_integrate_pkey"),
db.UniqueConstraint('account_id', 'provider', name='unique_account_provider'), db.UniqueConstraint("account_id", "provider", name="unique_account_provider"),
db.UniqueConstraint('provider', 'open_id', name='unique_provider_open_id') 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) account_id = db.Column(StringUUID, nullable=False)
provider = db.Column(db.String(16), nullable=False) provider = db.Column(db.String(16), nullable=False)
open_id = db.Column(db.String(255), nullable=False) open_id = db.Column(db.String(255), nullable=False)
encrypted_token = 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)')) 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)')) updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
class InvitationCode(db.Model): class InvitationCode(db.Model):
__tablename__ = 'invitation_codes' __tablename__ = "invitation_codes"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='invitation_code_pkey'), db.PrimaryKeyConstraint("id", name="invitation_code_pkey"),
db.Index('invitation_codes_batch_idx', 'batch'), db.Index("invitation_codes_batch_idx", "batch"),
db.Index('invitation_codes_code_idx', 'code', 'status') db.Index("invitation_codes_code_idx", "code", "status"),
) )
id = db.Column(db.Integer, nullable=False) id = db.Column(db.Integer, nullable=False)
@ -247,4 +258,4 @@ class InvitationCode(db.Model):
used_by_tenant_id = db.Column(StringUUID) used_by_tenant_id = db.Column(StringUUID)
used_by_account_id = db.Column(StringUUID) used_by_account_id = db.Column(StringUUID)
deprecated_at = db.Column(db.DateTime) 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)"))

18
api/models/api_based_extension.py
View File

@ -6,22 +6,22 @@ from .types import StringUUID
class APIBasedExtensionPoint(enum.Enum): class APIBasedExtensionPoint(enum.Enum):
APP_EXTERNAL_DATA_TOOL_QUERY = 'app.external_data_tool.query' APP_EXTERNAL_DATA_TOOL_QUERY = "app.external_data_tool.query"
PING = 'ping' PING = "ping"
APP_MODERATION_INPUT = 'app.moderation.input' APP_MODERATION_INPUT = "app.moderation.input"
APP_MODERATION_OUTPUT = 'app.moderation.output' APP_MODERATION_OUTPUT = "app.moderation.output"
class APIBasedExtension(db.Model): class APIBasedExtension(db.Model):
__tablename__ = 'api_based_extensions' __tablename__ = "api_based_extensions"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='api_based_extension_pkey'), db.PrimaryKeyConstraint("id", name="api_based_extension_pkey"),
db.Index('api_based_extension_tenant_idx', 'tenant_id'), 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) tenant_id = db.Column(StringUUID, nullable=False)
name = db.Column(db.String(255), nullable=False) name = db.Column(db.String(255), nullable=False)
api_endpoint = db.Column(db.String(255), nullable=False) api_endpoint = db.Column(db.String(255), nullable=False)
api_key = db.Column(db.Text, 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)"))

569
api/models/dataset.py
View File

@ -24,37 +24,34 @@ from .types import StringUUID
class DatasetPermissionEnum(str, enum.Enum): class DatasetPermissionEnum(str, enum.Enum):
ONLY_ME = 'only_me' ONLY_ME = "only_me"
ALL_TEAM = 'all_team_members' ALL_TEAM = "all_team_members"
PARTIAL_TEAM = 'partial_members' PARTIAL_TEAM = "partial_members"
class Dataset(db.Model): class Dataset(db.Model):
__tablename__ = 'datasets' __tablename__ = "datasets"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='dataset_pkey'), db.PrimaryKeyConstraint("id", name="dataset_pkey"),
db.Index('dataset_tenant_idx', 'tenant_id'), db.Index("dataset_tenant_idx", "tenant_id"),
db.Index('retrieval_model_idx', "retrieval_model", postgresql_using='gin') 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) tenant_id = db.Column(StringUUID, nullable=False)
name = db.Column(db.String(255), nullable=False) name = db.Column(db.String(255), nullable=False)
description = db.Column(db.Text, nullable=True) description = db.Column(db.Text, nullable=True)
provider = db.Column(db.String(255), nullable=False, provider = db.Column(db.String(255), nullable=False, server_default=db.text("'vendor'::character varying"))
server_default=db.text("'vendor'::character varying")) permission = db.Column(db.String(255), nullable=False, server_default=db.text("'only_me'::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)) data_source_type = db.Column(db.String(255))
indexing_technique = db.Column(db.String(255), nullable=True) indexing_technique = db.Column(db.String(255), nullable=True)
index_struct = db.Column(db.Text, nullable=True) index_struct = db.Column(db.Text, nullable=True)
created_by = db.Column(StringUUID, nullable=False) created_by = db.Column(StringUUID, nullable=False)
created_at = db.Column(db.DateTime, nullable=False, created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
server_default=db.text('CURRENT_TIMESTAMP(0)'))
updated_by = db.Column(StringUUID, nullable=True) updated_by = db.Column(StringUUID, nullable=True)
updated_at = db.Column(db.DateTime, nullable=False, updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
server_default=db.text('CURRENT_TIMESTAMP(0)'))
embedding_model = db.Column(db.String(255), nullable=True) embedding_model = db.Column(db.String(255), nullable=True)
embedding_model_provider = 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) collection_binding_id = db.Column(StringUUID, nullable=True)
@ -62,8 +59,9 @@ class Dataset(db.Model):
@property @property
def dataset_keyword_table(self): def dataset_keyword_table(self):
dataset_keyword_table = db.session.query(DatasetKeywordTable).filter( dataset_keyword_table = (
DatasetKeywordTable.dataset_id == self.id).first() db.session.query(DatasetKeywordTable).filter(DatasetKeywordTable.dataset_id == self.id).first()
)
if dataset_keyword_table: if dataset_keyword_table:
return dataset_keyword_table return dataset_keyword_table
@ -79,13 +77,19 @@ class Dataset(db.Model):
@property @property
def latest_process_rule(self): def latest_process_rule(self):
return DatasetProcessRule.query.filter(DatasetProcessRule.dataset_id == self.id) \ return (
.order_by(DatasetProcessRule.created_at.desc()).first() DatasetProcessRule.query.filter(DatasetProcessRule.dataset_id == self.id)
.order_by(DatasetProcessRule.created_at.desc())
.first()
)
@property @property
def app_count(self): def app_count(self):
return db.session.query(func.count(AppDatasetJoin.id)).filter(AppDatasetJoin.dataset_id == self.id, return (
App.id == AppDatasetJoin.app_id).scalar() db.session.query(func.count(AppDatasetJoin.id))
.filter(AppDatasetJoin.dataset_id == self.id, App.id == AppDatasetJoin.app_id)
.scalar()
)
@property @property
def document_count(self): def document_count(self):
@ -93,30 +97,40 @@ class Dataset(db.Model):
@property @property
def available_document_count(self): def available_document_count(self):
return db.session.query(func.count(Document.id)).filter( return (
db.session.query(func.count(Document.id))
.filter(
Document.dataset_id == self.id, Document.dataset_id == self.id,
Document.indexing_status == 'completed', Document.indexing_status == "completed",
Document.enabled == True, Document.enabled == True,
Document.archived == False Document.archived == False,
).scalar() )
.scalar()
)
@property @property
def available_segment_count(self): def available_segment_count(self):
return db.session.query(func.count(DocumentSegment.id)).filter( return (
db.session.query(func.count(DocumentSegment.id))
.filter(
DocumentSegment.dataset_id == self.id, DocumentSegment.dataset_id == self.id,
DocumentSegment.status == 'completed', DocumentSegment.status == "completed",
DocumentSegment.enabled == True DocumentSegment.enabled == True,
).scalar() )
.scalar()
)
@property @property
def word_count(self): def word_count(self):
return Document.query.with_entities(func.coalesce(func.sum(Document.word_count))) \ return (
.filter(Document.dataset_id == self.id).scalar() Document.query.with_entities(func.coalesce(func.sum(Document.word_count)))
.filter(Document.dataset_id == self.id)
.scalar()
)
@property @property
def doc_form(self): def doc_form(self):
document = db.session.query(Document).filter( document = db.session.query(Document).filter(Document.dataset_id == self.id).first()
Document.dataset_id == self.id).first()
if document: if document:
return document.doc_form return document.doc_form
return None return None
@ -124,76 +138,68 @@ class Dataset(db.Model):
@property @property
def retrieval_model_dict(self): def retrieval_model_dict(self):
default_retrieval_model = { default_retrieval_model = {
'search_method': RetrievalMethod.SEMANTIC_SEARCH.value, "search_method": RetrievalMethod.SEMANTIC_SEARCH.value,
'reranking_enable': False, "reranking_enable": False,
'reranking_model': { "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},
'reranking_provider_name': '', "top_k": 2,
'reranking_model_name': '' "score_threshold_enabled": False,
},
'top_k': 2,
'score_threshold_enabled': False
} }
return self.retrieval_model if self.retrieval_model else default_retrieval_model return self.retrieval_model if self.retrieval_model else default_retrieval_model
@property @property
def tags(self): def tags(self):
tags = db.session.query(Tag).join( tags = (
TagBinding, db.session.query(Tag)
Tag.id == TagBinding.tag_id .join(TagBinding, Tag.id == TagBinding.tag_id)
).filter( .filter(
TagBinding.target_id == self.id, TagBinding.target_id == self.id,
TagBinding.tenant_id == self.tenant_id, TagBinding.tenant_id == self.tenant_id,
Tag.tenant_id == self.tenant_id, Tag.tenant_id == self.tenant_id,
Tag.type == 'knowledge' Tag.type == "knowledge",
).all() )
.all()
)
return tags if tags else [] return tags if tags else []
@staticmethod @staticmethod
def gen_collection_name_by_id(dataset_id: str) -> str: def gen_collection_name_by_id(dataset_id: str) -> str:
normalized_dataset_id = dataset_id.replace("-", "_") 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): class DatasetProcessRule(db.Model):
__tablename__ = 'dataset_process_rules' __tablename__ = "dataset_process_rules"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='dataset_process_rule_pkey'), db.PrimaryKeyConstraint("id", name="dataset_process_rule_pkey"),
db.Index('dataset_process_rule_dataset_id_idx', 'dataset_id'), db.Index("dataset_process_rule_dataset_id_idx", "dataset_id"),
) )
id = db.Column(StringUUID, nullable=False, id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
server_default=db.text('uuid_generate_v4()'))
dataset_id = db.Column(StringUUID, nullable=False) dataset_id = db.Column(StringUUID, nullable=False)
mode = db.Column(db.String(255), nullable=False, mode = db.Column(db.String(255), nullable=False, server_default=db.text("'automatic'::character varying"))
server_default=db.text("'automatic'::character varying"))
rules = db.Column(db.Text, nullable=True) rules = db.Column(db.Text, nullable=True)
created_by = db.Column(StringUUID, nullable=False) created_by = db.Column(StringUUID, nullable=False)
created_at = db.Column(db.DateTime, nullable=False, created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
server_default=db.text('CURRENT_TIMESTAMP(0)'))
MODES = ['automatic', 'custom'] MODES = ["automatic", "custom"]
PRE_PROCESSING_RULES = ['remove_stopwords', 'remove_extra_spaces', 'remove_urls_emails'] PRE_PROCESSING_RULES = ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"]
AUTOMATIC_RULES = { AUTOMATIC_RULES = {
'pre_processing_rules': [ "pre_processing_rules": [
{'id': 'remove_extra_spaces', 'enabled': True}, {"id": "remove_extra_spaces", "enabled": True},
{'id': 'remove_urls_emails', 'enabled': False} {"id": "remove_urls_emails", "enabled": False},
], ],
'segmentation': { "segmentation": {"delimiter": "\n", "max_tokens": 500, "chunk_overlap": 50},
'delimiter': '\n',
'max_tokens': 500,
'chunk_overlap': 50
}
} }
def to_dict(self): def to_dict(self):
return { return {
'id': self.id, "id": self.id,
'dataset_id': self.dataset_id, "dataset_id": self.dataset_id,
'mode': self.mode, "mode": self.mode,
'rules': self.rules_dict, "rules": self.rules_dict,
'created_by': self.created_by, "created_by": self.created_by,
'created_at': self.created_at, "created_at": self.created_at,
} }
@property @property
@ -205,17 +211,16 @@ class DatasetProcessRule(db.Model):
class Document(db.Model): class Document(db.Model):
__tablename__ = 'documents' __tablename__ = "documents"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='document_pkey'), db.PrimaryKeyConstraint("id", name="document_pkey"),
db.Index('document_dataset_id_idx', 'dataset_id'), db.Index("document_dataset_id_idx", "dataset_id"),
db.Index('document_is_paused_idx', 'is_paused'), db.Index("document_is_paused_idx", "is_paused"),
db.Index('document_tenant_idx', 'tenant_id'), db.Index("document_tenant_idx", "tenant_id"),
) )
# initial fields # initial fields
id = db.Column(StringUUID, nullable=False, id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
server_default=db.text('uuid_generate_v4()'))
tenant_id = db.Column(StringUUID, nullable=False) tenant_id = db.Column(StringUUID, nullable=False)
dataset_id = db.Column(StringUUID, nullable=False) dataset_id = db.Column(StringUUID, nullable=False)
position = db.Column(db.Integer, 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_from = db.Column(db.String(255), nullable=False)
created_by = db.Column(StringUUID, nullable=False) created_by = db.Column(StringUUID, nullable=False)
created_api_request_id = db.Column(StringUUID, nullable=True) created_api_request_id = db.Column(StringUUID, nullable=True)
created_at = db.Column(db.DateTime, nullable=False, created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
server_default=db.text('CURRENT_TIMESTAMP(0)'))
# start processing # start processing
processing_started_at = db.Column(db.DateTime, nullable=True) 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) completed_at = db.Column(db.DateTime, nullable=True)
# pause # 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_by = db.Column(StringUUID, nullable=True)
paused_at = db.Column(db.DateTime, 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) stopped_at = db.Column(db.DateTime, nullable=True)
# basic fields # basic fields
indexing_status = db.Column(db.String( indexing_status = db.Column(db.String(255), nullable=False, server_default=db.text("'waiting'::character varying"))
255), nullable=False, server_default=db.text("'waiting'::character varying")) 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_at = db.Column(db.DateTime, nullable=True)
disabled_by = db.Column(StringUUID, nullable=True) disabled_by = db.Column(StringUUID, nullable=True)
archived = db.Column(db.Boolean, nullable=False, archived = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))
server_default=db.text('false'))
archived_reason = db.Column(db.String(255), nullable=True) archived_reason = db.Column(db.String(255), nullable=True)
archived_by = db.Column(StringUUID, nullable=True) archived_by = db.Column(StringUUID, nullable=True)
archived_at = db.Column(db.DateTime, nullable=True) archived_at = db.Column(db.DateTime, nullable=True)
updated_at = db.Column(db.DateTime, nullable=False, updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
server_default=db.text('CURRENT_TIMESTAMP(0)'))
doc_type = db.Column(db.String(40), nullable=True) doc_type = db.Column(db.String(40), nullable=True)
doc_metadata = db.Column(db.JSON, nullable=True) doc_metadata = db.Column(db.JSON, nullable=True)
doc_form = db.Column(db.String( doc_form = db.Column(db.String(255), nullable=False, server_default=db.text("'text_model'::character varying"))
255), nullable=False, server_default=db.text("'text_model'::character varying"))
doc_language = db.Column(db.String(255), nullable=True) 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 @property
def display_status(self): def display_status(self):
status = None status = None
if self.indexing_status == 'waiting': if self.indexing_status == "waiting":
status = 'queuing' status = "queuing"
elif self.indexing_status not in ['completed', 'error', 'waiting'] and self.is_paused: elif self.indexing_status not in ["completed", "error", "waiting"] and self.is_paused:
status = 'paused' status = "paused"
elif self.indexing_status in ['parsing', 'cleaning', 'splitting', 'indexing']: elif self.indexing_status in ["parsing", "cleaning", "splitting", "indexing"]:
status = 'indexing' status = "indexing"
elif self.indexing_status == 'error': elif self.indexing_status == "error":
status = 'error' status = "error"
elif self.indexing_status == 'completed' and not self.archived and self.enabled: elif self.indexing_status == "completed" and not self.archived and self.enabled:
status = 'available' status = "available"
elif self.indexing_status == 'completed' and not self.archived and not self.enabled: elif self.indexing_status == "completed" and not self.archived and not self.enabled:
status = 'disabled' status = "disabled"
elif self.indexing_status == 'completed' and self.archived: elif self.indexing_status == "completed" and self.archived:
status = 'archived' status = "archived"
return status return status
@property @property
@ -313,24 +312,26 @@ class Document(db.Model):
@property @property
def data_source_detail_dict(self): def data_source_detail_dict(self):
if self.data_source_info: 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) data_source_info_dict = json.loads(self.data_source_info)
file_detail = db.session.query(UploadFile). \ file_detail = (
filter(UploadFile.id == data_source_info_dict['upload_file_id']). \ db.session.query(UploadFile)
one_or_none() .filter(UploadFile.id == data_source_info_dict["upload_file_id"])
.one_or_none()
)
if file_detail: if file_detail:
return { return {
'upload_file': { "upload_file": {
'id': file_detail.id, "id": file_detail.id,
'name': file_detail.name, "name": file_detail.name,
'size': file_detail.size, "size": file_detail.size,
'extension': file_detail.extension, "extension": file_detail.extension,
'mime_type': file_detail.mime_type, "mime_type": file_detail.mime_type,
'created_by': file_detail.created_by, "created_by": file_detail.created_by,
'created_at': file_detail.created_at.timestamp() "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 json.loads(self.data_source_info)
return {} return {}
@ -356,120 +357,123 @@ class Document(db.Model):
@property @property
def hit_count(self): def hit_count(self):
return DocumentSegment.query.with_entities(func.coalesce(func.sum(DocumentSegment.hit_count))) \ return (
.filter(DocumentSegment.document_id == self.id).scalar() DocumentSegment.query.with_entities(func.coalesce(func.sum(DocumentSegment.hit_count)))
.filter(DocumentSegment.document_id == self.id)
.scalar()
)
def to_dict(self): def to_dict(self):
return { return {
'id': self.id, "id": self.id,
'tenant_id': self.tenant_id, "tenant_id": self.tenant_id,
'dataset_id': self.dataset_id, "dataset_id": self.dataset_id,
'position': self.position, "position": self.position,
'data_source_type': self.data_source_type, "data_source_type": self.data_source_type,
'data_source_info': self.data_source_info, "data_source_info": self.data_source_info,
'dataset_process_rule_id': self.dataset_process_rule_id, "dataset_process_rule_id": self.dataset_process_rule_id,
'batch': self.batch, "batch": self.batch,
'name': self.name, "name": self.name,
'created_from': self.created_from, "created_from": self.created_from,
'created_by': self.created_by, "created_by": self.created_by,
'created_api_request_id': self.created_api_request_id, "created_api_request_id": self.created_api_request_id,
'created_at': self.created_at, "created_at": self.created_at,
'processing_started_at': self.processing_started_at, "processing_started_at": self.processing_started_at,
'file_id': self.file_id, "file_id": self.file_id,
'word_count': self.word_count, "word_count": self.word_count,
'parsing_completed_at': self.parsing_completed_at, "parsing_completed_at": self.parsing_completed_at,
'cleaning_completed_at': self.cleaning_completed_at, "cleaning_completed_at": self.cleaning_completed_at,
'splitting_completed_at': self.splitting_completed_at, "splitting_completed_at": self.splitting_completed_at,
'tokens': self.tokens, "tokens": self.tokens,
'indexing_latency': self.indexing_latency, "indexing_latency": self.indexing_latency,
'completed_at': self.completed_at, "completed_at": self.completed_at,
'is_paused': self.is_paused, "is_paused": self.is_paused,
'paused_by': self.paused_by, "paused_by": self.paused_by,
'paused_at': self.paused_at, "paused_at": self.paused_at,
'error': self.error, "error": self.error,
'stopped_at': self.stopped_at, "stopped_at": self.stopped_at,
'indexing_status': self.indexing_status, "indexing_status": self.indexing_status,
'enabled': self.enabled, "enabled": self.enabled,
'disabled_at': self.disabled_at, "disabled_at": self.disabled_at,
'disabled_by': self.disabled_by, "disabled_by": self.disabled_by,
'archived': self.archived, "archived": self.archived,
'archived_reason': self.archived_reason, "archived_reason": self.archived_reason,
'archived_by': self.archived_by, "archived_by": self.archived_by,
'archived_at': self.archived_at, "archived_at": self.archived_at,
'updated_at': self.updated_at, "updated_at": self.updated_at,
'doc_type': self.doc_type, "doc_type": self.doc_type,
'doc_metadata': self.doc_metadata, "doc_metadata": self.doc_metadata,
'doc_form': self.doc_form, "doc_form": self.doc_form,
'doc_language': self.doc_language, "doc_language": self.doc_language,
'display_status': self.display_status, "display_status": self.display_status,
'data_source_info_dict': self.data_source_info_dict, "data_source_info_dict": self.data_source_info_dict,
'average_segment_length': self.average_segment_length, "average_segment_length": self.average_segment_length,
'dataset_process_rule': self.dataset_process_rule.to_dict() if self.dataset_process_rule else None, "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, "dataset": self.dataset.to_dict() if self.dataset else None,
'segment_count': self.segment_count, "segment_count": self.segment_count,
'hit_count': self.hit_count "hit_count": self.hit_count,
} }
@classmethod @classmethod
def from_dict(cls, data: dict): def from_dict(cls, data: dict):
return cls( return cls(
id=data.get('id'), id=data.get("id"),
tenant_id=data.get('tenant_id'), tenant_id=data.get("tenant_id"),
dataset_id=data.get('dataset_id'), dataset_id=data.get("dataset_id"),
position=data.get('position'), position=data.get("position"),
data_source_type=data.get('data_source_type'), data_source_type=data.get("data_source_type"),
data_source_info=data.get('data_source_info'), data_source_info=data.get("data_source_info"),
dataset_process_rule_id=data.get('dataset_process_rule_id'), dataset_process_rule_id=data.get("dataset_process_rule_id"),
batch=data.get('batch'), batch=data.get("batch"),
name=data.get('name'), name=data.get("name"),
created_from=data.get('created_from'), created_from=data.get("created_from"),
created_by=data.get('created_by'), created_by=data.get("created_by"),
created_api_request_id=data.get('created_api_request_id'), created_api_request_id=data.get("created_api_request_id"),
created_at=data.get('created_at'), created_at=data.get("created_at"),
processing_started_at=data.get('processing_started_at'), processing_started_at=data.get("processing_started_at"),
file_id=data.get('file_id'), file_id=data.get("file_id"),
word_count=data.get('word_count'), word_count=data.get("word_count"),
parsing_completed_at=data.get('parsing_completed_at'), parsing_completed_at=data.get("parsing_completed_at"),
cleaning_completed_at=data.get('cleaning_completed_at'), cleaning_completed_at=data.get("cleaning_completed_at"),
splitting_completed_at=data.get('splitting_completed_at'), splitting_completed_at=data.get("splitting_completed_at"),
tokens=data.get('tokens'), tokens=data.get("tokens"),
indexing_latency=data.get('indexing_latency'), indexing_latency=data.get("indexing_latency"),
completed_at=data.get('completed_at'), completed_at=data.get("completed_at"),
is_paused=data.get('is_paused'), is_paused=data.get("is_paused"),
paused_by=data.get('paused_by'), paused_by=data.get("paused_by"),
paused_at=data.get('paused_at'), paused_at=data.get("paused_at"),
error=data.get('error'), error=data.get("error"),
stopped_at=data.get('stopped_at'), stopped_at=data.get("stopped_at"),
indexing_status=data.get('indexing_status'), indexing_status=data.get("indexing_status"),
enabled=data.get('enabled'), enabled=data.get("enabled"),
disabled_at=data.get('disabled_at'), disabled_at=data.get("disabled_at"),
disabled_by=data.get('disabled_by'), disabled_by=data.get("disabled_by"),
archived=data.get('archived'), archived=data.get("archived"),
archived_reason=data.get('archived_reason'), archived_reason=data.get("archived_reason"),
archived_by=data.get('archived_by'), archived_by=data.get("archived_by"),
archived_at=data.get('archived_at'), archived_at=data.get("archived_at"),
updated_at=data.get('updated_at'), updated_at=data.get("updated_at"),
doc_type=data.get('doc_type'), doc_type=data.get("doc_type"),
doc_metadata=data.get('doc_metadata'), doc_metadata=data.get("doc_metadata"),
doc_form=data.get('doc_form'), doc_form=data.get("doc_form"),
doc_language=data.get('doc_language') doc_language=data.get("doc_language"),
) )
class DocumentSegment(db.Model): class DocumentSegment(db.Model):
__tablename__ = 'document_segments' __tablename__ = "document_segments"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='document_segment_pkey'), db.PrimaryKeyConstraint("id", name="document_segment_pkey"),
db.Index('document_segment_dataset_id_idx', 'dataset_id'), db.Index("document_segment_dataset_id_idx", "dataset_id"),
db.Index('document_segment_document_id_idx', 'document_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_dataset_idx", "dataset_id", "tenant_id"),
db.Index('document_segment_tenant_document_idx', 'document_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_dataset_node_idx", "dataset_id", "index_node_id"),
db.Index('document_segment_tenant_idx', 'tenant_id'), db.Index("document_segment_tenant_idx", "tenant_id"),
) )
# initial fields # initial fields
id = db.Column(StringUUID, nullable=False, id = db.Column(StringUUID, nullable=False, server_default=db.text("uuid_generate_v4()"))
server_default=db.text('uuid_generate_v4()'))
tenant_id = db.Column(StringUUID, nullable=False) tenant_id = db.Column(StringUUID, nullable=False)
dataset_id = db.Column(StringUUID, nullable=False) dataset_id = db.Column(StringUUID, nullable=False)
document_id = db.Column(StringUUID, nullable=False) document_id = db.Column(StringUUID, nullable=False)
@ -486,18 +490,14 @@ class DocumentSegment(db.Model):
# basic fields # basic fields
hit_count = db.Column(db.Integer, nullable=False, default=0) hit_count = db.Column(db.Integer, nullable=False, default=0)
enabled = db.Column(db.Boolean, nullable=False, enabled = db.Column(db.Boolean, nullable=False, server_default=db.text("true"))
server_default=db.text('true'))
disabled_at = db.Column(db.DateTime, nullable=True) disabled_at = db.Column(db.DateTime, nullable=True)
disabled_by = db.Column(StringUUID, nullable=True) disabled_by = db.Column(StringUUID, nullable=True)
status = db.Column(db.String(255), nullable=False, status = db.Column(db.String(255), nullable=False, server_default=db.text("'waiting'::character varying"))
server_default=db.text("'waiting'::character varying"))
created_by = db.Column(StringUUID, nullable=False) created_by = db.Column(StringUUID, nullable=False)
created_at = db.Column(db.DateTime, nullable=False, created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
server_default=db.text('CURRENT_TIMESTAMP(0)'))
updated_by = db.Column(StringUUID, nullable=True) updated_by = db.Column(StringUUID, nullable=True)
updated_at = db.Column(db.DateTime, nullable=False, updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
server_default=db.text('CURRENT_TIMESTAMP(0)'))
indexing_at = db.Column(db.DateTime, nullable=True) indexing_at = db.Column(db.DateTime, nullable=True)
completed_at = db.Column(db.DateTime, nullable=True) completed_at = db.Column(db.DateTime, nullable=True)
error = db.Column(db.Text, nullable=True) error = db.Column(db.Text, nullable=True)
@ -513,17 +513,19 @@ class DocumentSegment(db.Model):
@property @property
def previous_segment(self): def previous_segment(self):
return db.session.query(DocumentSegment).filter( return (
DocumentSegment.document_id == self.document_id, db.session.query(DocumentSegment)
DocumentSegment.position == self.position - 1 .filter(DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position - 1)
).first() .first()
)
@property @property
def next_segment(self): def next_segment(self):
return db.session.query(DocumentSegment).filter( return (
DocumentSegment.document_id == self.document_id, db.session.query(DocumentSegment)
DocumentSegment.position == self.position + 1 .filter(DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position + 1)
).first() .first()
)
def get_sign_content(self): def get_sign_content(self):
pattern = r"/files/([a-f0-9\-]+)/image-preview" pattern = r"/files/([a-f0-9\-]+)/image-preview"
@ -535,7 +537,7 @@ class DocumentSegment(db.Model):
nonce = os.urandom(16).hex() nonce = os.urandom(16).hex()
timestamp = str(int(time.time())) timestamp = str(int(time.time()))
data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}" 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() sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
encoded_sign = base64.urlsafe_b64encode(sign).decode() encoded_sign = base64.urlsafe_b64encode(sign).decode()
@ -546,21 +548,20 @@ class DocumentSegment(db.Model):
# Reconstruct the text with signed URLs # Reconstruct the text with signed URLs
offset = 0 offset = 0
for start, end, signed_url in signed_urls: 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) offset += len(signed_url) - (end - start)
return text return text
class AppDatasetJoin(db.Model): class AppDatasetJoin(db.Model):
__tablename__ = 'app_dataset_joins' __tablename__ = "app_dataset_joins"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='app_dataset_join_pkey'), db.PrimaryKeyConstraint("id", name="app_dataset_join_pkey"),
db.Index('app_dataset_join_app_dataset_idx', 'dataset_id', 'app_id'), 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) app_id = db.Column(StringUUID, nullable=False)
dataset_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()) 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): class DatasetQuery(db.Model):
__tablename__ = 'dataset_queries' __tablename__ = "dataset_queries"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='dataset_query_pkey'), db.PrimaryKeyConstraint("id", name="dataset_query_pkey"),
db.Index('dataset_query_dataset_id_idx', 'dataset_id'), 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) dataset_id = db.Column(StringUUID, nullable=False)
content = db.Column(db.Text, nullable=False) content = db.Column(db.Text, nullable=False)
source = db.Column(db.String(255), nullable=False) source = db.Column(db.String(255), nullable=False)
@ -588,17 +589,18 @@ class DatasetQuery(db.Model):
class DatasetKeywordTable(db.Model): class DatasetKeywordTable(db.Model):
__tablename__ = 'dataset_keyword_tables' __tablename__ = "dataset_keyword_tables"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='dataset_keyword_table_pkey'), db.PrimaryKeyConstraint("id", name="dataset_keyword_table_pkey"),
db.Index('dataset_keyword_table_dataset_id_idx', 'dataset_id'), 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) dataset_id = db.Column(StringUUID, nullable=False, unique=True)
keyword_table = db.Column(db.Text, nullable=False) keyword_table = db.Column(db.Text, nullable=False)
data_source_type = db.Column(db.String(255), nullable=False, data_source_type = db.Column(
server_default=db.text("'database'::character varying")) db.String(255), nullable=False, server_default=db.text("'database'::character varying")
)
@property @property
def keyword_table_dict(self): def keyword_table_dict(self):
@ -614,19 +616,17 @@ class DatasetKeywordTable(db.Model):
return dct return dct
# get dataset # get dataset
dataset = Dataset.query.filter_by( dataset = Dataset.query.filter_by(id=self.dataset_id).first()
id=self.dataset_id
).first()
if not dataset: if not dataset:
return None 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 return json.loads(self.keyword_table, cls=SetDecoder) if self.keyword_table else None
else: else:
file_key = 'keyword_files/' + dataset.tenant_id + '/' + self.dataset_id + '.txt' file_key = "keyword_files/" + dataset.tenant_id + "/" + self.dataset_id + ".txt"
try: try:
keyword_table_text = storage.load_once(file_key) keyword_table_text = storage.load_once(file_key)
if keyword_table_text: 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 return None
except Exception as e: except Exception as e:
logging.exception(str(e)) logging.exception(str(e))
@ -634,21 +634,21 @@ class DatasetKeywordTable(db.Model):
class Embedding(db.Model): class Embedding(db.Model):
__tablename__ = 'embeddings' __tablename__ = "embeddings"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='embedding_pkey'), db.PrimaryKeyConstraint("id", name="embedding_pkey"),
db.UniqueConstraint('model_name', 'hash', 'provider_name', name='embedding_hash_idx'), db.UniqueConstraint("model_name", "hash", "provider_name", name="embedding_hash_idx"),
db.Index('created_at_idx', 'created_at') db.Index("created_at_idx", "created_at"),
) )
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()"))
model_name = db.Column(db.String(255), nullable=False, model_name = db.Column(
server_default=db.text("'text-embedding-ada-002'::character varying")) db.String(255), nullable=False, server_default=db.text("'text-embedding-ada-002'::character varying")
)
hash = db.Column(db.String(64), nullable=False) hash = db.Column(db.String(64), nullable=False)
embedding = db.Column(db.LargeBinary, nullable=False) embedding = db.Column(db.LargeBinary, 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)"))
provider_name = db.Column(db.String(255), nullable=False, provider_name = db.Column(db.String(255), nullable=False, server_default=db.text("''::character varying"))
server_default=db.text("''::character varying"))
def set_embedding(self, embedding_data: list[float]): def set_embedding(self, embedding_data: list[float]):
self.embedding = pickle.dumps(embedding_data, protocol=pickle.HIGHEST_PROTOCOL) self.embedding = pickle.dumps(embedding_data, protocol=pickle.HIGHEST_PROTOCOL)
@ -658,33 +658,32 @@ class Embedding(db.Model):
class DatasetCollectionBinding(db.Model): class DatasetCollectionBinding(db.Model):
__tablename__ = 'dataset_collection_bindings' __tablename__ = "dataset_collection_bindings"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='dataset_collection_bindings_pkey'), db.PrimaryKeyConstraint("id", name="dataset_collection_bindings_pkey"),
db.Index('provider_model_name_idx', 'provider_name', 'model_name') 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) provider_name = db.Column(db.String(40), nullable=False)
model_name = db.Column(db.String(255), 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) type = db.Column(db.String(40), server_default=db.text("'dataset'::character varying"), nullable=False)
collection_name = db.Column(db.String(64), 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): class DatasetPermission(db.Model):
__tablename__ = 'dataset_permissions' __tablename__ = "dataset_permissions"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='dataset_permission_pkey'), db.PrimaryKeyConstraint("id", name="dataset_permission_pkey"),
db.Index('idx_dataset_permissions_dataset_id', 'dataset_id'), db.Index("idx_dataset_permissions_dataset_id", "dataset_id"),
db.Index('idx_dataset_permissions_account_id', 'account_id'), db.Index("idx_dataset_permissions_account_id", "account_id"),
db.Index('idx_dataset_permissions_tenant_id', 'tenant_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) dataset_id = db.Column(StringUUID, nullable=False)
account_id = db.Column(StringUUID, nullable=False) account_id = db.Column(StringUUID, nullable=False)
tenant_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')) 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)')) created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))

871
api/models/model.py
View File

File diff suppressed because it is too large Load Diff

118
api/models/provider.py
View File

@ -6,8 +6,8 @@ from .types import StringUUID
class ProviderType(Enum): class ProviderType(Enum):
CUSTOM = 'custom' CUSTOM = "custom"
SYSTEM = 'system' SYSTEM = "system"
@staticmethod @staticmethod
def value_of(value): def value_of(value):
@ -18,13 +18,13 @@ class ProviderType(Enum):
class ProviderQuotaType(Enum): class ProviderQuotaType(Enum):
PAID = 'paid' PAID = "paid"
"""hosted paid quota""" """hosted paid quota"""
FREE = 'free' FREE = "free"
"""third-party free quota""" """third-party free quota"""
TRIAL = 'trial' TRIAL = "trial"
"""hosted trial quota""" """hosted trial quota"""
@staticmethod @staticmethod
@ -39,27 +39,30 @@ class Provider(db.Model):
""" """
Provider model representing the API providers and their configurations. Provider model representing the API providers and their configurations.
""" """
__tablename__ = 'providers'
__tablename__ = "providers"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='provider_pkey'), db.PrimaryKeyConstraint("id", name="provider_pkey"),
db.Index('provider_tenant_id_provider_idx', 'tenant_id', 'provider_name'), 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.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) tenant_id = db.Column(StringUUID, nullable=False)
provider_name = db.Column(db.String(255), 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")) provider_type = db.Column(db.String(40), nullable=False, server_default=db.text("'custom'::character varying"))
encrypted_config = db.Column(db.Text, nullable=True) 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) last_used = db.Column(db.DateTime, nullable=True)
quota_type = db.Column(db.String(40), nullable=True, server_default=db.text("''::character varying")) quota_type = db.Column(db.String(40), nullable=True, server_default=db.text("''::character varying"))
quota_limit = db.Column(db.BigInteger, nullable=True) quota_limit = db.Column(db.BigInteger, nullable=True)
quota_used = db.Column(db.BigInteger, default=0) quota_used = db.Column(db.BigInteger, default=0)
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_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): def __repr__(self):
return f"<Provider(id={self.id}, tenant_id={self.tenant_id}, provider_name='{self.provider_name}', provider_type='{self.provider_type}')>" return f"<Provider(id={self.id}, tenant_id={self.tenant_id}, provider_name='{self.provider_name}', provider_type='{self.provider_type}')>"
@ -86,118 +89,123 @@ class ProviderModel(db.Model):
""" """
Provider model representing the API provider_models and their configurations. Provider model representing the API provider_models and their configurations.
""" """
__tablename__ = 'provider_models'
__tablename__ = "provider_models"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='provider_model_pkey'), db.PrimaryKeyConstraint("id", name="provider_model_pkey"),
db.Index('provider_model_tenant_id_provider_idx', 'tenant_id', 'provider_name'), 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.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) tenant_id = db.Column(StringUUID, nullable=False)
provider_name = db.Column(db.String(255), nullable=False) provider_name = db.Column(db.String(255), nullable=False)
model_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) model_type = db.Column(db.String(40), nullable=False)
encrypted_config = db.Column(db.Text, nullable=True) 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"))
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_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): class TenantDefaultModel(db.Model):
__tablename__ = 'tenant_default_models' __tablename__ = "tenant_default_models"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='tenant_default_model_pkey'), 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.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) tenant_id = db.Column(StringUUID, nullable=False)
provider_name = db.Column(db.String(255), nullable=False) provider_name = db.Column(db.String(255), nullable=False)
model_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) model_type = db.Column(db.String(40), 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_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): class TenantPreferredModelProvider(db.Model):
__tablename__ = 'tenant_preferred_model_providers' __tablename__ = "tenant_preferred_model_providers"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='tenant_preferred_model_provider_pkey'), db.PrimaryKeyConstraint("id", name="tenant_preferred_model_provider_pkey"),
db.Index('tenant_preferred_model_provider_tenant_provider_idx', 'tenant_id', 'provider_name'), 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) tenant_id = db.Column(StringUUID, nullable=False)
provider_name = db.Column(db.String(255), nullable=False) provider_name = db.Column(db.String(255), nullable=False)
preferred_provider_type = db.Column(db.String(40), 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)')) 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)')) updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
class ProviderOrder(db.Model): class ProviderOrder(db.Model):
__tablename__ = 'provider_orders' __tablename__ = "provider_orders"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='provider_order_pkey'), db.PrimaryKeyConstraint("id", name="provider_order_pkey"),
db.Index('provider_order_tenant_provider_idx', 'tenant_id', 'provider_name'), 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) tenant_id = db.Column(StringUUID, nullable=False)
provider_name = db.Column(db.String(255), nullable=False) provider_name = db.Column(db.String(255), nullable=False)
account_id = db.Column(StringUUID, nullable=False) account_id = db.Column(StringUUID, nullable=False)
payment_product_id = db.Column(db.String(191), nullable=False) payment_product_id = db.Column(db.String(191), nullable=False)
payment_id = db.Column(db.String(191)) payment_id = db.Column(db.String(191))
transaction_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)) currency = db.Column(db.String(40))
total_amount = db.Column(db.Integer) total_amount = db.Column(db.Integer)
payment_status = db.Column(db.String(40), nullable=False, server_default=db.text("'wait_pay'::character varying")) payment_status = db.Column(db.String(40), nullable=False, server_default=db.text("'wait_pay'::character varying"))
paid_at = db.Column(db.DateTime) paid_at = db.Column(db.DateTime)
pay_failed_at = db.Column(db.DateTime) pay_failed_at = db.Column(db.DateTime)
refunded_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)')) 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)')) updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
class ProviderModelSetting(db.Model): class ProviderModelSetting(db.Model):
""" """
Provider model settings for record the model enabled status and load balancing status. Provider model settings for record the model enabled status and load balancing status.
""" """
__tablename__ = 'provider_model_settings'
__tablename__ = "provider_model_settings"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='provider_model_setting_pkey'), db.PrimaryKeyConstraint("id", name="provider_model_setting_pkey"),
db.Index('provider_model_setting_tenant_provider_model_idx', 'tenant_id', 'provider_name', 'model_type'), 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) tenant_id = db.Column(StringUUID, nullable=False)
provider_name = db.Column(db.String(255), nullable=False) provider_name = db.Column(db.String(255), nullable=False)
model_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) model_type = db.Column(db.String(40), nullable=False)
enabled = db.Column(db.Boolean, nullable=False, server_default=db.text('true')) 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')) 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)')) 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)')) updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
class LoadBalancingModelConfig(db.Model): class LoadBalancingModelConfig(db.Model):
""" """
Configurations for load balancing models. Configurations for load balancing models.
""" """
__tablename__ = 'load_balancing_model_configs'
__tablename__ = "load_balancing_model_configs"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='load_balancing_model_config_pkey'), 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.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) tenant_id = db.Column(StringUUID, nullable=False)
provider_name = db.Column(db.String(255), nullable=False) provider_name = db.Column(db.String(255), nullable=False)
model_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) model_type = db.Column(db.String(40), nullable=False)
name = db.Column(db.String(255), nullable=False) name = db.Column(db.String(255), nullable=False)
encrypted_config = db.Column(db.Text, nullable=True) encrypted_config = db.Column(db.Text, nullable=True)
enabled = db.Column(db.Boolean, nullable=False, server_default=db.text('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)')) 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)')) updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))

48
api/models/source.py
View File

@ -8,48 +8,48 @@ from .types import StringUUID
class DataSourceOauthBinding(db.Model): class DataSourceOauthBinding(db.Model):
__tablename__ = 'data_source_oauth_bindings' __tablename__ = "data_source_oauth_bindings"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='source_binding_pkey'), db.PrimaryKeyConstraint("id", name="source_binding_pkey"),
db.Index('source_binding_tenant_id_idx', 'tenant_id'), db.Index("source_binding_tenant_id_idx", "tenant_id"),
db.Index('source_info_idx', "source_info", postgresql_using='gin') 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) tenant_id = db.Column(StringUUID, nullable=False)
access_token = db.Column(db.String(255), nullable=False) access_token = db.Column(db.String(255), nullable=False)
provider = db.Column(db.String(255), nullable=False) provider = db.Column(db.String(255), nullable=False)
source_info = db.Column(JSONB, nullable=False) source_info = db.Column(JSONB, 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_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')) disabled = db.Column(db.Boolean, nullable=True, server_default=db.text("false"))
class DataSourceApiKeyAuthBinding(db.Model): class DataSourceApiKeyAuthBinding(db.Model):
__tablename__ = 'data_source_api_key_auth_bindings' __tablename__ = "data_source_api_key_auth_bindings"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='data_source_api_key_auth_binding_pkey'), 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_tenant_id_idx", "tenant_id"),
db.Index('data_source_api_key_auth_binding_provider_idx', 'provider'), 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) tenant_id = db.Column(StringUUID, nullable=False)
category = db.Column(db.String(255), nullable=False) category = db.Column(db.String(255), nullable=False)
provider = db.Column(db.String(255), nullable=False) provider = db.Column(db.String(255), nullable=False)
credentials = db.Column(db.Text, nullable=True) # JSON credentials = db.Column(db.Text, nullable=True) # 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)"))
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)"))
disabled = db.Column(db.Boolean, nullable=True, server_default=db.text('false')) disabled = db.Column(db.Boolean, nullable=True, server_default=db.text("false"))
def to_dict(self): def to_dict(self):
return { return {
'id': self.id, "id": self.id,
'tenant_id': self.tenant_id, "tenant_id": self.tenant_id,
'category': self.category, "category": self.category,
'provider': self.provider, "provider": self.provider,
'credentials': json.loads(self.credentials), "credentials": json.loads(self.credentials),
'created_at': self.created_at.timestamp(), "created_at": self.created_at.timestamp(),
'updated_at': self.updated_at.timestamp(), "updated_at": self.updated_at.timestamp(),
'disabled': self.disabled "disabled": self.disabled,
} }

21
api/models/task.py
View File

@ -8,15 +8,18 @@ from extensions.ext_database import db
class CeleryTask(db.Model): class CeleryTask(db.Model):
"""Task result/status.""" """Task result/status."""
__tablename__ = 'celery_taskmeta' __tablename__ = "celery_taskmeta"
id = db.Column(db.Integer, db.Sequence('task_id_sequence'), id = db.Column(db.Integer, db.Sequence("task_id_sequence"), primary_key=True, autoincrement=True)
primary_key=True, autoincrement=True)
task_id = db.Column(db.String(155), unique=True) task_id = db.Column(db.String(155), unique=True)
status = db.Column(db.String(50), default=states.PENDING) status = db.Column(db.String(50), default=states.PENDING)
result = db.Column(db.PickleType, nullable=True) result = db.Column(db.PickleType, nullable=True)
date_done = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc).replace(tzinfo=None), date_done = db.Column(
onupdate=lambda: datetime.now(timezone.utc).replace(tzinfo=None), nullable=True) 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) traceback = db.Column(db.Text, nullable=True)
name = db.Column(db.String(155), nullable=True) name = db.Column(db.String(155), nullable=True)
args = db.Column(db.LargeBinary, nullable=True) args = db.Column(db.LargeBinary, nullable=True)
@ -29,11 +32,9 @@ class CeleryTask(db.Model):
class CeleryTaskSet(db.Model): class CeleryTaskSet(db.Model):
"""TaskSet result.""" """TaskSet result."""
__tablename__ = 'celery_tasksetmeta' __tablename__ = "celery_tasksetmeta"
id = db.Column(db.Integer, db.Sequence('taskset_id_sequence'), id = db.Column(db.Integer, db.Sequence("taskset_id_sequence"), autoincrement=True, primary_key=True)
autoincrement=True, primary_key=True)
taskset_id = db.Column(db.String(155), unique=True) taskset_id = db.Column(db.String(155), unique=True)
result = db.Column(db.PickleType, nullable=True) result = db.Column(db.PickleType, nullable=True)
date_done = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc).replace(tzinfo=None), date_done = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc).replace(tzinfo=None), nullable=True)
nullable=True)

16
api/models/tool.py
View File

@ -7,7 +7,7 @@ from .types import StringUUID
class ToolProviderName(Enum): class ToolProviderName(Enum):
SERPAPI = 'serpapi' SERPAPI = "serpapi"
@staticmethod @staticmethod
def value_of(value): def value_of(value):
@ -18,19 +18,19 @@ class ToolProviderName(Enum):
class ToolProvider(db.Model): class ToolProvider(db.Model):
__tablename__ = 'tool_providers' __tablename__ = "tool_providers"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='tool_provider_pkey'), db.PrimaryKeyConstraint("id", name="tool_provider_pkey"),
db.UniqueConstraint('tenant_id', 'tool_name', name='unique_tool_provider_tool_name') 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) tenant_id = db.Column(StringUUID, nullable=False)
tool_name = db.Column(db.String(40), nullable=False) tool_name = db.Column(db.String(40), nullable=False)
encrypted_credentials = db.Column(db.Text, nullable=True) encrypted_credentials = db.Column(db.Text, nullable=True)
is_enabled = db.Column(db.Boolean, nullable=False, server_default=db.text('false')) 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)')) 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)')) updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
@property @property
def credentials_is_set(self): def credentials_is_set(self):

126
api/models/tools.py
View File

@ -15,15 +15,16 @@ class BuiltinToolProvider(db.Model):
""" """
This table stores the tool provider information for built-in tools for each tenant. This table stores the tool provider information for built-in tools for each tenant.
""" """
__tablename__ = 'tool_builtin_providers'
__tablename__ = "tool_builtin_providers"
__table_args__ = ( __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 # 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 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 # id of the tenant
tenant_id = db.Column(StringUUID, nullable=True) tenant_id = db.Column(StringUUID, nullable=True)
# who created this tool provider # who created this tool provider
@ -32,27 +33,29 @@ class BuiltinToolProvider(db.Model):
provider = db.Column(db.String(40), nullable=False) provider = db.Column(db.String(40), nullable=False)
# credential of the tool provider # credential of the tool provider
encrypted_credentials = db.Column(db.Text, nullable=True) encrypted_credentials = db.Column(db.Text, 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_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 @property
def credentials(self) -> dict: def credentials(self) -> dict:
return json.loads(self.encrypted_credentials) return json.loads(self.encrypted_credentials)
class PublishedAppTool(db.Model): class PublishedAppTool(db.Model):
""" """
The table stores the apps published as a tool for each person. The table stores the apps published as a tool for each person.
""" """
__tablename__ = 'tool_published_apps'
__tablename__ = "tool_published_apps"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='published_app_tool_pkey'), db.PrimaryKeyConstraint("id", name="published_app_tool_pkey"),
db.UniqueConstraint('app_id', 'user_id', name='unique_published_app_tool') db.UniqueConstraint("app_id", "user_id", name="unique_published_app_tool"),
) )
# id of the 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 app # 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 # who published this tool
user_id = db.Column(StringUUID, nullable=False) user_id = db.Column(StringUUID, nullable=False)
# description of the tool, stored in i18n format, for human # description of the tool, stored in i18n format, for human
@ -67,8 +70,8 @@ class PublishedAppTool(db.Model):
tool_name = db.Column(db.String(40), nullable=False) tool_name = db.Column(db.String(40), nullable=False)
# author # author
author = db.Column(db.String(40), nullable=False) author = db.Column(db.String(40), 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_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 @property
def description_i18n(self) -> I18nObject: def description_i18n(self) -> I18nObject:
@ -78,17 +81,19 @@ class PublishedAppTool(db.Model):
def app(self) -> App: def app(self) -> App:
return db.session.query(App).filter(App.id == self.app_id).first() return db.session.query(App).filter(App.id == self.app_id).first()
class ApiToolProvider(db.Model): class ApiToolProvider(db.Model):
""" """
The table stores the api providers. The table stores the api providers.
""" """
__tablename__ = 'tool_api_providers'
__tablename__ = "tool_api_providers"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='tool_api_provider_pkey'), db.PrimaryKeyConstraint("id", name="tool_api_provider_pkey"),
db.UniqueConstraint('name', 'tenant_id', name='unique_api_tool_provider') 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 of the api provider
name = db.Column(db.String(40), nullable=False) name = db.Column(db.String(40), nullable=False)
# icon # icon
@ -111,8 +116,8 @@ class ApiToolProvider(db.Model):
# custom_disclaimer # custom_disclaimer
custom_disclaimer = db.Column(db.String(255), nullable=True) custom_disclaimer = db.Column(db.String(255), 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_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 @property
def schema_type(self) -> ApiProviderSchemaType: def schema_type(self) -> ApiProviderSchemaType:
@ -134,17 +139,19 @@ class ApiToolProvider(db.Model):
def tenant(self) -> Tenant: def tenant(self) -> Tenant:
return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first() return db.session.query(Tenant).filter(Tenant.id == self.tenant_id).first()
class ToolLabelBinding(db.Model): class ToolLabelBinding(db.Model):
""" """
The table stores the labels for tools. The table stores the labels for tools.
""" """
__tablename__ = 'tool_label_bindings'
__tablename__ = "tool_label_bindings"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='tool_label_bind_pkey'), db.PrimaryKeyConstraint("id", name="tool_label_bind_pkey"),
db.UniqueConstraint('tool_id', 'label_name', name='unique_tool_label_bind'), 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
tool_id = db.Column(db.String(64), nullable=False) tool_id = db.Column(db.String(64), nullable=False)
# tool type # tool type
@ -152,28 +159,30 @@ class ToolLabelBinding(db.Model):
# label name # label name
label_name = db.Column(db.String(40), nullable=False) label_name = db.Column(db.String(40), nullable=False)
class WorkflowToolProvider(db.Model): class WorkflowToolProvider(db.Model):
""" """
The table stores the workflow providers. The table stores the workflow providers.
""" """
__tablename__ = 'tool_workflow_providers'
__tablename__ = "tool_workflow_providers"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='tool_workflow_provider_pkey'), db.PrimaryKeyConstraint("id", name="tool_workflow_provider_pkey"),
db.UniqueConstraint('name', 'tenant_id', name='unique_workflow_tool_provider'), db.UniqueConstraint("name", "tenant_id", name="unique_workflow_tool_provider"),
db.UniqueConstraint('tenant_id', 'app_id', name='unique_workflow_tool_provider_app_id'), 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 of the workflow provider
name = db.Column(db.String(40), nullable=False) name = db.Column(db.String(40), nullable=False)
# label of the workflow provider # 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
icon = db.Column(db.String(255), nullable=False) icon = db.Column(db.String(255), nullable=False)
# app id of the workflow provider # app id of the workflow provider
app_id = db.Column(StringUUID, nullable=False) app_id = db.Column(StringUUID, nullable=False)
# version of the workflow provider # 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 # who created this tool
user_id = db.Column(StringUUID, nullable=False) user_id = db.Column(StringUUID, nullable=False)
# tenant id # tenant id
@ -181,12 +190,12 @@ class WorkflowToolProvider(db.Model):
# description of the provider # description of the provider
description = db.Column(db.Text, nullable=False) description = db.Column(db.Text, nullable=False)
# parameter configuration # 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
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)')) 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)')) updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))
@property @property
def schema_type(self) -> ApiProviderSchemaType: def schema_type(self) -> ApiProviderSchemaType:
@ -202,25 +211,22 @@ class WorkflowToolProvider(db.Model):
@property @property
def parameter_configurations(self) -> list[WorkflowToolParameterConfiguration]: def parameter_configurations(self) -> list[WorkflowToolParameterConfiguration]:
return [ return [WorkflowToolParameterConfiguration(**config) for config in json.loads(self.parameter_configuration)]
WorkflowToolParameterConfiguration(**config)
for config in json.loads(self.parameter_configuration)
]
@property @property
def app(self) -> App: def app(self) -> App:
return db.session.query(App).filter(App.id == self.app_id).first() return db.session.query(App).filter(App.id == self.app_id).first()
class ToolModelInvoke(db.Model): class ToolModelInvoke(db.Model):
""" """
store the invoke logs from tool invoke 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 # who invoke this tool
user_id = db.Column(StringUUID, nullable=False) user_id = db.Column(StringUUID, nullable=False)
# tenant id # tenant id
@ -238,29 +244,31 @@ class ToolModelInvoke(db.Model):
# invoke response # invoke response
model_response = db.Column(db.Text, nullable=False) model_response = db.Column(db.Text, nullable=False)
prompt_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_tokens = db.Column(db.Integer, nullable=False, server_default=db.text("0"))
answer_unit_price = db.Column(db.Numeric(10, 4), nullable=False) 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')) 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')) provider_response_latency = db.Column(db.Float, nullable=False, server_default=db.text("0"))
total_price = db.Column(db.Numeric(10, 7)) total_price = db.Column(db.Numeric(10, 7))
currency = db.Column(db.String(255), nullable=False) currency = 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)"))
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 ToolConversationVariables(db.Model): class ToolConversationVariables(db.Model):
""" """
store the conversation variables from tool invoke store the conversation variables from tool invoke
""" """
__tablename__ = "tool_conversation_variables" __tablename__ = "tool_conversation_variables"
__table_args__ = ( __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 # add index for user_id and conversation_id
db.Index('user_id_idx', 'user_id'), db.Index("user_id_idx", "user_id"),
db.Index('conversation_id_idx', 'conversation_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 # conversation user id
user_id = db.Column(StringUUID, nullable=False) user_id = db.Column(StringUUID, nullable=False)
# tenant id # tenant id
@ -270,25 +278,27 @@ class ToolConversationVariables(db.Model):
# variables pool # variables pool
variables_str = db.Column(db.Text, nullable=False) variables_str = 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)"))
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 @property
def variables(self) -> dict: def variables(self) -> dict:
return json.loads(self.variables_str) return json.loads(self.variables_str)
class ToolFile(db.Model): class ToolFile(db.Model):
""" """
store the file created by agent store the file created by agent
""" """
__tablename__ = "tool_files" __tablename__ = "tool_files"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='tool_file_pkey'), db.PrimaryKeyConstraint("id", name="tool_file_pkey"),
# add index for conversation_id # 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 # conversation user id
user_id = db.Column(StringUUID, nullable=False) user_id = db.Column(StringUUID, nullable=False)
# tenant id # tenant id

4
api/models/types.py
View File

@ -9,13 +9,13 @@ class StringUUID(TypeDecorator):
def process_bind_param(self, value, dialect): def process_bind_param(self, value, dialect):
if value is None: if value is None:
return value return value
elif dialect.name == 'postgresql': elif dialect.name == "postgresql":
return str(value) return str(value)
else: else:
return value.hex return value.hex
def load_dialect_impl(self, dialect): def load_dialect_impl(self, dialect):
if dialect.name == 'postgresql': if dialect.name == "postgresql":
return dialect.type_descriptor(UUID()) return dialect.type_descriptor(UUID())
else: else:
return dialect.type_descriptor(CHAR(36)) return dialect.type_descriptor(CHAR(36))

21
api/models/web.py
View File

@ -1,4 +1,3 @@
from extensions.ext_database import db from extensions.ext_database import db
from .model import Message from .model import Message
@ -6,18 +5,18 @@ from .types import StringUUID
class SavedMessage(db.Model): class SavedMessage(db.Model):
__tablename__ = 'saved_messages' __tablename__ = "saved_messages"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='saved_message_pkey'), db.PrimaryKeyConstraint("id", name="saved_message_pkey"),
db.Index('saved_message_message_idx', 'app_id', 'message_id', 'created_by_role', 'created_by'), 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) app_id = db.Column(StringUUID, nullable=False)
message_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_role = db.Column(db.String(255), nullable=False, server_default=db.text("'end_user'::character varying"))
created_by = db.Column(StringUUID, 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 @property
def message(self): def message(self):
@ -25,15 +24,15 @@ class SavedMessage(db.Model):
class PinnedConversation(db.Model): class PinnedConversation(db.Model):
__tablename__ = 'pinned_conversations' __tablename__ = "pinned_conversations"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='pinned_conversation_pkey'), db.PrimaryKeyConstraint("id", name="pinned_conversation_pkey"),
db.Index('pinned_conversation_conversation_idx', 'app_id', 'conversation_id', 'created_by_role', 'created_by'), 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) app_id = db.Column(StringUUID, nullable=False)
conversation_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_role = db.Column(db.String(255), nullable=False, server_default=db.text("'end_user'::character varying"))
created_by = db.Column(StringUUID, 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)"))

357
api/models/workflow.py
View File

@ -22,11 +22,12 @@ class CreatedByRole(Enum):
""" """
Created By Role Enum Created By Role Enum
""" """
ACCOUNT = 'account'
END_USER = 'end_user' ACCOUNT = "account"
END_USER = "end_user"
@classmethod @classmethod
def value_of(cls, value: str) -> 'CreatedByRole': def value_of(cls, value: str) -> "CreatedByRole":
""" """
Get value of given mode. Get value of given mode.
@ -36,18 +37,19 @@ class CreatedByRole(Enum):
for mode in cls: for mode in cls:
if mode.value == value: if mode.value == value:
return mode return mode
raise ValueError(f'invalid created by role value {value}') raise ValueError(f"invalid created by role value {value}")
class WorkflowType(Enum): class WorkflowType(Enum):
""" """
Workflow Type Enum Workflow Type Enum
""" """
WORKFLOW = 'workflow'
CHAT = 'chat' WORKFLOW = "workflow"
CHAT = "chat"
@classmethod @classmethod
def value_of(cls, value: str) -> 'WorkflowType': def value_of(cls, value: str) -> "WorkflowType":
""" """
Get value of given mode. Get value of given mode.
@ -57,10 +59,10 @@ class WorkflowType(Enum):
for mode in cls: for mode in cls:
if mode.value == value: if mode.value == value:
return mode return mode
raise ValueError(f'invalid workflow type value {value}') raise ValueError(f"invalid workflow type value {value}")
@classmethod @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. Get workflow type from app mode.
@ -68,6 +70,7 @@ class WorkflowType(Enum):
:return: workflow type :return: workflow type
""" """
from models.model import AppMode from models.model import AppMode
app_mode = app_mode if isinstance(app_mode, AppMode) else AppMode.value_of(app_mode) 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 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 - updated_at (timestamp) `optional` Last update time
""" """
__tablename__ = 'workflows' __tablename__ = "workflows"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='workflow_pkey'), db.PrimaryKeyConstraint("id", name="workflow_pkey"),
db.Index('workflow_version_idx', 'tenant_id', 'app_id', 'version'), 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) tenant_id: Mapped[str] = db.Column(StringUUID, nullable=False)
app_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) 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) graph: Mapped[str] = db.Column(db.Text)
features: Mapped[str] = db.Column(db.Text) features: Mapped[str] = db.Column(db.Text)
created_by: Mapped[str] = db.Column(StringUUID, nullable=False) 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_by: Mapped[str] = db.Column(StringUUID)
updated_at: Mapped[datetime] = db.Column(db.DateTime) updated_at: Mapped[datetime] = db.Column(db.DateTime)
_environment_variables: Mapped[str] = db.Column('environment_variables', db.Text, nullable=False, server_default='{}') _environment_variables: Mapped[str] = db.Column(
_conversation_variables: Mapped[str] = db.Column('conversation_variables', db.Text, nullable=False, server_default='{}') "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, def __init__(
features: str, created_by: str, environment_variables: Sequence[Variable], self,
conversation_variables: Sequence[Variable]): *,
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.tenant_id = tenant_id
self.app_id = app_id self.app_id = app_id
self.type = type self.type = type
@ -160,22 +179,20 @@ class Workflow(db.Model):
return [] return []
graph_dict = self.graph_dict graph_dict = self.graph_dict
if 'nodes' not in graph_dict: if "nodes" not in graph_dict:
return [] 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: if not start_node:
return [] return []
# get user_input_form from start node # 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: if to_old_structure:
old_structure_variables = [] old_structure_variables = []
for variable in variables: for variable in variables:
old_structure_variables.append({ old_structure_variables.append({variable["type"]: variable})
variable['type']: variable
})
return old_structure_variables return old_structure_variables
@ -188,25 +205,24 @@ class Workflow(db.Model):
:return: hash :return: hash
""" """
entity = { entity = {"graph": self.graph_dict, "features": self.features_dict}
'graph': self.graph_dict,
'features': self.features_dict
}
return helper.generate_text_hash(json.dumps(entity, sort_keys=True)) return helper.generate_text_hash(json.dumps(entity, sort_keys=True))
@property @property
def tool_published(self) -> bool: def tool_published(self) -> bool:
from models.tools import WorkflowToolProvider from models.tools import WorkflowToolProvider
return db.session.query(WorkflowToolProvider).filter(
WorkflowToolProvider.app_id == self.app_id return (
).first() is not None db.session.query(WorkflowToolProvider).filter(WorkflowToolProvider.app_id == self.app_id).first()
is not None
)
@property @property
def environment_variables(self) -> Sequence[Variable]: def environment_variables(self) -> Sequence[Variable]:
# TODO: find some way to init `self._environment_variables` when instance created. # TODO: find some way to init `self._environment_variables` when instance created.
if self._environment_variables is None: if self._environment_variables is None:
self._environment_variables = '{}' self._environment_variables = "{}"
tenant_id = contexts.tenant_id.get() tenant_id = contexts.tenant_id.get()
@ -215,9 +231,7 @@ class Workflow(db.Model):
# decrypt secret variables value # decrypt secret variables value
decrypt_func = ( decrypt_func = (
lambda var: var.model_copy( lambda var: var.model_copy(update={"value": encrypter.decrypt_token(tenant_id=tenant_id, token=var.value)})
update={'value': encrypter.decrypt_token(tenant_id=tenant_id, token=var.value)}
)
if isinstance(var, SecretVariable) if isinstance(var, SecretVariable)
else var else var
) )
@ -230,19 +244,17 @@ class Workflow(db.Model):
value = list(value) value = list(value)
if any(var for var in value if not var.id): 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`). # 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} origin_variables_dictionary = {var.id: var for var in self.environment_variables}
for i, variable in enumerate(value): for i, variable in enumerate(value):
if variable.id in origin_variables_dictionary and variable.value == HIDDEN_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 secret variables value
encrypt_func = ( encrypt_func = (
lambda var: var.model_copy( lambda var: var.model_copy(update={"value": encrypter.encrypt_token(tenant_id=tenant_id, token=var.value)})
update={'value': encrypter.encrypt_token(tenant_id=tenant_id, token=var.value)}
)
if isinstance(var, SecretVariable) if isinstance(var, SecretVariable)
else var else var
) )
@ -256,15 +268,15 @@ class Workflow(db.Model):
def to_dict(self, *, include_secret: bool = False) -> Mapping[str, Any]: def to_dict(self, *, include_secret: bool = False) -> Mapping[str, Any]:
environment_variables = list(self.environment_variables) environment_variables = list(self.environment_variables)
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 for v in environment_variables
] ]
result = { result = {
'graph': self.graph_dict, "graph": self.graph_dict,
'features': self.features_dict, "features": self.features_dict,
'environment_variables': [var.model_dump(mode='json') for var in environment_variables], "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], "conversation_variables": [var.model_dump(mode="json") for var in self.conversation_variables],
} }
return result return result
@ -272,7 +284,7 @@ class Workflow(db.Model):
def conversation_variables(self) -> Sequence[Variable]: def conversation_variables(self) -> Sequence[Variable]:
# TODO: find some way to init `self._conversation_variables` when instance created. # TODO: find some way to init `self._conversation_variables` when instance created.
if self._conversation_variables is None: if self._conversation_variables is None:
self._conversation_variables = '{}' self._conversation_variables = "{}"
variables_dict: dict[str, Any] = json.loads(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()] 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 Workflow Run Triggered From Enum
""" """
DEBUGGING = 'debugging'
APP_RUN = 'app-run' DEBUGGING = "debugging"
APP_RUN = "app-run"
@classmethod @classmethod
def value_of(cls, value: str) -> 'WorkflowRunTriggeredFrom': def value_of(cls, value: str) -> "WorkflowRunTriggeredFrom":
""" """
Get value of given mode. Get value of given mode.
@ -304,20 +317,21 @@ class WorkflowRunTriggeredFrom(Enum):
for mode in cls: for mode in cls:
if mode.value == value: if mode.value == value:
return mode 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): class WorkflowRunStatus(Enum):
""" """
Workflow Run Status Enum Workflow Run Status Enum
""" """
RUNNING = 'running'
SUCCEEDED = 'succeeded' RUNNING = "running"
FAILED = 'failed' SUCCEEDED = "succeeded"
STOPPED = 'stopped' FAILED = "failed"
STOPPED = "stopped"
@classmethod @classmethod
def value_of(cls, value: str) -> 'WorkflowRunStatus': def value_of(cls, value: str) -> "WorkflowRunStatus":
""" """
Get value of given mode. Get value of given mode.
@ -327,7 +341,7 @@ class WorkflowRunStatus(Enum):
for mode in cls: for mode in cls:
if mode.value == value: if mode.value == value:
return mode return mode
raise ValueError(f'invalid workflow run status value {value}') raise ValueError(f"invalid workflow run status value {value}")
class WorkflowRun(db.Model): class WorkflowRun(db.Model):
@ -368,14 +382,14 @@ class WorkflowRun(db.Model):
- finished_at (timestamp) End time - finished_at (timestamp) End time
""" """
__tablename__ = 'workflow_runs' __tablename__ = "workflow_runs"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='workflow_run_pkey'), db.PrimaryKeyConstraint("id", name="workflow_run_pkey"),
db.Index('workflow_run_triggerd_from_idx', 'tenant_id', 'app_id', 'triggered_from'), 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.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) tenant_id = db.Column(StringUUID, nullable=False)
app_id = db.Column(StringUUID, nullable=False) app_id = db.Column(StringUUID, nullable=False)
sequence_number = db.Column(db.Integer, 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) status = db.Column(db.String(255), nullable=False)
outputs = db.Column(db.Text) outputs = db.Column(db.Text)
error = db.Column(db.Text) 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"))
total_tokens = db.Column(db.Integer, 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')) total_steps = db.Column(db.Integer, server_default=db.text("0"))
created_by_role = 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_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) finished_at = db.Column(db.DateTime)
@property @property
def created_by_account(self): def created_by_account(self):
created_by_role = CreatedByRole.value_of(self.created_by_role) created_by_role = CreatedByRole.value_of(self.created_by_role)
return db.session.get(Account, self.created_by) \ return db.session.get(Account, self.created_by) if created_by_role == CreatedByRole.ACCOUNT else None
if created_by_role == CreatedByRole.ACCOUNT else None
@property @property
def created_by_end_user(self): def created_by_end_user(self):
from models.model import EndUser from models.model import EndUser
created_by_role = CreatedByRole.value_of(self.created_by_role) created_by_role = CreatedByRole.value_of(self.created_by_role)
return db.session.get(EndUser, self.created_by) \ return db.session.get(EndUser, self.created_by) if created_by_role == CreatedByRole.END_USER else None
if created_by_role == CreatedByRole.END_USER else None
@property @property
def graph_dict(self): def graph_dict(self):
@ -422,12 +435,12 @@ class WorkflowRun(db.Model):
return json.loads(self.outputs) if self.outputs else None return json.loads(self.outputs) if self.outputs else None
@property @property
def message(self) -> Optional['Message']: def message(self) -> Optional["Message"]:
from models.model import Message from models.model import Message
return db.session.query(Message).filter(
Message.app_id == self.app_id, return (
Message.workflow_run_id == self.id db.session.query(Message).filter(Message.app_id == self.app_id, Message.workflow_run_id == self.id).first()
).first() )
@property @property
def workflow(self): def workflow(self):
@ -435,51 +448,51 @@ class WorkflowRun(db.Model):
def to_dict(self): def to_dict(self):
return { return {
'id': self.id, "id": self.id,
'tenant_id': self.tenant_id, "tenant_id": self.tenant_id,
'app_id': self.app_id, "app_id": self.app_id,
'sequence_number': self.sequence_number, "sequence_number": self.sequence_number,
'workflow_id': self.workflow_id, "workflow_id": self.workflow_id,
'type': self.type, "type": self.type,
'triggered_from': self.triggered_from, "triggered_from": self.triggered_from,
'version': self.version, "version": self.version,
'graph': self.graph_dict, "graph": self.graph_dict,
'inputs': self.inputs_dict, "inputs": self.inputs_dict,
'status': self.status, "status": self.status,
'outputs': self.outputs_dict, "outputs": self.outputs_dict,
'error': self.error, "error": self.error,
'elapsed_time': self.elapsed_time, "elapsed_time": self.elapsed_time,
'total_tokens': self.total_tokens, "total_tokens": self.total_tokens,
'total_steps': self.total_steps, "total_steps": self.total_steps,
'created_by_role': self.created_by_role, "created_by_role": self.created_by_role,
'created_by': self.created_by, "created_by": self.created_by,
'created_at': self.created_at, "created_at": self.created_at,
'finished_at': self.finished_at, "finished_at": self.finished_at,
} }
@classmethod @classmethod
def from_dict(cls, data: dict) -> 'WorkflowRun': def from_dict(cls, data: dict) -> "WorkflowRun":
return cls( return cls(
id=data.get('id'), id=data.get("id"),
tenant_id=data.get('tenant_id'), tenant_id=data.get("tenant_id"),
app_id=data.get('app_id'), app_id=data.get("app_id"),
sequence_number=data.get('sequence_number'), sequence_number=data.get("sequence_number"),
workflow_id=data.get('workflow_id'), workflow_id=data.get("workflow_id"),
type=data.get('type'), type=data.get("type"),
triggered_from=data.get('triggered_from'), triggered_from=data.get("triggered_from"),
version=data.get('version'), version=data.get("version"),
graph=json.dumps(data.get('graph')), graph=json.dumps(data.get("graph")),
inputs=json.dumps(data.get('inputs')), inputs=json.dumps(data.get("inputs")),
status=data.get('status'), status=data.get("status"),
outputs=json.dumps(data.get('outputs')), outputs=json.dumps(data.get("outputs")),
error=data.get('error'), error=data.get("error"),
elapsed_time=data.get('elapsed_time'), elapsed_time=data.get("elapsed_time"),
total_tokens=data.get('total_tokens'), total_tokens=data.get("total_tokens"),
total_steps=data.get('total_steps'), total_steps=data.get("total_steps"),
created_by_role=data.get('created_by_role'), created_by_role=data.get("created_by_role"),
created_by=data.get('created_by'), created_by=data.get("created_by"),
created_at=data.get('created_at'), created_at=data.get("created_at"),
finished_at=data.get('finished_at'), finished_at=data.get("finished_at"),
) )
@ -487,11 +500,12 @@ class WorkflowNodeExecutionTriggeredFrom(Enum):
""" """
Workflow Node Execution Triggered From Enum Workflow Node Execution Triggered From Enum
""" """
SINGLE_STEP = 'single-step'
WORKFLOW_RUN = 'workflow-run' SINGLE_STEP = "single-step"
WORKFLOW_RUN = "workflow-run"
@classmethod @classmethod
def value_of(cls, value: str) -> 'WorkflowNodeExecutionTriggeredFrom': def value_of(cls, value: str) -> "WorkflowNodeExecutionTriggeredFrom":
""" """
Get value of given mode. Get value of given mode.
@ -501,19 +515,20 @@ class WorkflowNodeExecutionTriggeredFrom(Enum):
for mode in cls: for mode in cls:
if mode.value == value: if mode.value == value:
return mode 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): class WorkflowNodeExecutionStatus(Enum):
""" """
Workflow Node Execution Status Enum Workflow Node Execution Status Enum
""" """
RUNNING = 'running'
SUCCEEDED = 'succeeded' RUNNING = "running"
FAILED = 'failed' SUCCEEDED = "succeeded"
FAILED = "failed"
@classmethod @classmethod
def value_of(cls, value: str) -> 'WorkflowNodeExecutionStatus': def value_of(cls, value: str) -> "WorkflowNodeExecutionStatus":
""" """
Get value of given mode. Get value of given mode.
@ -523,7 +538,7 @@ class WorkflowNodeExecutionStatus(Enum):
for mode in cls: for mode in cls:
if mode.value == value: if mode.value == value:
return mode 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): class WorkflowNodeExecution(db.Model):
@ -574,18 +589,31 @@ class WorkflowNodeExecution(db.Model):
- finished_at (timestamp) End time - finished_at (timestamp) End time
""" """
__tablename__ = 'workflow_node_executions' __tablename__ = "workflow_node_executions"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='workflow_node_execution_pkey'), db.PrimaryKeyConstraint("id", name="workflow_node_execution_pkey"),
db.Index('workflow_node_execution_workflow_run_idx', 'tenant_id', 'app_id', 'workflow_id', db.Index(
'triggered_from', 'workflow_run_id'), "workflow_node_execution_workflow_run_idx",
db.Index('workflow_node_execution_node_run_idx', 'tenant_id', 'app_id', 'workflow_id', "tenant_id",
'triggered_from', 'node_id'), "app_id",
db.Index('workflow_node_execution_id_idx', 'tenant_id', 'app_id', 'workflow_id', "workflow_id",
'triggered_from', 'node_execution_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) tenant_id = db.Column(StringUUID, nullable=False)
app_id = db.Column(StringUUID, nullable=False) app_id = db.Column(StringUUID, nullable=False)
workflow_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) outputs = db.Column(db.Text)
status = db.Column(db.String(255), nullable=False) status = db.Column(db.String(255), nullable=False)
error = db.Column(db.Text) 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) 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_role = db.Column(db.String(255), nullable=False)
created_by = db.Column(StringUUID, nullable=False) created_by = db.Column(StringUUID, nullable=False)
finished_at = db.Column(db.DateTime) finished_at = db.Column(db.DateTime)
@ -612,15 +640,14 @@ class WorkflowNodeExecution(db.Model):
@property @property
def created_by_account(self): def created_by_account(self):
created_by_role = CreatedByRole.value_of(self.created_by_role) created_by_role = CreatedByRole.value_of(self.created_by_role)
return db.session.get(Account, self.created_by) \ return db.session.get(Account, self.created_by) if created_by_role == CreatedByRole.ACCOUNT else None
if created_by_role == CreatedByRole.ACCOUNT else None
@property @property
def created_by_end_user(self): def created_by_end_user(self):
from models.model import EndUser from models.model import EndUser
created_by_role = CreatedByRole.value_of(self.created_by_role) created_by_role = CreatedByRole.value_of(self.created_by_role)
return db.session.get(EndUser, self.created_by) \ return db.session.get(EndUser, self.created_by) if created_by_role == CreatedByRole.END_USER else None
if created_by_role == CreatedByRole.END_USER else None
@property @property
def inputs_dict(self): def inputs_dict(self):
@ -641,15 +668,17 @@ class WorkflowNodeExecution(db.Model):
@property @property
def extras(self): def extras(self):
from core.tools.tool_manager import ToolManager from core.tools.tool_manager import ToolManager
extras = {} extras = {}
if self.execution_metadata_dict: if self.execution_metadata_dict:
from core.workflow.entities.node_entities import NodeType 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'] if self.node_type == NodeType.TOOL.value and "tool_info" in self.execution_metadata_dict:
extras['icon'] = ToolManager.get_tool_icon( tool_info = self.execution_metadata_dict["tool_info"]
extras["icon"] = ToolManager.get_tool_icon(
tenant_id=self.tenant_id, tenant_id=self.tenant_id,
provider_type=tool_info['provider_type'], provider_type=tool_info["provider_type"],
provider_id=tool_info['provider_id'] provider_id=tool_info["provider_id"],
) )
return extras return extras
@ -659,12 +688,13 @@ class WorkflowAppLogCreatedFrom(Enum):
""" """
Workflow App Log Created From Enum Workflow App Log Created From Enum
""" """
SERVICE_API = 'service-api'
WEB_APP = 'web-app' SERVICE_API = "service-api"
INSTALLED_APP = 'installed-app' WEB_APP = "web-app"
INSTALLED_APP = "installed-app"
@classmethod @classmethod
def value_of(cls, value: str) -> 'WorkflowAppLogCreatedFrom': def value_of(cls, value: str) -> "WorkflowAppLogCreatedFrom":
""" """
Get value of given mode. Get value of given mode.
@ -674,7 +704,7 @@ class WorkflowAppLogCreatedFrom(Enum):
for mode in cls: for mode in cls:
if mode.value == value: if mode.value == value:
return mode 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): class WorkflowAppLog(db.Model):
@ -706,13 +736,13 @@ class WorkflowAppLog(db.Model):
- created_at (timestamp) Creation time - created_at (timestamp) Creation time
""" """
__tablename__ = 'workflow_app_logs' __tablename__ = "workflow_app_logs"
__table_args__ = ( __table_args__ = (
db.PrimaryKeyConstraint('id', name='workflow_app_log_pkey'), db.PrimaryKeyConstraint("id", name="workflow_app_log_pkey"),
db.Index('workflow_app_log_app_idx', 'tenant_id', 'app_id'), 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) tenant_id = db.Column(StringUUID, nullable=False)
app_id = db.Column(StringUUID, nullable=False) app_id = db.Column(StringUUID, nullable=False)
workflow_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_from = db.Column(db.String(255), nullable=False)
created_by_role = 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_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 @property
def workflow_run(self): def workflow_run(self):
@ -729,26 +759,27 @@ class WorkflowAppLog(db.Model):
@property @property
def created_by_account(self): def created_by_account(self):
created_by_role = CreatedByRole.value_of(self.created_by_role) created_by_role = CreatedByRole.value_of(self.created_by_role)
return db.session.get(Account, self.created_by) \ return db.session.get(Account, self.created_by) if created_by_role == CreatedByRole.ACCOUNT else None
if created_by_role == CreatedByRole.ACCOUNT else None
@property @property
def created_by_end_user(self): def created_by_end_user(self):
from models.model import EndUser from models.model import EndUser
created_by_role = CreatedByRole.value_of(self.created_by_role) created_by_role = CreatedByRole.value_of(self.created_by_role)
return db.session.get(EndUser, self.created_by) \ return db.session.get(EndUser, self.created_by) if created_by_role == CreatedByRole.END_USER else None
if created_by_role == CreatedByRole.END_USER else None
class ConversationVariable(db.Model): class ConversationVariable(db.Model):
__tablename__ = 'workflow_conversation_variables' __tablename__ = "workflow_conversation_variables"
id: Mapped[str] = db.Column(StringUUID, primary_key=True) id: Mapped[str] = db.Column(StringUUID, primary_key=True)
conversation_id: Mapped[str] = db.Column(StringUUID, nullable=False, 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) app_id: Mapped[str] = db.Column(StringUUID, nullable=False, index=True)
data = db.Column(db.Text, nullable=False) data = db.Column(db.Text, nullable=False)
created_at = db.Column(db.DateTime, nullable=False, index=True, server_default=db.text('CURRENT_TIMESTAMP(0)')) 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()) 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: def __init__(self, *, id: str, app_id: str, conversation_id: str, data: str) -> None:
self.id = id self.id = id
@ -757,7 +788,7 @@ class ConversationVariable(db.Model):
self.data = data self.data = data
@classmethod @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( obj = cls(
id=variable.id, id=variable.id,
app_id=app_id, app_id=app_id,

1
api/pyproject.toml
View File

@ -68,7 +68,6 @@ ignore = [
[tool.ruff.format] [tool.ruff.format]
exclude = [ exclude = [
"models/**/*.py",
"migrations/**/*", "migrations/**/*",
] ]