import base64 import hashlib import hmac import json import logging import os import pickle import re import time from collections.abc import Sequence from datetime import datetime from json import JSONDecodeError from typing import Any, ClassVar, TypedDict, cast, override from uuid import uuid4 import sqlalchemy as sa from sqlalchemy import DateTime, String, func, select from sqlalchemy.orm import Mapped, Session, mapped_column, scoped_session from configs import dify_config from core.rag.entities import ParentMode, Rule from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from core.rag.index_processor.constant.query_type import QueryType from core.rag.retrieval.retrieval_methods import RetrievalMethod from core.tools.signature import sign_upload_file_preview_url from extensions.ext_storage import storage from libs.uuid_utils import uuidv7 from .account import Account from .base import Base, TypeBase from .engine import db from .enums import ( CollectionBindingType, CreatorUserRole, DatasetMetadataType, DatasetQuerySource, DatasetRuntimeMode, DataSourceType, DocumentCreatedFrom, DocumentDocType, IndexingStatus, ProcessRuleMode, SegmentStatus, SegmentType, SummaryStatus, TidbAuthBindingStatus, ) from .model import App, Tag, TagBinding, UploadFile from .types import AdjustedJSON, BinaryData, EnumText, LongText, StringUUID, adjusted_json_index logger = logging.getLogger(__name__) class PreProcessingRuleItem(TypedDict): id: str enabled: bool class SegmentationConfig(TypedDict): delimiter: str max_tokens: int chunk_overlap: int class AutomaticRulesConfig(TypedDict): pre_processing_rules: list[PreProcessingRuleItem] segmentation: SegmentationConfig class ProcessRuleDict(TypedDict): id: str dataset_id: str mode: str rules: dict[str, Any] | None class DocMetadataDetailItem(TypedDict): id: str name: str type: str value: Any class AttachmentItem(TypedDict): id: str name: str size: int extension: str mime_type: str source_url: str class DatasetBindingItem(TypedDict): id: str name: str from models.enums import PermissionEnum # Backward-compatible alias — new code should import PermissionEnum from models.enums DatasetPermissionEnum = PermissionEnum class Dataset(Base): __tablename__ = "datasets" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="dataset_pkey"), sa.Index("dataset_tenant_idx", "tenant_id"), sa.Index("dataset_tenant_maintainer_idx", "tenant_id", "maintainer"), adjusted_json_index("retrieval_model_idx", "retrieval_model"), ) INDEXING_TECHNIQUE_LIST = ["high_quality", "economy", None] PROVIDER_LIST = ["vendor", "external", None] DOC_FORM_LIST = [member.value for member in IndexStructureType] id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4())) tenant_id: Mapped[str] = mapped_column(StringUUID) name: Mapped[str] = mapped_column(String(255)) description = mapped_column(LongText, nullable=True) provider: Mapped[str] = mapped_column(String(255), server_default=sa.text("'vendor'")) permission: Mapped[DatasetPermissionEnum] = mapped_column( EnumText(DatasetPermissionEnum, length=255), default=DatasetPermissionEnum.ONLY_ME, ) data_source_type = mapped_column(EnumText(DataSourceType, length=255)) indexing_technique: Mapped[IndexTechniqueType | None] = mapped_column(EnumText(IndexTechniqueType, length=255)) index_struct = mapped_column(LongText, nullable=True) created_by = mapped_column(StringUUID, nullable=False) maintainer: Mapped[str | None] = mapped_column(StringUUID, nullable=True) created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp()) updated_by = mapped_column(StringUUID, nullable=True) updated_at = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp() ) embedding_model = mapped_column(sa.String(255), nullable=True) embedding_model_provider = mapped_column(sa.String(255), nullable=True) keyword_number = mapped_column(sa.Integer, nullable=True, server_default=sa.text("10")) collection_binding_id = mapped_column(StringUUID, nullable=True) retrieval_model = mapped_column(AdjustedJSON, nullable=True) summary_index_setting = mapped_column(AdjustedJSON, nullable=True) built_in_field_enabled = mapped_column(sa.Boolean, nullable=False, server_default=sa.false()) icon_info = mapped_column(AdjustedJSON, nullable=True) runtime_mode = mapped_column( EnumText(DatasetRuntimeMode, length=255), nullable=True, server_default=sa.text("'general'") ) pipeline_id = mapped_column(StringUUID, nullable=True) chunk_structure = mapped_column(sa.String(255), nullable=True) enable_api = mapped_column(sa.Boolean, nullable=False, server_default=sa.true()) is_multimodal = mapped_column(sa.Boolean, default=False, nullable=False) def get_total_documents(self, *, session: Session) -> int: return self.get_document_count(session=session) def get_total_available_documents(self, *, session: Session) -> int: return ( session.scalar( select(func.count(Document.id)).where( Document.dataset_id == self.id, Document.indexing_status == "completed", Document.enabled == True, Document.archived == False, ) ) or 0 ) def get_dataset_keyword_table(self, *, session: Session) -> "DatasetKeywordTable | None": return session.scalar(select(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == self.id)) @property def index_struct_dict(self): return json.loads(self.index_struct) if self.index_struct else None @property def external_retrieval_model(self): default_retrieval_model = { "top_k": 2, "score_threshold": 0.0, } return self.retrieval_model or default_retrieval_model def get_created_by_account(self, *, session: Session) -> Account | None: return session.get(Account, self.created_by) def get_author_name(self, *, session: Session) -> str | None: account = self.get_created_by_account(session=session) if account: return account.name return None def get_latest_process_rule(self, *, session: Session) -> "DatasetProcessRule | None": return session.scalar( select(DatasetProcessRule) .where(DatasetProcessRule.dataset_id == self.id) .order_by(DatasetProcessRule.created_at.desc()) .limit(1) ) def get_app_count(self, *, session: Session) -> int: return ( session.scalar( select(func.count(AppDatasetJoin.id)).where( AppDatasetJoin.dataset_id == self.id, App.id == AppDatasetJoin.app_id ) ) or 0 ) def get_document_count(self, *, session: Session) -> int: return session.scalar(select(func.count(Document.id)).where(Document.dataset_id == self.id)) or 0 def get_available_document_count(self, *, session: Session) -> int: return ( session.scalar( select(func.count(Document.id)).where( Document.dataset_id == self.id, Document.indexing_status == "completed", Document.enabled == True, Document.archived == False, ) ) or 0 ) @property def available_segment_count(self): return ( db.session.scalar( select(func.count(DocumentSegment.id)).where( DocumentSegment.dataset_id == self.id, DocumentSegment.status == "completed", DocumentSegment.enabled == True, ) ) or 0 ) def get_word_count(self, *, session: Session) -> int: return ( session.scalar( select(func.coalesce(func.sum(Document.word_count), 0)).where(Document.dataset_id == self.id) ) or 0 ) def get_doc_form(self, *, session: Session) -> str | None: if self.chunk_structure: return self.chunk_structure return session.scalar( select(Document.doc_form) .where(Document.dataset_id == self.id, Document.tenant_id == self.tenant_id) .limit(1) ) @property def retrieval_model_dict(self): """Return a normalized retrieval model payload for API responses. Older rows may only persist a partial retrieval model dict. Merge the stored value over the current defaults so response validation still sees the required baseline fields. """ default_retrieval_model = { "search_method": RetrievalMethod.SEMANTIC_SEARCH, "reranking_enable": False, "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""}, "top_k": 2, "score_threshold_enabled": False, } if not self.retrieval_model: return default_retrieval_model return {**default_retrieval_model, **self.retrieval_model} def get_tags(self, *, session: Session) -> Sequence[Tag]: tags = session.scalars( select(Tag) .join(TagBinding, Tag.id == TagBinding.tag_id) .where( TagBinding.target_id == self.id, TagBinding.tenant_id == self.tenant_id, Tag.tenant_id == self.tenant_id, Tag.type == "knowledge", ) ).all() return tags or [] def get_external_knowledge_info(self, *, session: Session) -> dict[str, Any] | None: if self.provider != "external": return None external_knowledge_binding = session.scalar( select(ExternalKnowledgeBindings).where( ExternalKnowledgeBindings.dataset_id == self.id, ExternalKnowledgeBindings.tenant_id == self.tenant_id, ) ) if not external_knowledge_binding: return None external_knowledge_api = session.scalar( select(ExternalKnowledgeApis).where( ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id, ExternalKnowledgeApis.tenant_id == self.tenant_id, ) ) return self.build_external_knowledge_info(external_knowledge_binding, external_knowledge_api) @staticmethod def build_external_knowledge_info( external_knowledge_binding: "ExternalKnowledgeBindings | None", external_knowledge_api: "ExternalKnowledgeApis | None", ) -> dict[str, Any] | None: """Format an external knowledge binding and its API into the response payload.""" if external_knowledge_binding is None: return None if external_knowledge_api is None or external_knowledge_api.settings is None: return None return { "external_knowledge_id": external_knowledge_binding.external_knowledge_id, "external_knowledge_api_id": external_knowledge_api.id, "external_knowledge_api_name": external_knowledge_api.name, "external_knowledge_api_endpoint": json.loads(external_knowledge_api.settings).get("endpoint", ""), } def get_is_published(self, *, session: Session) -> bool: if self.pipeline_id: pipeline = session.scalar(select(Pipeline).where(Pipeline.id == self.pipeline_id)) if pipeline: return pipeline.is_published return False def get_doc_metadata(self, *, session: Session) -> list[dict[str, str]]: dataset_metadatas = session.scalars(select(DatasetMetadata).where(DatasetMetadata.dataset_id == self.id)).all() return self.build_doc_metadata(dataset_metadatas) def build_doc_metadata(self, dataset_metadatas: "Sequence[DatasetMetadata]") -> list[dict[str, str]]: """Format metadata rows, appending the built-in fields when they are enabled.""" doc_metadata = [ { "id": dataset_metadata.id, "name": dataset_metadata.name, "type": dataset_metadata.type, } for dataset_metadata in dataset_metadatas ] if self.built_in_field_enabled: doc_metadata.append( { "id": "built-in", "name": BuiltInField.document_name, "type": "string", } ) doc_metadata.append( { "id": "built-in", "name": BuiltInField.uploader, "type": "string", } ) doc_metadata.append( { "id": "built-in", "name": BuiltInField.upload_date, "type": "time", } ) doc_metadata.append( { "id": "built-in", "name": BuiltInField.last_update_date, "type": "time", } ) doc_metadata.append( { "id": "built-in", "name": BuiltInField.source, "type": "string", } ) return doc_metadata @staticmethod def gen_collection_name_by_id(dataset_id: str) -> str: normalized_dataset_id = dataset_id.replace("-", "_") return f"{dify_config.VECTOR_INDEX_NAME_PREFIX}_{normalized_dataset_id}_Node" class DatasetProcessRule(TypeBase): __tablename__ = "dataset_process_rules" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="dataset_process_rule_pkey"), sa.Index("dataset_process_rule_dataset_id_idx", "dataset_id"), ) id: Mapped[str] = mapped_column(StringUUID, nullable=False, default_factory=lambda: str(uuid4()), init=False) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) mode: Mapped[ProcessRuleMode] = mapped_column( EnumText(ProcessRuleMode, length=255), nullable=False, server_default=sa.text("'automatic'") ) rules: Mapped[str | None] = mapped_column(LongText, nullable=True) created_by: Mapped[str] = mapped_column(StringUUID, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) MODES = ["automatic", "custom", "hierarchical"] PRE_PROCESSING_RULES = ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"] AUTOMATIC_RULES: ClassVar[AutomaticRulesConfig] = { "pre_processing_rules": [ {"id": "remove_extra_spaces", "enabled": True}, {"id": "remove_urls_emails", "enabled": False}, ], "segmentation": {"delimiter": "\n", "max_tokens": 500, "chunk_overlap": 50}, } def to_dict(self) -> ProcessRuleDict: return { "id": self.id, "dataset_id": self.dataset_id, "mode": self.mode, "rules": self.rules_dict, } @property def rules_dict(self) -> dict[str, Any] | None: try: return json.loads(self.rules) if self.rules else None except JSONDecodeError: return None class Document(Base): __tablename__ = "documents" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="document_pkey"), sa.Index("document_dataset_id_idx", "dataset_id"), sa.Index("document_is_paused_idx", "is_paused"), sa.Index("document_tenant_idx", "tenant_id"), adjusted_json_index("document_metadata_idx", "doc_metadata"), ) # initial fields id = mapped_column(StringUUID, nullable=False, default=lambda: str(uuid4())) tenant_id = mapped_column(StringUUID, nullable=False) dataset_id = mapped_column(StringUUID, nullable=False) position: Mapped[int] = mapped_column(sa.Integer, nullable=False) data_source_type: Mapped[str] = mapped_column(EnumText(DataSourceType, length=255), nullable=False) data_source_info = mapped_column(LongText, nullable=True) dataset_process_rule_id = mapped_column(StringUUID, nullable=True) batch: Mapped[str] = mapped_column(String(255), nullable=False) name: Mapped[str] = mapped_column(String(255), nullable=False) created_from: Mapped[str] = mapped_column(EnumText(DocumentCreatedFrom, length=255), nullable=False) created_by = mapped_column(StringUUID, nullable=False) created_api_request_id = mapped_column(StringUUID, nullable=True) created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp()) # start processing processing_started_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) # parsing file_id = mapped_column(LongText, nullable=True) word_count: Mapped[int | None] = mapped_column(sa.Integer, nullable=True) # TODO: make this not nullable parsing_completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) # cleaning cleaning_completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) # split splitting_completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) # indexing tokens: Mapped[int | None] = mapped_column(sa.Integer, nullable=True) indexing_latency: Mapped[float | None] = mapped_column(sa.Float, nullable=True) completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) # pause is_paused: Mapped[bool | None] = mapped_column(sa.Boolean, nullable=True, server_default=sa.false()) paused_by = mapped_column(StringUUID, nullable=True) paused_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) # error error = mapped_column(LongText, nullable=True) stopped_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) # basic fields indexing_status = mapped_column( EnumText(IndexingStatus, length=255), nullable=False, server_default=sa.text("'waiting'") ) enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.true()) disabled_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) disabled_by = mapped_column(StringUUID, nullable=True) archived: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.false()) archived_reason = mapped_column(String(255), nullable=True) archived_by = mapped_column(StringUUID, nullable=True) archived_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True) updated_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp() ) doc_type = mapped_column(EnumText(DocumentDocType, length=40), nullable=True) doc_metadata = mapped_column(AdjustedJSON, nullable=True) doc_form: Mapped[IndexStructureType] = mapped_column( EnumText(IndexStructureType, length=255), nullable=False, server_default=sa.text("'text_model'") ) doc_language = mapped_column(String(255), nullable=True) need_summary: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.false()) DATA_SOURCES = ["upload_file", "notion_import", "website_crawl"] @property def display_status(self): status = None if self.indexing_status == "waiting": status = "queuing" elif self.indexing_status not in {"completed", "error", "waiting"} and self.is_paused: status = "paused" elif self.indexing_status in {"parsing", "cleaning", "splitting", "indexing"}: status = "indexing" elif self.indexing_status == "error": status = "error" elif self.indexing_status == "completed" and not self.archived and self.enabled: status = "available" elif self.indexing_status == "completed" and not self.archived and not self.enabled: status = "disabled" elif self.indexing_status == "completed" and self.archived: status = "archived" return status @property def data_source_info_dict(self) -> dict[str, Any]: if self.data_source_info: try: data_source_info_dict: dict[str, Any] = json.loads(self.data_source_info) except JSONDecodeError: data_source_info_dict = {} return data_source_info_dict return {} def get_data_source_detail_dict(self, *, session: Session) -> dict[str, Any]: if self.data_source_info: if self.data_source_type == "upload_file": data_source_info_dict: dict[str, Any] = json.loads(self.data_source_info) file_detail = session.scalar( select(UploadFile).where(UploadFile.id == data_source_info_dict["upload_file_id"]) ) if file_detail: return { "upload_file": { "id": file_detail.id, "name": file_detail.name, "size": file_detail.size, "extension": file_detail.extension, "mime_type": file_detail.mime_type, "created_by": file_detail.created_by, "created_at": file_detail.created_at.timestamp(), } } elif self.data_source_type in {"notion_import", "website_crawl"}: result: dict[str, Any] = json.loads(self.data_source_info) return result return {} def get_dataset_process_rule(self, *, session: Session) -> "DatasetProcessRule | None": if self.dataset_process_rule_id: return session.get(DatasetProcessRule, self.dataset_process_rule_id) return None def get_dataset(self, *, session: Session) -> Dataset | None: """Load the owning dataset with the caller-owned database session.""" return session.get(Dataset, self.dataset_id) def get_segment_count(self, *, session: Session) -> int: return session.scalar(select(func.count(DocumentSegment.id)).where(DocumentSegment.document_id == self.id)) or 0 def get_hit_count(self, *, session: Session) -> int: return ( session.scalar( select(func.coalesce(func.sum(DocumentSegment.hit_count), 0)).where( DocumentSegment.document_id == self.id ) ) or 0 ) def get_uploader(self, *, session: Session) -> str | None: user = session.scalar(select(Account).where(Account.id == self.created_by)) return user.name if user else None @property def upload_date(self): return self.created_at @property def last_update_date(self): return self.updated_at def get_doc_metadata_details(self, *, session: Session) -> list[DocMetadataDetailItem] | None: if self.doc_metadata: document_metadatas = session.scalars( select(DatasetMetadata) .join(DatasetMetadataBinding, DatasetMetadataBinding.metadata_id == DatasetMetadata.id) .where( DatasetMetadata.tenant_id == self.tenant_id, DatasetMetadata.dataset_id == self.dataset_id, DatasetMetadataBinding.tenant_id == self.tenant_id, DatasetMetadataBinding.dataset_id == self.dataset_id, DatasetMetadataBinding.document_id == self.id, ) ).all() metadata_list: list[DocMetadataDetailItem] = [] for metadata in document_metadatas: metadata_dict: DocMetadataDetailItem = { "id": metadata.id, "name": metadata.name, "type": metadata.type, "value": self.doc_metadata.get(metadata.name), } metadata_list.append(metadata_dict) # deal built-in fields metadata_list.extend(self.get_built_in_fields(session=session)) return metadata_list return None def get_built_in_fields(self, *, session: Session) -> list[DocMetadataDetailItem]: built_in_fields: list[DocMetadataDetailItem] = [] built_in_fields.append( { "id": "built-in", "name": BuiltInField.document_name, "type": "string", "value": self.name, } ) built_in_fields.append( { "id": "built-in", "name": BuiltInField.uploader, "type": "string", "value": self.get_uploader(session=session), } ) built_in_fields.append( { "id": "built-in", "name": BuiltInField.upload_date, "type": "time", "value": str(self.created_at.timestamp()), } ) built_in_fields.append( { "id": "built-in", "name": BuiltInField.last_update_date, "type": "time", "value": str(self.updated_at.timestamp()), } ) built_in_fields.append( { "id": "built-in", "name": BuiltInField.source, "type": "string", "value": MetadataDataSource[self.data_source_type], } ) return built_in_fields @classmethod def from_dict(cls, data: dict[str, Any]): return cls( id=data.get("id"), tenant_id=data.get("tenant_id"), dataset_id=data.get("dataset_id"), position=data.get("position"), data_source_type=data.get("data_source_type"), data_source_info=data.get("data_source_info"), dataset_process_rule_id=data.get("dataset_process_rule_id"), batch=data.get("batch"), name=data.get("name"), created_from=data.get("created_from"), created_by=data.get("created_by"), created_api_request_id=data.get("created_api_request_id"), created_at=data.get("created_at"), processing_started_at=data.get("processing_started_at"), file_id=data.get("file_id"), word_count=data.get("word_count"), parsing_completed_at=data.get("parsing_completed_at"), cleaning_completed_at=data.get("cleaning_completed_at"), splitting_completed_at=data.get("splitting_completed_at"), tokens=data.get("tokens"), indexing_latency=data.get("indexing_latency"), completed_at=data.get("completed_at"), is_paused=data.get("is_paused"), paused_by=data.get("paused_by"), paused_at=data.get("paused_at"), error=data.get("error"), stopped_at=data.get("stopped_at"), indexing_status=data.get("indexing_status"), enabled=data.get("enabled"), disabled_at=data.get("disabled_at"), disabled_by=data.get("disabled_by"), archived=data.get("archived"), archived_reason=data.get("archived_reason"), archived_by=data.get("archived_by"), archived_at=data.get("archived_at"), updated_at=data.get("updated_at"), doc_type=data.get("doc_type"), doc_metadata=data.get("doc_metadata"), doc_form=data.get("doc_form"), doc_language=data.get("doc_language"), ) class DocumentSegment(TypeBase): __tablename__ = "document_segments" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="document_segment_pkey"), sa.Index("document_segment_dataset_id_idx", "dataset_id"), sa.Index("document_segment_document_id_idx", "document_id"), sa.Index("document_segment_tenant_dataset_idx", "dataset_id", "tenant_id"), sa.Index("document_segment_tenant_document_idx", "document_id", "tenant_id"), sa.Index("document_segment_node_dataset_idx", "index_node_id", "dataset_id"), sa.Index("document_segment_tenant_idx", "tenant_id"), ) # initial fields id: Mapped[str] = mapped_column(StringUUID, nullable=False, default_factory=lambda: str(uuid4()), init=False) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) document_id: Mapped[str] = mapped_column(StringUUID, nullable=False) position: Mapped[int] = mapped_column(sa.Integer, nullable=False) content: Mapped[str] = mapped_column(LongText, nullable=False) word_count: Mapped[int] = mapped_column(sa.Integer, nullable=False) tokens: Mapped[int] = mapped_column(sa.Integer, nullable=False) created_by: Mapped[str] = mapped_column(StringUUID, nullable=False) # basic fields # indexing fields index_node_id: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None) index_node_hash: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None) enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=True) answer: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) keywords: Mapped[Any] = mapped_column(sa.JSON, nullable=True, default=None) disabled_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) disabled_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) status: Mapped[SegmentStatus] = mapped_column(EnumText(SegmentStatus, length=255), default=SegmentStatus.WAITING) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) updated_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) indexing_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) completed_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) error: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) stopped_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) hit_count: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0) def get_dataset(self, *, session: Session) -> Dataset | None: """Load the owning dataset with the caller-owned database session.""" return session.get(Dataset, self.dataset_id) def get_document(self, *, session: Session) -> Document | None: """Load the owning document with the caller-owned database session.""" return session.get(Document, self.document_id) def previous_segment(self, session: Session) -> "DocumentSegment | None": return session.scalar( select(DocumentSegment).where( DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position - 1 ) ) def next_segment(self, session: Session) -> "DocumentSegment | None": return session.scalar( select(DocumentSegment).where( DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position + 1 ) ) def get_child_chunks(self, *, session: Session, include_full_doc: bool = True) -> Sequence["ChildChunk"]: """Load hierarchical child chunks with the caller-owned database session.""" document = session.get(Document, self.document_id) if not document: return [] process_rule = document.get_dataset_process_rule(session=session) if process_rule and process_rule.mode == "hierarchical": rules_dict = process_rule.rules_dict if rules_dict: rules = Rule.model_validate(rules_dict) if rules.parent_mode and (include_full_doc or rules.parent_mode != ParentMode.FULL_DOC): child_chunks = session.scalars( select(ChildChunk).where(ChildChunk.segment_id == self.id).order_by(ChildChunk.position.asc()) ).all() return child_chunks or [] return [] @property def sign_content(self) -> str: return self.get_sign_content() def get_sign_content(self) -> str: signed_urls: list[tuple[int, int, str]] = [] text = self.content # For data before v0.10.0 pattern = r"/files/([a-f0-9\-]+)/image-preview(?:\?.*?)?" matches = re.finditer(pattern, text) for match in matches: upload_file_id = match.group(1) nonce = os.urandom(16).hex() timestamp = str(int(time.time())) data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}" secret_key = dify_config.SECRET_KEY.encode() sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() encoded_sign = base64.urlsafe_b64encode(sign).decode() params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" base_url = f"/files/{upload_file_id}/image-preview" signed_url = f"{base_url}?{params}" signed_urls.append((match.start(), match.end(), signed_url)) # For data after v0.10.0 pattern = r"/files/([a-f0-9\-]+)/file-preview(?:\?.*?)?" matches = re.finditer(pattern, text) for match in matches: upload_file_id = match.group(1) nonce = os.urandom(16).hex() timestamp = str(int(time.time())) data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}" secret_key = dify_config.SECRET_KEY.encode() sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() encoded_sign = base64.urlsafe_b64encode(sign).decode() params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" base_url = f"/files/{upload_file_id}/file-preview" signed_url = f"{base_url}?{params}" signed_urls.append((match.start(), match.end(), signed_url)) # For tools directory - direct file formats (e.g., .png, .jpg, etc.) # Match URL including any query parameters up to common URL boundaries (space, parenthesis, quotes) pattern = r"/files/tools/([a-f0-9\-]+)\.([a-zA-Z0-9]+)(?:\?[^\s\)\"\']*)?" matches = re.finditer(pattern, text) for match in matches: upload_file_id = match.group(1) file_extension = match.group(2) nonce = os.urandom(16).hex() timestamp = str(int(time.time())) data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}" secret_key = dify_config.SECRET_KEY.encode() sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() encoded_sign = base64.urlsafe_b64encode(sign).decode() params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" base_url = f"/files/tools/{upload_file_id}.{file_extension}" signed_url = f"{base_url}?{params}" signed_urls.append((match.start(), match.end(), signed_url)) # Reconstruct the text with signed URLs offset = 0 for start, end, signed_url in signed_urls: text = text[: start + offset] + signed_url + text[end + offset :] offset += len(signed_url) - (end - start) return text def get_attachments(self, *, session: Session) -> list[AttachmentItem]: """Load attachment metadata with the caller-owned database session.""" # Use JOIN to fetch attachments in a single query instead of two separate queries attachments_with_bindings = session.execute( select(SegmentAttachmentBinding, UploadFile) .join(UploadFile, UploadFile.id == SegmentAttachmentBinding.attachment_id) .where( SegmentAttachmentBinding.tenant_id == self.tenant_id, SegmentAttachmentBinding.dataset_id == self.dataset_id, SegmentAttachmentBinding.document_id == self.document_id, SegmentAttachmentBinding.segment_id == self.id, ) ).all() if not attachments_with_bindings: return [] attachment_list: list[AttachmentItem] = [] for _, attachment in attachments_with_bindings: upload_file_id = attachment.id nonce = os.urandom(16).hex() timestamp = str(int(time.time())) data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}" secret_key = dify_config.SECRET_KEY.encode() sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest() encoded_sign = base64.urlsafe_b64encode(sign).decode() params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}" reference_url = dify_config.FILES_URL or dify_config.CONSOLE_API_URL or "" base_url = f"{reference_url}/files/{upload_file_id}/image-preview" source_url = f"{base_url}?{params}" attachment_list.append( { "id": attachment.id, "name": attachment.name, "size": attachment.size, "extension": attachment.extension, "mime_type": attachment.mime_type, "source_url": source_url, } ) return attachment_list class ChildChunk(TypeBase): __tablename__ = "child_chunks" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="child_chunk_pkey"), sa.Index("child_chunk_dataset_id_idx", "tenant_id", "dataset_id", "document_id", "segment_id", "index_node_id"), sa.Index("child_chunks_node_idx", "index_node_id", "dataset_id"), sa.Index("child_chunks_segment_idx", "segment_id"), ) # initial fields id: Mapped[str] = mapped_column(StringUUID, nullable=False, default_factory=lambda: str(uuid4()), init=False) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) document_id: Mapped[str] = mapped_column(StringUUID, nullable=False) segment_id: Mapped[str] = mapped_column(StringUUID, nullable=False) position: Mapped[int] = mapped_column(sa.Integer, nullable=False) content: Mapped[str] = mapped_column(LongText, nullable=False) word_count: Mapped[int] = mapped_column(sa.Integer, nullable=False) # indexing fields created_by: Mapped[str] = mapped_column(StringUUID, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False ) updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, init=False) updated_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=sa.func.current_timestamp(), onupdate=func.current_timestamp(), init=False, ) indexing_at: Mapped[datetime | None] = mapped_column( DateTime, nullable=True, insert_default=None, server_default=None, init=False ) completed_at: Mapped[datetime | None] = mapped_column( DateTime, nullable=True, insert_default=None, server_default=None, init=False ) index_node_id: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None) index_node_hash: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None) type: Mapped[SegmentType] = mapped_column( EnumText(SegmentType, length=255), nullable=False, default=SegmentType.AUTOMATIC, ) error: Mapped[str | None] = mapped_column(LongText, nullable=True, init=False) def dataset(self, session: Session) -> Dataset | None: return session.scalar(select(Dataset).where(Dataset.id == self.dataset_id)) def document(self, session: Session) -> Document | None: return session.scalar(select(Document).where(Document.id == self.document_id)) def segment(self, session: Session) -> DocumentSegment | None: return session.scalar(select(DocumentSegment).where(DocumentSegment.id == self.segment_id)) class AppDatasetJoin(TypeBase): __tablename__ = "app_dataset_joins" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="app_dataset_join_pkey"), sa.Index("app_dataset_join_app_dataset_idx", "dataset_id", "app_id"), ) id: Mapped[str] = mapped_column( StringUUID, primary_key=True, nullable=False, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False, ) app_id: Mapped[str] = mapped_column(StringUUID, nullable=False) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False ) def app(self, session: Session) -> App | None: return session.get(App, self.app_id) class DatasetQuery(TypeBase): __tablename__ = "dataset_queries" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="dataset_query_pkey"), sa.Index("dataset_query_dataset_id_idx", "dataset_id"), ) id: Mapped[str] = mapped_column( StringUUID, primary_key=True, nullable=False, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False, ) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) content: Mapped[str] = mapped_column(LongText, nullable=False) source: Mapped[str] = mapped_column(EnumText(DatasetQuerySource, length=255), nullable=False) source_app_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) created_by_role: Mapped[CreatorUserRole] = mapped_column(EnumText(CreatorUserRole, length=255), nullable=False) created_by: Mapped[str] = mapped_column(StringUUID, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False ) def get_queries(self, *, session: Session) -> list[dict[str, Any]]: try: queries = json.loads(self.content) if isinstance(queries, list): for query in queries: if query["content_type"] == QueryType.IMAGE_QUERY: file_info = session.scalar(select(UploadFile).where(UploadFile.id == query["content"])) if file_info: query["file_info"] = { "id": file_info.id, "name": file_info.name, "size": file_info.size, "extension": file_info.extension, "mime_type": file_info.mime_type, "source_url": sign_upload_file_preview_url(file_info.id, file_info.extension), } else: query["file_info"] = None return queries else: return [queries] except JSONDecodeError: return [ { "content_type": QueryType.TEXT_QUERY, "content": self.content, "file_info": None, } ] class DatasetKeywordTable(TypeBase): __tablename__ = "dataset_keyword_tables" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="dataset_keyword_table_pkey"), sa.Index("dataset_keyword_table_dataset_id_idx", "dataset_id"), ) id: Mapped[str] = mapped_column( StringUUID, primary_key=True, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False, ) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False, unique=True) keyword_table: Mapped[str] = mapped_column(LongText, nullable=False) data_source_type: Mapped[str] = mapped_column(String(255), nullable=False, default="database") def get_keyword_table_dict(self, *, session: Session) -> dict[str, set[Any]] | None: class SetDecoder(json.JSONDecoder): def __init__(self, *args: Any, **kwargs: Any) -> None: def object_hook(dct: Any) -> Any: if isinstance(dct, dict): result: dict[str, Any] = {} items = cast(dict[str, Any], dct).items() for keyword, node_idxs in items: if isinstance(node_idxs, list): result[keyword] = set(node_idxs) else: result[keyword] = node_idxs return result return dct super().__init__(object_hook=object_hook, *args, **kwargs) # get dataset dataset = session.scalar(select(Dataset).where(Dataset.id == self.dataset_id)) if not dataset: return None if self.data_source_type == "database": return json.loads(self.keyword_table, cls=SetDecoder) if self.keyword_table else None else: file_key = "keyword_files/" + dataset.tenant_id + "/" + self.dataset_id + ".txt" try: keyword_table_text = storage.load_once(file_key) if keyword_table_text: return json.loads(keyword_table_text.decode("utf-8"), cls=SetDecoder) return None except Exception: logger.exception("Failed to load keyword table from file: %s", file_key) return None class Embedding(TypeBase): __tablename__ = "embeddings" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="embedding_pkey"), sa.UniqueConstraint("model_name", "hash", "provider_name", name="embedding_hash_idx"), sa.Index("created_at_idx", "created_at"), ) id: Mapped[str] = mapped_column( StringUUID, primary_key=True, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False, ) model_name: Mapped[str] = mapped_column( String(255), nullable=False, server_default=sa.text("'text-embedding-ada-002'") ) hash: Mapped[str] = mapped_column(String(64), nullable=False) embedding: Mapped[bytes] = mapped_column(BinaryData, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) provider_name: Mapped[str] = mapped_column(String(255), nullable=False, server_default=sa.text("''")) def set_embedding(self, embedding_data: list[float]): self.embedding = pickle.dumps(embedding_data, protocol=pickle.HIGHEST_PROTOCOL) def get_embedding(self) -> list[float]: return cast(list[float], pickle.loads(self.embedding)) # noqa: S301 class DatasetCollectionBinding(TypeBase): __tablename__ = "dataset_collection_bindings" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="dataset_collection_bindings_pkey"), sa.Index("provider_model_name_idx", "provider_name", "model_name"), ) id: Mapped[str] = mapped_column( StringUUID, primary_key=True, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False, ) provider_name: Mapped[str] = mapped_column(String(255), nullable=False) model_name: Mapped[str] = mapped_column(String(255), nullable=False) type: Mapped[str] = mapped_column( EnumText(CollectionBindingType, length=40), server_default=sa.text("'dataset'"), nullable=False ) collection_name: Mapped[str] = mapped_column(String(64), nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) class TidbAuthBinding(TypeBase): __tablename__ = "tidb_auth_bindings" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="tidb_auth_bindings_pkey"), sa.Index("tidb_auth_bindings_tenant_idx", "tenant_id"), sa.Index("tidb_auth_bindings_active_idx", "active"), sa.Index("tidb_auth_bindings_created_at_idx", "created_at"), sa.Index("tidb_auth_bindings_status_idx", "status"), ) id: Mapped[str] = mapped_column( StringUUID, primary_key=True, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False, ) tenant_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) cluster_id: Mapped[str] = mapped_column(String(255), nullable=False) cluster_name: Mapped[str] = mapped_column(String(255), nullable=False) active: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.false()) status: Mapped[TidbAuthBindingStatus] = mapped_column( EnumText(TidbAuthBindingStatus, length=255), nullable=False, server_default=sa.text("'CREATING'") ) account: Mapped[str] = mapped_column(String(255), nullable=False) password: Mapped[str] = mapped_column(String(255), nullable=False) qdrant_endpoint: Mapped[str | None] = mapped_column(String(512), nullable=True, default=None) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) class Whitelist(TypeBase): __tablename__ = "whitelists" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="whitelists_pkey"), sa.Index("whitelists_tenant_idx", "tenant_id"), ) id: Mapped[str] = mapped_column( StringUUID, primary_key=True, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False, ) tenant_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True) category: Mapped[str] = mapped_column(String(255), nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) class DatasetPermission(TypeBase): __tablename__ = "dataset_permissions" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="dataset_permission_pkey"), sa.Index("idx_dataset_permissions_dataset_id", "dataset_id"), sa.Index("idx_dataset_permissions_account_id", "account_id"), sa.Index("idx_dataset_permissions_tenant_id", "tenant_id"), ) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), primary_key=True, init=False, ) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) account_id: Mapped[str] = mapped_column(StringUUID, nullable=False) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) has_permission: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=True) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) class ExternalKnowledgeApis(TypeBase): __tablename__ = "external_knowledge_apis" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="external_knowledge_apis_pkey"), sa.Index("external_knowledge_apis_tenant_idx", "tenant_id"), sa.Index("external_knowledge_apis_name_idx", "name"), ) id: Mapped[str] = mapped_column( StringUUID, nullable=False, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False, ) name: Mapped[str] = mapped_column(String(255), nullable=False) description: Mapped[str] = mapped_column(String(255), nullable=False) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) settings: Mapped[str | None] = mapped_column(LongText, nullable=True) created_by: Mapped[str] = mapped_column(StringUUID, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True) updated_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False ) @property def settings_dict(self) -> dict[str, Any] | None: try: return json.loads(self.settings) if self.settings else None except JSONDecodeError: return None def get_dataset_bindings(self, *, session: Session) -> list[DatasetBindingItem]: external_knowledge_bindings = session.scalars( select(ExternalKnowledgeBindings).where(ExternalKnowledgeBindings.external_knowledge_api_id == self.id) ).all() dataset_ids = [binding.dataset_id for binding in external_knowledge_bindings] datasets = session.scalars(select(Dataset).where(Dataset.id.in_(dataset_ids))).all() dataset_bindings: list[DatasetBindingItem] = [] for dataset in datasets: dataset_bindings.append({"id": dataset.id, "name": dataset.name}) return dataset_bindings class ExternalKnowledgeBindings(TypeBase): __tablename__ = "external_knowledge_bindings" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="external_knowledge_bindings_pkey"), sa.Index("external_knowledge_bindings_tenant_idx", "tenant_id"), sa.Index("external_knowledge_bindings_dataset_idx", "dataset_id"), sa.Index("external_knowledge_bindings_external_knowledge_idx", "external_knowledge_id"), sa.Index("external_knowledge_bindings_external_knowledge_api_idx", "external_knowledge_api_id"), ) id: Mapped[str] = mapped_column( StringUUID, nullable=False, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False, ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) external_knowledge_api_id: Mapped[str] = mapped_column(StringUUID, nullable=False) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) external_knowledge_id: Mapped[str] = mapped_column(String(512), nullable=False) created_by: Mapped[str] = mapped_column(StringUUID, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None, init=False) updated_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False ) class DatasetAutoDisableLog(TypeBase): __tablename__ = "dataset_auto_disable_logs" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="dataset_auto_disable_log_pkey"), sa.Index("dataset_auto_disable_log_tenant_idx", "tenant_id"), sa.Index("dataset_auto_disable_log_dataset_idx", "dataset_id"), sa.Index("dataset_auto_disable_log_created_atx", "created_at"), ) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) document_id: Mapped[str] = mapped_column(StringUUID, nullable=False) notified: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False ) class RateLimitLog(TypeBase): __tablename__ = "rate_limit_logs" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="rate_limit_log_pkey"), sa.Index("rate_limit_log_tenant_idx", "tenant_id"), sa.Index("rate_limit_log_operation_idx", "operation"), ) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) subscription_plan: Mapped[str] = mapped_column(String(255), nullable=False) operation: Mapped[str] = mapped_column(String(255), nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) class DatasetMetadata(TypeBase): __tablename__ = "dataset_metadatas" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="dataset_metadata_pkey"), sa.Index("dataset_metadata_tenant_idx", "tenant_id"), sa.Index("dataset_metadata_dataset_idx", "dataset_id"), ) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) type: Mapped[str] = mapped_column(EnumText(DatasetMetadataType, length=255), nullable=False) name: Mapped[str] = mapped_column(String(255), nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False ) updated_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=sa.func.current_timestamp(), onupdate=func.current_timestamp(), init=False, ) created_by: Mapped[str] = mapped_column(StringUUID, nullable=False) updated_by: Mapped[str] = mapped_column(StringUUID, nullable=True, default=None) class DatasetMetadataBinding(TypeBase): __tablename__ = "dataset_metadata_bindings" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="dataset_metadata_binding_pkey"), sa.Index("dataset_metadata_binding_tenant_idx", "tenant_id"), sa.Index("dataset_metadata_binding_dataset_idx", "dataset_id"), sa.Index("dataset_metadata_binding_metadata_idx", "metadata_id"), sa.Index("dataset_metadata_binding_document_idx", "document_id"), ) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) metadata_id: Mapped[str] = mapped_column(StringUUID, nullable=False) document_id: Mapped[str] = mapped_column(StringUUID, nullable=False) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) created_by: Mapped[str] = mapped_column(StringUUID, nullable=False) class PipelineBuiltInTemplate(TypeBase): __tablename__ = "pipeline_built_in_templates" __table_args__ = (sa.PrimaryKeyConstraint("id", name="pipeline_built_in_template_pkey"),) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False ) name: Mapped[str] = mapped_column(sa.String(255), nullable=False) description: Mapped[str] = mapped_column(LongText, nullable=False) chunk_structure: Mapped[str] = mapped_column(sa.String(255), nullable=False) icon: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False) yaml_content: Mapped[str] = mapped_column(LongText, nullable=False) copyright: Mapped[str] = mapped_column(sa.String(255), nullable=False) privacy_policy: Mapped[str] = mapped_column(sa.String(255), nullable=False) position: Mapped[int] = mapped_column(sa.Integer, nullable=False) install_count: Mapped[int] = mapped_column(sa.Integer, nullable=False) language: Mapped[str] = mapped_column(sa.String(255), nullable=False) created_at: Mapped[datetime] = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) updated_at: Mapped[datetime] = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False, ) class PipelineCustomizedTemplate(TypeBase): __tablename__ = "pipeline_customized_templates" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="pipeline_customized_template_pkey"), sa.Index("pipeline_customized_template_tenant_idx", "tenant_id"), ) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) name: Mapped[str] = mapped_column(sa.String(255), nullable=False) description: Mapped[str] = mapped_column(LongText, nullable=False) chunk_structure: Mapped[str] = mapped_column(sa.String(255), nullable=False) icon: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False) position: Mapped[int] = mapped_column(sa.Integer, nullable=False) yaml_content: Mapped[str] = mapped_column(LongText, nullable=False) install_count: Mapped[int] = mapped_column(sa.Integer, nullable=False) language: Mapped[str] = mapped_column(sa.String(255), nullable=False) created_by: Mapped[str] = mapped_column(StringUUID, nullable=False) updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None, init=False) created_at: Mapped[datetime] = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) updated_at: Mapped[datetime] = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False, ) def created_user_name(self, session: Session) -> str: account = session.scalar(select(Account).where(Account.id == self.created_by)) if account: return account.name return "" class Pipeline(TypeBase): __tablename__ = "pipelines" __table_args__ = (sa.PrimaryKeyConstraint("id", name="pipeline_pkey"),) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) name: Mapped[str] = mapped_column(sa.String(255), nullable=False) description: Mapped[str] = mapped_column(LongText, nullable=False, default=sa.text("''")) workflow_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) is_public: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=False) is_published: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=False) created_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) created_at: Mapped[datetime] = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) updated_at: Mapped[datetime] = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False, ) def retrieve_dataset(self, session: Session | scoped_session): return session.scalar( select(Dataset).where(Dataset.pipeline_id == self.id, Dataset.tenant_id == self.tenant_id) ) class DocumentPipelineExecutionLog(TypeBase): __tablename__ = "document_pipeline_execution_logs" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="document_pipeline_execution_log_pkey"), sa.Index("document_pipeline_execution_logs_document_id_idx", "document_id"), ) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False ) pipeline_id: Mapped[str] = mapped_column(StringUUID, nullable=False) document_id: Mapped[str] = mapped_column(StringUUID, nullable=False) datasource_type: Mapped[str] = mapped_column(sa.String(255), nullable=False) datasource_info: Mapped[str] = mapped_column(LongText, nullable=False) datasource_node_id: Mapped[str] = mapped_column(sa.String(255), nullable=False) input_data: Mapped[dict[str, Any]] = mapped_column(sa.JSON, nullable=False) created_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True) created_at: Mapped[datetime] = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) class PipelineRecommendedPlugin(TypeBase): __tablename__ = "pipeline_recommended_plugins" __table_args__ = (sa.PrimaryKeyConstraint("id", name="pipeline_recommended_plugin_pkey"),) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False ) plugin_id: Mapped[str] = mapped_column(LongText, nullable=False) provider_name: Mapped[str] = mapped_column(LongText, nullable=False) type: Mapped[str] = mapped_column(sa.String(50), nullable=False, server_default=sa.text("'tool'")) position: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0) active: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=True) created_at: Mapped[datetime] = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) updated_at: Mapped[datetime] = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False, ) class SegmentAttachmentBinding(TypeBase): __tablename__ = "segment_attachment_bindings" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="segment_attachment_binding_pkey"), sa.Index( "segment_attachment_binding_tenant_dataset_document_segment_idx", "tenant_id", "dataset_id", "document_id", "segment_id", ), sa.Index("segment_attachment_binding_attachment_idx", "attachment_id"), ) id: Mapped[str] = mapped_column( StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False ) tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) document_id: Mapped[str] = mapped_column(StringUUID, nullable=False) segment_id: Mapped[str] = mapped_column(StringUUID, nullable=False) attachment_id: Mapped[str] = mapped_column(StringUUID, nullable=False) created_at: Mapped[datetime] = mapped_column( sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) class DocumentSegmentSummary(TypeBase): __tablename__ = "document_segment_summaries" __table_args__ = ( sa.PrimaryKeyConstraint("id", name="document_segment_summaries_pkey"), sa.Index("document_segment_summaries_dataset_id_idx", "dataset_id"), sa.Index("document_segment_summaries_document_id_idx", "document_id"), sa.Index("document_segment_summaries_chunk_id_idx", "chunk_id"), sa.Index("document_segment_summaries_status_idx", "status"), ) id: Mapped[str] = mapped_column( StringUUID, nullable=False, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False, ) dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False) document_id: Mapped[str] = mapped_column(StringUUID, nullable=False) # corresponds to DocumentSegment.id or parent chunk id chunk_id: Mapped[str] = mapped_column(StringUUID, nullable=False) summary_content: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) summary_index_node_id: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None) summary_index_node_hash: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None) tokens: Mapped[int | None] = mapped_column(sa.Integer, nullable=True, default=None) status: Mapped[SummaryStatus] = mapped_column( EnumText(SummaryStatus, length=32), nullable=False, default=SummaryStatus.GENERATING, ) error: Mapped[str | None] = mapped_column(LongText, nullable=True, default=None) enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, default=True) disabled_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None) disabled_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None) created_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), init=False ) updated_at: Mapped[datetime] = mapped_column( DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp(), init=False, ) @override def __repr__(self): return f""