You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

dataset.py 55KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321
  1. import base64
  2. import enum
  3. import hashlib
  4. import hmac
  5. import json
  6. import logging
  7. import os
  8. import pickle
  9. import re
  10. import time
  11. from datetime import datetime
  12. from json import JSONDecodeError
  13. from typing import Any, Optional, cast
  14. import sqlalchemy as sa
  15. from sqlalchemy import DateTime, String, func, select
  16. from sqlalchemy.dialects.postgresql import JSONB
  17. from sqlalchemy.orm import Mapped, mapped_column
  18. from configs import dify_config
  19. from core.rag.index_processor.constant.built_in_field import BuiltInField, MetadataDataSource
  20. from core.rag.retrieval.retrieval_methods import RetrievalMethod
  21. from extensions.ext_storage import storage
  22. from services.entities.knowledge_entities.knowledge_entities import ParentMode, Rule
  23. from .account import Account
  24. from .base import Base
  25. from .engine import db
  26. from .model import App, Tag, TagBinding, UploadFile
  27. from .types import StringUUID
  28. logger = logging.getLogger(__name__)
  29. class DatasetPermissionEnum(enum.StrEnum):
  30. ONLY_ME = "only_me"
  31. ALL_TEAM = "all_team_members"
  32. PARTIAL_TEAM = "partial_members"
  33. class Dataset(Base):
  34. __tablename__ = "datasets"
  35. __table_args__ = (
  36. sa.PrimaryKeyConstraint("id", name="dataset_pkey"),
  37. sa.Index("dataset_tenant_idx", "tenant_id"),
  38. sa.Index("retrieval_model_idx", "retrieval_model", postgresql_using="gin"),
  39. )
  40. INDEXING_TECHNIQUE_LIST = ["high_quality", "economy", None]
  41. PROVIDER_LIST = ["vendor", "external", None]
  42. id = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
  43. tenant_id: Mapped[str] = mapped_column(StringUUID)
  44. name: Mapped[str] = mapped_column(String(255))
  45. description = mapped_column(sa.Text, nullable=True)
  46. provider: Mapped[str] = mapped_column(String(255), server_default=sa.text("'vendor'::character varying"))
  47. permission: Mapped[str] = mapped_column(String(255), server_default=sa.text("'only_me'::character varying"))
  48. data_source_type = mapped_column(String(255))
  49. indexing_technique: Mapped[Optional[str]] = mapped_column(String(255))
  50. index_struct = mapped_column(sa.Text, nullable=True)
  51. created_by = mapped_column(StringUUID, nullable=False)
  52. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  53. updated_by = mapped_column(StringUUID, nullable=True)
  54. updated_at = mapped_column(db.DateTime, nullable=False, server_default=func.current_timestamp())
  55. embedding_model = mapped_column(db.String(255), nullable=True)
  56. embedding_model_provider = mapped_column(db.String(255), nullable=True)
  57. keyword_number = db.Column(db.Integer, nullable=True, server_default=db.text("10"))
  58. collection_binding_id = mapped_column(StringUUID, nullable=True)
  59. retrieval_model = mapped_column(JSONB, nullable=True)
  60. built_in_field_enabled = mapped_column(db.Boolean, nullable=False, server_default=db.text("false"))
  61. icon_info = db.Column(JSONB, nullable=True)
  62. runtime_mode = db.Column(db.String(255), nullable=True, server_default=db.text("'general'::character varying"))
  63. pipeline_id = db.Column(StringUUID, nullable=True)
  64. chunk_structure = db.Column(db.String(255), nullable=True)
  65. @property
  66. def total_documents(self):
  67. return db.session.query(func.count(Document.id)).filter(Document.dataset_id == self.id).scalar()
  68. @property
  69. def total_available_documents(self):
  70. return (
  71. db.session.query(func.count(Document.id))
  72. .filter(
  73. Document.dataset_id == self.id,
  74. Document.indexing_status == "completed",
  75. Document.enabled == True,
  76. Document.archived == False,
  77. )
  78. .scalar()
  79. )
  80. @property
  81. def dataset_keyword_table(self):
  82. dataset_keyword_table = (
  83. db.session.query(DatasetKeywordTable).where(DatasetKeywordTable.dataset_id == self.id).first()
  84. )
  85. if dataset_keyword_table:
  86. return dataset_keyword_table
  87. return None
  88. @property
  89. def index_struct_dict(self):
  90. return json.loads(self.index_struct) if self.index_struct else None
  91. @property
  92. def external_retrieval_model(self):
  93. default_retrieval_model = {
  94. "top_k": 2,
  95. "score_threshold": 0.0,
  96. }
  97. return self.retrieval_model or default_retrieval_model
  98. @property
  99. def created_by_account(self):
  100. return db.session.get(Account, self.created_by)
  101. @property
  102. def latest_process_rule(self):
  103. return (
  104. db.session.query(DatasetProcessRule)
  105. .where(DatasetProcessRule.dataset_id == self.id)
  106. .order_by(DatasetProcessRule.created_at.desc())
  107. .first()
  108. )
  109. @property
  110. def app_count(self):
  111. return (
  112. db.session.query(func.count(AppDatasetJoin.id))
  113. .where(AppDatasetJoin.dataset_id == self.id, App.id == AppDatasetJoin.app_id)
  114. .scalar()
  115. )
  116. @property
  117. def document_count(self):
  118. return db.session.query(func.count(Document.id)).where(Document.dataset_id == self.id).scalar()
  119. @property
  120. def available_document_count(self):
  121. return (
  122. db.session.query(func.count(Document.id))
  123. .where(
  124. Document.dataset_id == self.id,
  125. Document.indexing_status == "completed",
  126. Document.enabled == True,
  127. Document.archived == False,
  128. )
  129. .scalar()
  130. )
  131. @property
  132. def available_segment_count(self):
  133. return (
  134. db.session.query(func.count(DocumentSegment.id))
  135. .where(
  136. DocumentSegment.dataset_id == self.id,
  137. DocumentSegment.status == "completed",
  138. DocumentSegment.enabled == True,
  139. )
  140. .scalar()
  141. )
  142. @property
  143. def word_count(self):
  144. return (
  145. db.session.query(Document)
  146. .with_entities(func.coalesce(func.sum(Document.word_count), 0))
  147. .where(Document.dataset_id == self.id)
  148. .scalar()
  149. )
  150. @property
  151. def doc_form(self) -> Optional[str]:
  152. if self.chunk_structure:
  153. return self.chunk_structure
  154. document = db.session.query(Document).filter(Document.dataset_id == self.id).first()
  155. if document:
  156. return document.doc_form
  157. return None
  158. @property
  159. def retrieval_model_dict(self):
  160. default_retrieval_model = {
  161. "search_method": RetrievalMethod.SEMANTIC_SEARCH.value,
  162. "reranking_enable": False,
  163. "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},
  164. "top_k": 2,
  165. "score_threshold_enabled": False,
  166. }
  167. return self.retrieval_model or default_retrieval_model
  168. @property
  169. def tags(self):
  170. tags = (
  171. db.session.query(Tag)
  172. .join(TagBinding, Tag.id == TagBinding.tag_id)
  173. .where(
  174. TagBinding.target_id == self.id,
  175. TagBinding.tenant_id == self.tenant_id,
  176. Tag.tenant_id == self.tenant_id,
  177. Tag.type == "knowledge",
  178. )
  179. .all()
  180. )
  181. return tags or []
  182. @property
  183. def external_knowledge_info(self):
  184. if self.provider != "external":
  185. return None
  186. external_knowledge_binding = (
  187. db.session.query(ExternalKnowledgeBindings).where(ExternalKnowledgeBindings.dataset_id == self.id).first()
  188. )
  189. if not external_knowledge_binding:
  190. return None
  191. external_knowledge_api = db.session.scalar(
  192. select(ExternalKnowledgeApis).where(
  193. ExternalKnowledgeApis.id == external_knowledge_binding.external_knowledge_api_id
  194. )
  195. )
  196. if not external_knowledge_api:
  197. return None
  198. return {
  199. "external_knowledge_id": external_knowledge_binding.external_knowledge_id,
  200. "external_knowledge_api_id": external_knowledge_api.id,
  201. "external_knowledge_api_name": external_knowledge_api.name,
  202. "external_knowledge_api_endpoint": json.loads(external_knowledge_api.settings).get("endpoint", ""),
  203. }
  204. @property
  205. def is_published(self):
  206. if self.pipeline_id:
  207. pipeline = db.session.query(Pipeline).filter(Pipeline.id == self.pipeline_id).first()
  208. if pipeline:
  209. return pipeline.is_published
  210. return False
  211. @property
  212. def doc_metadata(self):
  213. dataset_metadatas = db.session.query(DatasetMetadata).where(DatasetMetadata.dataset_id == self.id).all()
  214. doc_metadata = [
  215. {
  216. "id": dataset_metadata.id,
  217. "name": dataset_metadata.name,
  218. "type": dataset_metadata.type,
  219. }
  220. for dataset_metadata in dataset_metadatas
  221. ]
  222. if self.built_in_field_enabled:
  223. doc_metadata.append(
  224. {
  225. "id": "built-in",
  226. "name": BuiltInField.document_name.value,
  227. "type": "string",
  228. }
  229. )
  230. doc_metadata.append(
  231. {
  232. "id": "built-in",
  233. "name": BuiltInField.uploader.value,
  234. "type": "string",
  235. }
  236. )
  237. doc_metadata.append(
  238. {
  239. "id": "built-in",
  240. "name": BuiltInField.upload_date.value,
  241. "type": "time",
  242. }
  243. )
  244. doc_metadata.append(
  245. {
  246. "id": "built-in",
  247. "name": BuiltInField.last_update_date.value,
  248. "type": "time",
  249. }
  250. )
  251. doc_metadata.append(
  252. {
  253. "id": "built-in",
  254. "name": BuiltInField.source.value,
  255. "type": "string",
  256. }
  257. )
  258. return doc_metadata
  259. @staticmethod
  260. def gen_collection_name_by_id(dataset_id: str) -> str:
  261. normalized_dataset_id = dataset_id.replace("-", "_")
  262. return f"{dify_config.VECTOR_INDEX_NAME_PREFIX}_{normalized_dataset_id}_Node"
  263. class DatasetProcessRule(Base):
  264. __tablename__ = "dataset_process_rules"
  265. __table_args__ = (
  266. sa.PrimaryKeyConstraint("id", name="dataset_process_rule_pkey"),
  267. sa.Index("dataset_process_rule_dataset_id_idx", "dataset_id"),
  268. )
  269. id = mapped_column(StringUUID, nullable=False, server_default=sa.text("uuid_generate_v4()"))
  270. dataset_id = mapped_column(StringUUID, nullable=False)
  271. mode = mapped_column(String(255), nullable=False, server_default=sa.text("'automatic'::character varying"))
  272. rules = mapped_column(sa.Text, nullable=True)
  273. created_by = mapped_column(StringUUID, nullable=False)
  274. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  275. MODES = ["automatic", "custom", "hierarchical"]
  276. PRE_PROCESSING_RULES = ["remove_stopwords", "remove_extra_spaces", "remove_urls_emails"]
  277. AUTOMATIC_RULES: dict[str, Any] = {
  278. "pre_processing_rules": [
  279. {"id": "remove_extra_spaces", "enabled": True},
  280. {"id": "remove_urls_emails", "enabled": False},
  281. ],
  282. "segmentation": {"delimiter": "\n", "max_tokens": 500, "chunk_overlap": 50},
  283. }
  284. def to_dict(self):
  285. return {
  286. "id": self.id,
  287. "dataset_id": self.dataset_id,
  288. "mode": self.mode,
  289. "rules": self.rules_dict,
  290. }
  291. @property
  292. def rules_dict(self):
  293. try:
  294. return json.loads(self.rules) if self.rules else None
  295. except JSONDecodeError:
  296. return None
  297. class Document(Base):
  298. __tablename__ = "documents"
  299. __table_args__ = (
  300. sa.PrimaryKeyConstraint("id", name="document_pkey"),
  301. sa.Index("document_dataset_id_idx", "dataset_id"),
  302. sa.Index("document_is_paused_idx", "is_paused"),
  303. sa.Index("document_tenant_idx", "tenant_id"),
  304. sa.Index("document_metadata_idx", "doc_metadata", postgresql_using="gin"),
  305. )
  306. # initial fields
  307. id = mapped_column(StringUUID, nullable=False, server_default=sa.text("uuid_generate_v4()"))
  308. tenant_id = mapped_column(StringUUID, nullable=False)
  309. dataset_id = mapped_column(StringUUID, nullable=False)
  310. position: Mapped[int] = mapped_column(sa.Integer, nullable=False)
  311. data_source_type: Mapped[str] = mapped_column(String(255), nullable=False)
  312. data_source_info = mapped_column(sa.Text, nullable=True)
  313. dataset_process_rule_id = mapped_column(StringUUID, nullable=True)
  314. batch: Mapped[str] = mapped_column(String(255), nullable=False)
  315. name: Mapped[str] = mapped_column(String(255), nullable=False)
  316. created_from: Mapped[str] = mapped_column(String(255), nullable=False)
  317. created_by = mapped_column(StringUUID, nullable=False)
  318. created_api_request_id = mapped_column(StringUUID, nullable=True)
  319. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  320. # start processing
  321. processing_started_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  322. # parsing
  323. file_id = mapped_column(sa.Text, nullable=True)
  324. word_count: Mapped[Optional[int]] = mapped_column(sa.Integer, nullable=True) # TODO: make this not nullable
  325. parsing_completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  326. # cleaning
  327. cleaning_completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  328. # split
  329. splitting_completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  330. # indexing
  331. tokens: Mapped[Optional[int]] = mapped_column(sa.Integer, nullable=True)
  332. indexing_latency: Mapped[Optional[float]] = mapped_column(sa.Float, nullable=True)
  333. completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  334. # pause
  335. is_paused: Mapped[Optional[bool]] = mapped_column(sa.Boolean, nullable=True, server_default=sa.text("false"))
  336. paused_by = mapped_column(StringUUID, nullable=True)
  337. paused_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  338. # error
  339. error = mapped_column(sa.Text, nullable=True)
  340. stopped_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  341. # basic fields
  342. indexing_status = mapped_column(String(255), nullable=False, server_default=sa.text("'waiting'::character varying"))
  343. enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("true"))
  344. disabled_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  345. disabled_by = mapped_column(StringUUID, nullable=True)
  346. archived: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
  347. archived_reason = mapped_column(String(255), nullable=True)
  348. archived_by = mapped_column(StringUUID, nullable=True)
  349. archived_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  350. updated_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  351. doc_type = mapped_column(String(40), nullable=True)
  352. doc_metadata = mapped_column(JSONB, nullable=True)
  353. doc_form = mapped_column(String(255), nullable=False, server_default=sa.text("'text_model'::character varying"))
  354. doc_language = mapped_column(String(255), nullable=True)
  355. DATA_SOURCES = ["upload_file", "notion_import", "website_crawl"]
  356. @property
  357. def display_status(self):
  358. status = None
  359. if self.indexing_status == "waiting":
  360. status = "queuing"
  361. elif self.indexing_status not in {"completed", "error", "waiting"} and self.is_paused:
  362. status = "paused"
  363. elif self.indexing_status in {"parsing", "cleaning", "splitting", "indexing"}:
  364. status = "indexing"
  365. elif self.indexing_status == "error":
  366. status = "error"
  367. elif self.indexing_status == "completed" and not self.archived and self.enabled:
  368. status = "available"
  369. elif self.indexing_status == "completed" and not self.archived and not self.enabled:
  370. status = "disabled"
  371. elif self.indexing_status == "completed" and self.archived:
  372. status = "archived"
  373. return status
  374. @property
  375. def data_source_info_dict(self) -> dict[str, Any]:
  376. if self.data_source_info:
  377. try:
  378. data_source_info_dict = json.loads(self.data_source_info)
  379. except JSONDecodeError:
  380. data_source_info_dict = {}
  381. return data_source_info_dict
  382. return {}
  383. @property
  384. def data_source_detail_dict(self):
  385. if self.data_source_info:
  386. if self.data_source_type == "upload_file":
  387. data_source_info_dict = json.loads(self.data_source_info)
  388. file_detail = (
  389. db.session.query(UploadFile)
  390. .where(UploadFile.id == data_source_info_dict["upload_file_id"])
  391. .one_or_none()
  392. )
  393. if file_detail:
  394. return {
  395. "upload_file": {
  396. "id": file_detail.id,
  397. "name": file_detail.name,
  398. "size": file_detail.size,
  399. "extension": file_detail.extension,
  400. "mime_type": file_detail.mime_type,
  401. "created_by": file_detail.created_by,
  402. "created_at": file_detail.created_at.timestamp(),
  403. }
  404. }
  405. elif self.data_source_type in {"notion_import", "website_crawl"}:
  406. return json.loads(self.data_source_info)
  407. return {}
  408. @property
  409. def average_segment_length(self):
  410. if self.word_count and self.word_count != 0 and self.segment_count and self.segment_count != 0:
  411. return self.word_count // self.segment_count
  412. return 0
  413. @property
  414. def dataset_process_rule(self):
  415. if self.dataset_process_rule_id:
  416. return db.session.get(DatasetProcessRule, self.dataset_process_rule_id)
  417. return None
  418. @property
  419. def dataset(self):
  420. return db.session.query(Dataset).where(Dataset.id == self.dataset_id).one_or_none()
  421. @property
  422. def segment_count(self):
  423. return db.session.query(DocumentSegment).where(DocumentSegment.document_id == self.id).count()
  424. @property
  425. def hit_count(self):
  426. return (
  427. db.session.query(DocumentSegment)
  428. .with_entities(func.coalesce(func.sum(DocumentSegment.hit_count), 0))
  429. .where(DocumentSegment.document_id == self.id)
  430. .scalar()
  431. )
  432. @property
  433. def uploader(self):
  434. user = db.session.query(Account).where(Account.id == self.created_by).first()
  435. return user.name if user else None
  436. @property
  437. def upload_date(self):
  438. return self.created_at
  439. @property
  440. def last_update_date(self):
  441. return self.updated_at
  442. @property
  443. def doc_metadata_details(self):
  444. if self.doc_metadata:
  445. document_metadatas = (
  446. db.session.query(DatasetMetadata)
  447. .join(DatasetMetadataBinding, DatasetMetadataBinding.metadata_id == DatasetMetadata.id)
  448. .where(
  449. DatasetMetadataBinding.dataset_id == self.dataset_id, DatasetMetadataBinding.document_id == self.id
  450. )
  451. .all()
  452. )
  453. metadata_list = []
  454. for metadata in document_metadatas:
  455. metadata_dict = {
  456. "id": metadata.id,
  457. "name": metadata.name,
  458. "type": metadata.type,
  459. "value": self.doc_metadata.get(metadata.name),
  460. }
  461. metadata_list.append(metadata_dict)
  462. # deal built-in fields
  463. metadata_list.extend(self.get_built_in_fields())
  464. return metadata_list
  465. return None
  466. @property
  467. def process_rule_dict(self):
  468. if self.dataset_process_rule_id:
  469. return self.dataset_process_rule.to_dict()
  470. return None
  471. def get_built_in_fields(self):
  472. built_in_fields = []
  473. built_in_fields.append(
  474. {
  475. "id": "built-in",
  476. "name": BuiltInField.document_name,
  477. "type": "string",
  478. "value": self.name,
  479. }
  480. )
  481. built_in_fields.append(
  482. {
  483. "id": "built-in",
  484. "name": BuiltInField.uploader,
  485. "type": "string",
  486. "value": self.uploader,
  487. }
  488. )
  489. built_in_fields.append(
  490. {
  491. "id": "built-in",
  492. "name": BuiltInField.upload_date,
  493. "type": "time",
  494. "value": str(self.created_at.timestamp()),
  495. }
  496. )
  497. built_in_fields.append(
  498. {
  499. "id": "built-in",
  500. "name": BuiltInField.last_update_date,
  501. "type": "time",
  502. "value": str(self.updated_at.timestamp()),
  503. }
  504. )
  505. built_in_fields.append(
  506. {
  507. "id": "built-in",
  508. "name": BuiltInField.source,
  509. "type": "string",
  510. "value": MetadataDataSource[self.data_source_type].value,
  511. }
  512. )
  513. return built_in_fields
  514. def to_dict(self):
  515. return {
  516. "id": self.id,
  517. "tenant_id": self.tenant_id,
  518. "dataset_id": self.dataset_id,
  519. "position": self.position,
  520. "data_source_type": self.data_source_type,
  521. "data_source_info": self.data_source_info,
  522. "dataset_process_rule_id": self.dataset_process_rule_id,
  523. "batch": self.batch,
  524. "name": self.name,
  525. "created_from": self.created_from,
  526. "created_by": self.created_by,
  527. "created_api_request_id": self.created_api_request_id,
  528. "created_at": self.created_at,
  529. "processing_started_at": self.processing_started_at,
  530. "file_id": self.file_id,
  531. "word_count": self.word_count,
  532. "parsing_completed_at": self.parsing_completed_at,
  533. "cleaning_completed_at": self.cleaning_completed_at,
  534. "splitting_completed_at": self.splitting_completed_at,
  535. "tokens": self.tokens,
  536. "indexing_latency": self.indexing_latency,
  537. "completed_at": self.completed_at,
  538. "is_paused": self.is_paused,
  539. "paused_by": self.paused_by,
  540. "paused_at": self.paused_at,
  541. "error": self.error,
  542. "stopped_at": self.stopped_at,
  543. "indexing_status": self.indexing_status,
  544. "enabled": self.enabled,
  545. "disabled_at": self.disabled_at,
  546. "disabled_by": self.disabled_by,
  547. "archived": self.archived,
  548. "archived_reason": self.archived_reason,
  549. "archived_by": self.archived_by,
  550. "archived_at": self.archived_at,
  551. "updated_at": self.updated_at,
  552. "doc_type": self.doc_type,
  553. "doc_metadata": self.doc_metadata,
  554. "doc_form": self.doc_form,
  555. "doc_language": self.doc_language,
  556. "display_status": self.display_status,
  557. "data_source_info_dict": self.data_source_info_dict,
  558. "average_segment_length": self.average_segment_length,
  559. "dataset_process_rule": self.dataset_process_rule.to_dict() if self.dataset_process_rule else None,
  560. "dataset": self.dataset.to_dict() if self.dataset else None,
  561. "segment_count": self.segment_count,
  562. "hit_count": self.hit_count,
  563. }
  564. @classmethod
  565. def from_dict(cls, data: dict):
  566. return cls(
  567. id=data.get("id"),
  568. tenant_id=data.get("tenant_id"),
  569. dataset_id=data.get("dataset_id"),
  570. position=data.get("position"),
  571. data_source_type=data.get("data_source_type"),
  572. data_source_info=data.get("data_source_info"),
  573. dataset_process_rule_id=data.get("dataset_process_rule_id"),
  574. batch=data.get("batch"),
  575. name=data.get("name"),
  576. created_from=data.get("created_from"),
  577. created_by=data.get("created_by"),
  578. created_api_request_id=data.get("created_api_request_id"),
  579. created_at=data.get("created_at"),
  580. processing_started_at=data.get("processing_started_at"),
  581. file_id=data.get("file_id"),
  582. word_count=data.get("word_count"),
  583. parsing_completed_at=data.get("parsing_completed_at"),
  584. cleaning_completed_at=data.get("cleaning_completed_at"),
  585. splitting_completed_at=data.get("splitting_completed_at"),
  586. tokens=data.get("tokens"),
  587. indexing_latency=data.get("indexing_latency"),
  588. completed_at=data.get("completed_at"),
  589. is_paused=data.get("is_paused"),
  590. paused_by=data.get("paused_by"),
  591. paused_at=data.get("paused_at"),
  592. error=data.get("error"),
  593. stopped_at=data.get("stopped_at"),
  594. indexing_status=data.get("indexing_status"),
  595. enabled=data.get("enabled"),
  596. disabled_at=data.get("disabled_at"),
  597. disabled_by=data.get("disabled_by"),
  598. archived=data.get("archived"),
  599. archived_reason=data.get("archived_reason"),
  600. archived_by=data.get("archived_by"),
  601. archived_at=data.get("archived_at"),
  602. updated_at=data.get("updated_at"),
  603. doc_type=data.get("doc_type"),
  604. doc_metadata=data.get("doc_metadata"),
  605. doc_form=data.get("doc_form"),
  606. doc_language=data.get("doc_language"),
  607. )
  608. class DocumentSegment(Base):
  609. __tablename__ = "document_segments"
  610. __table_args__ = (
  611. sa.PrimaryKeyConstraint("id", name="document_segment_pkey"),
  612. sa.Index("document_segment_dataset_id_idx", "dataset_id"),
  613. sa.Index("document_segment_document_id_idx", "document_id"),
  614. sa.Index("document_segment_tenant_dataset_idx", "dataset_id", "tenant_id"),
  615. sa.Index("document_segment_tenant_document_idx", "document_id", "tenant_id"),
  616. sa.Index("document_segment_node_dataset_idx", "index_node_id", "dataset_id"),
  617. sa.Index("document_segment_tenant_idx", "tenant_id"),
  618. )
  619. # initial fields
  620. id = mapped_column(StringUUID, nullable=False, server_default=sa.text("uuid_generate_v4()"))
  621. tenant_id = mapped_column(StringUUID, nullable=False)
  622. dataset_id = mapped_column(StringUUID, nullable=False)
  623. document_id = mapped_column(StringUUID, nullable=False)
  624. position: Mapped[int]
  625. content = mapped_column(sa.Text, nullable=False)
  626. answer = mapped_column(sa.Text, nullable=True)
  627. word_count: Mapped[int]
  628. tokens: Mapped[int]
  629. # indexing fields
  630. keywords = mapped_column(sa.JSON, nullable=True)
  631. index_node_id = mapped_column(String(255), nullable=True)
  632. index_node_hash = mapped_column(String(255), nullable=True)
  633. # basic fields
  634. hit_count: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0)
  635. enabled: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("true"))
  636. disabled_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  637. disabled_by = mapped_column(StringUUID, nullable=True)
  638. status: Mapped[str] = mapped_column(String(255), server_default=sa.text("'waiting'::character varying"))
  639. created_by = mapped_column(StringUUID, nullable=False)
  640. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  641. updated_by = mapped_column(StringUUID, nullable=True)
  642. updated_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  643. indexing_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  644. completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  645. error = mapped_column(sa.Text, nullable=True)
  646. stopped_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  647. @property
  648. def dataset(self):
  649. return db.session.scalar(select(Dataset).where(Dataset.id == self.dataset_id))
  650. @property
  651. def document(self):
  652. return db.session.scalar(select(Document).where(Document.id == self.document_id))
  653. @property
  654. def previous_segment(self):
  655. return db.session.scalar(
  656. select(DocumentSegment).where(
  657. DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position - 1
  658. )
  659. )
  660. @property
  661. def next_segment(self):
  662. return db.session.scalar(
  663. select(DocumentSegment).where(
  664. DocumentSegment.document_id == self.document_id, DocumentSegment.position == self.position + 1
  665. )
  666. )
  667. @property
  668. def child_chunks(self):
  669. process_rule = self.document.dataset_process_rule
  670. if process_rule.mode == "hierarchical":
  671. rules = Rule(**process_rule.rules_dict)
  672. if rules.parent_mode and rules.parent_mode != ParentMode.FULL_DOC:
  673. child_chunks = (
  674. db.session.query(ChildChunk)
  675. .where(ChildChunk.segment_id == self.id)
  676. .order_by(ChildChunk.position.asc())
  677. .all()
  678. )
  679. return child_chunks or []
  680. else:
  681. return []
  682. else:
  683. return []
  684. def get_child_chunks(self):
  685. process_rule = self.document.dataset_process_rule
  686. if process_rule.mode == "hierarchical":
  687. rules = Rule(**process_rule.rules_dict)
  688. if rules.parent_mode:
  689. child_chunks = (
  690. db.session.query(ChildChunk)
  691. .where(ChildChunk.segment_id == self.id)
  692. .order_by(ChildChunk.position.asc())
  693. .all()
  694. )
  695. return child_chunks or []
  696. else:
  697. return []
  698. else:
  699. return []
  700. @property
  701. def sign_content(self):
  702. return self.get_sign_content()
  703. def get_sign_content(self):
  704. signed_urls = []
  705. text = self.content
  706. # For data before v0.10.0
  707. pattern = r"/files/([a-f0-9\-]+)/image-preview(?:\?.*?)?"
  708. matches = re.finditer(pattern, text)
  709. for match in matches:
  710. upload_file_id = match.group(1)
  711. nonce = os.urandom(16).hex()
  712. timestamp = str(int(time.time()))
  713. data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}"
  714. secret_key = dify_config.SECRET_KEY.encode() if dify_config.SECRET_KEY else b""
  715. sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
  716. encoded_sign = base64.urlsafe_b64encode(sign).decode()
  717. params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
  718. base_url = f"/files/{upload_file_id}/image-preview"
  719. signed_url = f"{base_url}?{params}"
  720. signed_urls.append((match.start(), match.end(), signed_url))
  721. # For data after v0.10.0
  722. pattern = r"/files/([a-f0-9\-]+)/file-preview(?:\?.*?)?"
  723. matches = re.finditer(pattern, text)
  724. for match in matches:
  725. upload_file_id = match.group(1)
  726. nonce = os.urandom(16).hex()
  727. timestamp = str(int(time.time()))
  728. data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}"
  729. secret_key = dify_config.SECRET_KEY.encode() if dify_config.SECRET_KEY else b""
  730. sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
  731. encoded_sign = base64.urlsafe_b64encode(sign).decode()
  732. params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
  733. base_url = f"/files/{upload_file_id}/file-preview"
  734. signed_url = f"{base_url}?{params}"
  735. signed_urls.append((match.start(), match.end(), signed_url))
  736. # For tools directory - direct file formats (e.g., .png, .jpg, etc.)
  737. pattern = r"/files/tools/([a-f0-9\-]+)\.([a-zA-Z0-9]+)(?:\?.*?)?"
  738. matches = re.finditer(pattern, text)
  739. for match in matches:
  740. upload_file_id = match.group(1)
  741. file_extension = match.group(2)
  742. nonce = os.urandom(16).hex()
  743. timestamp = str(int(time.time()))
  744. data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}"
  745. secret_key = dify_config.SECRET_KEY.encode() if dify_config.SECRET_KEY else b""
  746. sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
  747. encoded_sign = base64.urlsafe_b64encode(sign).decode()
  748. params = f"timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
  749. base_url = f"/files/tools/{upload_file_id}.{file_extension}"
  750. signed_url = f"{base_url}?{params}"
  751. signed_urls.append((match.start(), match.end(), signed_url))
  752. # Reconstruct the text with signed URLs
  753. offset = 0
  754. for start, end, signed_url in signed_urls:
  755. text = text[: start + offset] + signed_url + text[end + offset :]
  756. offset += len(signed_url) - (end - start)
  757. return text
  758. class ChildChunk(Base):
  759. __tablename__ = "child_chunks"
  760. __table_args__ = (
  761. sa.PrimaryKeyConstraint("id", name="child_chunk_pkey"),
  762. sa.Index("child_chunk_dataset_id_idx", "tenant_id", "dataset_id", "document_id", "segment_id", "index_node_id"),
  763. sa.Index("child_chunks_node_idx", "index_node_id", "dataset_id"),
  764. sa.Index("child_chunks_segment_idx", "segment_id"),
  765. )
  766. # initial fields
  767. id = mapped_column(StringUUID, nullable=False, server_default=sa.text("uuid_generate_v4()"))
  768. tenant_id = mapped_column(StringUUID, nullable=False)
  769. dataset_id = mapped_column(StringUUID, nullable=False)
  770. document_id = mapped_column(StringUUID, nullable=False)
  771. segment_id = mapped_column(StringUUID, nullable=False)
  772. position: Mapped[int] = mapped_column(sa.Integer, nullable=False)
  773. content = mapped_column(sa.Text, nullable=False)
  774. word_count: Mapped[int] = mapped_column(sa.Integer, nullable=False)
  775. # indexing fields
  776. index_node_id = mapped_column(String(255), nullable=True)
  777. index_node_hash = mapped_column(String(255), nullable=True)
  778. type = mapped_column(String(255), nullable=False, server_default=sa.text("'automatic'::character varying"))
  779. created_by = mapped_column(StringUUID, nullable=False)
  780. created_at: Mapped[datetime] = mapped_column(
  781. DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
  782. )
  783. updated_by = mapped_column(StringUUID, nullable=True)
  784. updated_at: Mapped[datetime] = mapped_column(
  785. DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
  786. )
  787. indexing_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  788. completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime, nullable=True)
  789. error = mapped_column(sa.Text, nullable=True)
  790. @property
  791. def dataset(self):
  792. return db.session.query(Dataset).where(Dataset.id == self.dataset_id).first()
  793. @property
  794. def document(self):
  795. return db.session.query(Document).where(Document.id == self.document_id).first()
  796. @property
  797. def segment(self):
  798. return db.session.query(DocumentSegment).where(DocumentSegment.id == self.segment_id).first()
  799. class AppDatasetJoin(Base):
  800. __tablename__ = "app_dataset_joins"
  801. __table_args__ = (
  802. sa.PrimaryKeyConstraint("id", name="app_dataset_join_pkey"),
  803. sa.Index("app_dataset_join_app_dataset_idx", "dataset_id", "app_id"),
  804. )
  805. id = mapped_column(StringUUID, primary_key=True, nullable=False, server_default=sa.text("uuid_generate_v4()"))
  806. app_id = mapped_column(StringUUID, nullable=False)
  807. dataset_id = mapped_column(StringUUID, nullable=False)
  808. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=db.func.current_timestamp())
  809. @property
  810. def app(self):
  811. return db.session.get(App, self.app_id)
  812. class DatasetQuery(Base):
  813. __tablename__ = "dataset_queries"
  814. __table_args__ = (
  815. sa.PrimaryKeyConstraint("id", name="dataset_query_pkey"),
  816. sa.Index("dataset_query_dataset_id_idx", "dataset_id"),
  817. )
  818. id = mapped_column(StringUUID, primary_key=True, nullable=False, server_default=sa.text("uuid_generate_v4()"))
  819. dataset_id = mapped_column(StringUUID, nullable=False)
  820. content = mapped_column(sa.Text, nullable=False)
  821. source: Mapped[str] = mapped_column(String(255), nullable=False)
  822. source_app_id = mapped_column(StringUUID, nullable=True)
  823. created_by_role = mapped_column(String, nullable=False)
  824. created_by = mapped_column(StringUUID, nullable=False)
  825. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=db.func.current_timestamp())
  826. class DatasetKeywordTable(Base):
  827. __tablename__ = "dataset_keyword_tables"
  828. __table_args__ = (
  829. sa.PrimaryKeyConstraint("id", name="dataset_keyword_table_pkey"),
  830. sa.Index("dataset_keyword_table_dataset_id_idx", "dataset_id"),
  831. )
  832. id = mapped_column(StringUUID, primary_key=True, server_default=sa.text("uuid_generate_v4()"))
  833. dataset_id = mapped_column(StringUUID, nullable=False, unique=True)
  834. keyword_table = mapped_column(sa.Text, nullable=False)
  835. data_source_type = mapped_column(
  836. String(255), nullable=False, server_default=sa.text("'database'::character varying")
  837. )
  838. @property
  839. def keyword_table_dict(self):
  840. class SetDecoder(json.JSONDecoder):
  841. def __init__(self, *args, **kwargs):
  842. super().__init__(object_hook=self.object_hook, *args, **kwargs)
  843. def object_hook(self, dct):
  844. if isinstance(dct, dict):
  845. for keyword, node_idxs in dct.items():
  846. if isinstance(node_idxs, list):
  847. dct[keyword] = set(node_idxs)
  848. return dct
  849. # get dataset
  850. dataset = db.session.query(Dataset).filter_by(id=self.dataset_id).first()
  851. if not dataset:
  852. return None
  853. if self.data_source_type == "database":
  854. return json.loads(self.keyword_table, cls=SetDecoder) if self.keyword_table else None
  855. else:
  856. file_key = "keyword_files/" + dataset.tenant_id + "/" + self.dataset_id + ".txt"
  857. try:
  858. keyword_table_text = storage.load_once(file_key)
  859. if keyword_table_text:
  860. return json.loads(keyword_table_text.decode("utf-8"), cls=SetDecoder)
  861. return None
  862. except Exception as e:
  863. logger.exception("Failed to load keyword table from file: %s", file_key)
  864. return None
  865. class Embedding(Base):
  866. __tablename__ = "embeddings"
  867. __table_args__ = (
  868. sa.PrimaryKeyConstraint("id", name="embedding_pkey"),
  869. sa.UniqueConstraint("model_name", "hash", "provider_name", name="embedding_hash_idx"),
  870. sa.Index("created_at_idx", "created_at"),
  871. )
  872. id = mapped_column(StringUUID, primary_key=True, server_default=sa.text("uuid_generate_v4()"))
  873. model_name = mapped_column(
  874. String(255), nullable=False, server_default=sa.text("'text-embedding-ada-002'::character varying")
  875. )
  876. hash = mapped_column(String(64), nullable=False)
  877. embedding = mapped_column(sa.LargeBinary, nullable=False)
  878. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  879. provider_name = mapped_column(String(255), nullable=False, server_default=sa.text("''::character varying"))
  880. def set_embedding(self, embedding_data: list[float]):
  881. self.embedding = pickle.dumps(embedding_data, protocol=pickle.HIGHEST_PROTOCOL)
  882. def get_embedding(self) -> list[float]:
  883. return cast(list[float], pickle.loads(self.embedding)) # noqa: S301
  884. class DatasetCollectionBinding(Base):
  885. __tablename__ = "dataset_collection_bindings"
  886. __table_args__ = (
  887. sa.PrimaryKeyConstraint("id", name="dataset_collection_bindings_pkey"),
  888. sa.Index("provider_model_name_idx", "provider_name", "model_name"),
  889. )
  890. id = mapped_column(StringUUID, primary_key=True, server_default=sa.text("uuid_generate_v4()"))
  891. provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
  892. model_name: Mapped[str] = mapped_column(String(255), nullable=False)
  893. type = mapped_column(String(40), server_default=sa.text("'dataset'::character varying"), nullable=False)
  894. collection_name = mapped_column(String(64), nullable=False)
  895. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  896. class TidbAuthBinding(Base):
  897. __tablename__ = "tidb_auth_bindings"
  898. __table_args__ = (
  899. sa.PrimaryKeyConstraint("id", name="tidb_auth_bindings_pkey"),
  900. sa.Index("tidb_auth_bindings_tenant_idx", "tenant_id"),
  901. sa.Index("tidb_auth_bindings_active_idx", "active"),
  902. sa.Index("tidb_auth_bindings_created_at_idx", "created_at"),
  903. sa.Index("tidb_auth_bindings_status_idx", "status"),
  904. )
  905. id = mapped_column(StringUUID, primary_key=True, server_default=sa.text("uuid_generate_v4()"))
  906. tenant_id = mapped_column(StringUUID, nullable=True)
  907. cluster_id: Mapped[str] = mapped_column(String(255), nullable=False)
  908. cluster_name: Mapped[str] = mapped_column(String(255), nullable=False)
  909. active: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=db.text("false"))
  910. status = mapped_column(String(255), nullable=False, server_default=db.text("'CREATING'::character varying"))
  911. account: Mapped[str] = mapped_column(String(255), nullable=False)
  912. password: Mapped[str] = mapped_column(String(255), nullable=False)
  913. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  914. class Whitelist(Base):
  915. __tablename__ = "whitelists"
  916. __table_args__ = (
  917. sa.PrimaryKeyConstraint("id", name="whitelists_pkey"),
  918. sa.Index("whitelists_tenant_idx", "tenant_id"),
  919. )
  920. id = mapped_column(StringUUID, primary_key=True, server_default=sa.text("uuid_generate_v4()"))
  921. tenant_id = mapped_column(StringUUID, nullable=True)
  922. category: Mapped[str] = mapped_column(String(255), nullable=False)
  923. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  924. class DatasetPermission(Base):
  925. __tablename__ = "dataset_permissions"
  926. __table_args__ = (
  927. sa.PrimaryKeyConstraint("id", name="dataset_permission_pkey"),
  928. sa.Index("idx_dataset_permissions_dataset_id", "dataset_id"),
  929. sa.Index("idx_dataset_permissions_account_id", "account_id"),
  930. sa.Index("idx_dataset_permissions_tenant_id", "tenant_id"),
  931. )
  932. id = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"), primary_key=True)
  933. dataset_id = mapped_column(StringUUID, nullable=False)
  934. account_id = mapped_column(StringUUID, nullable=False)
  935. tenant_id = mapped_column(StringUUID, nullable=False)
  936. has_permission: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("true"))
  937. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  938. class ExternalKnowledgeApis(Base):
  939. __tablename__ = "external_knowledge_apis"
  940. __table_args__ = (
  941. sa.PrimaryKeyConstraint("id", name="external_knowledge_apis_pkey"),
  942. sa.Index("external_knowledge_apis_tenant_idx", "tenant_id"),
  943. sa.Index("external_knowledge_apis_name_idx", "name"),
  944. )
  945. id = mapped_column(StringUUID, nullable=False, server_default=sa.text("uuid_generate_v4()"))
  946. name: Mapped[str] = mapped_column(String(255), nullable=False)
  947. description: Mapped[str] = mapped_column(String(255), nullable=False)
  948. tenant_id = mapped_column(StringUUID, nullable=False)
  949. settings = mapped_column(sa.Text, nullable=True)
  950. created_by = mapped_column(StringUUID, nullable=False)
  951. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  952. updated_by = mapped_column(StringUUID, nullable=True)
  953. updated_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  954. def to_dict(self):
  955. return {
  956. "id": self.id,
  957. "tenant_id": self.tenant_id,
  958. "name": self.name,
  959. "description": self.description,
  960. "settings": self.settings_dict,
  961. "dataset_bindings": self.dataset_bindings,
  962. "created_by": self.created_by,
  963. "created_at": self.created_at.isoformat(),
  964. }
  965. @property
  966. def settings_dict(self):
  967. try:
  968. return json.loads(self.settings) if self.settings else None
  969. except JSONDecodeError:
  970. return None
  971. @property
  972. def dataset_bindings(self):
  973. external_knowledge_bindings = (
  974. db.session.query(ExternalKnowledgeBindings)
  975. .where(ExternalKnowledgeBindings.external_knowledge_api_id == self.id)
  976. .all()
  977. )
  978. dataset_ids = [binding.dataset_id for binding in external_knowledge_bindings]
  979. datasets = db.session.query(Dataset).where(Dataset.id.in_(dataset_ids)).all()
  980. dataset_bindings = []
  981. for dataset in datasets:
  982. dataset_bindings.append({"id": dataset.id, "name": dataset.name})
  983. return dataset_bindings
  984. class ExternalKnowledgeBindings(Base):
  985. __tablename__ = "external_knowledge_bindings"
  986. __table_args__ = (
  987. sa.PrimaryKeyConstraint("id", name="external_knowledge_bindings_pkey"),
  988. sa.Index("external_knowledge_bindings_tenant_idx", "tenant_id"),
  989. sa.Index("external_knowledge_bindings_dataset_idx", "dataset_id"),
  990. sa.Index("external_knowledge_bindings_external_knowledge_idx", "external_knowledge_id"),
  991. sa.Index("external_knowledge_bindings_external_knowledge_api_idx", "external_knowledge_api_id"),
  992. )
  993. id = mapped_column(StringUUID, nullable=False, server_default=sa.text("uuid_generate_v4()"))
  994. tenant_id = mapped_column(StringUUID, nullable=False)
  995. external_knowledge_api_id = mapped_column(StringUUID, nullable=False)
  996. dataset_id = mapped_column(StringUUID, nullable=False)
  997. external_knowledge_id = mapped_column(sa.Text, nullable=False)
  998. created_by = mapped_column(StringUUID, nullable=False)
  999. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  1000. updated_by = mapped_column(StringUUID, nullable=True)
  1001. updated_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  1002. class DatasetAutoDisableLog(Base):
  1003. __tablename__ = "dataset_auto_disable_logs"
  1004. __table_args__ = (
  1005. sa.PrimaryKeyConstraint("id", name="dataset_auto_disable_log_pkey"),
  1006. sa.Index("dataset_auto_disable_log_tenant_idx", "tenant_id"),
  1007. sa.Index("dataset_auto_disable_log_dataset_idx", "dataset_id"),
  1008. sa.Index("dataset_auto_disable_log_created_atx", "created_at"),
  1009. )
  1010. id = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
  1011. tenant_id = mapped_column(StringUUID, nullable=False)
  1012. dataset_id = mapped_column(StringUUID, nullable=False)
  1013. document_id = mapped_column(StringUUID, nullable=False)
  1014. notified: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
  1015. created_at: Mapped[datetime] = mapped_column(
  1016. DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
  1017. )
  1018. class RateLimitLog(Base):
  1019. __tablename__ = "rate_limit_logs"
  1020. __table_args__ = (
  1021. sa.PrimaryKeyConstraint("id", name="rate_limit_log_pkey"),
  1022. sa.Index("rate_limit_log_tenant_idx", "tenant_id"),
  1023. sa.Index("rate_limit_log_operation_idx", "operation"),
  1024. )
  1025. id = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
  1026. tenant_id = mapped_column(StringUUID, nullable=False)
  1027. subscription_plan: Mapped[str] = mapped_column(String(255), nullable=False)
  1028. operation: Mapped[str] = mapped_column(String(255), nullable=False)
  1029. created_at: Mapped[datetime] = mapped_column(
  1030. DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
  1031. )
  1032. class DatasetMetadata(Base):
  1033. __tablename__ = "dataset_metadatas"
  1034. __table_args__ = (
  1035. sa.PrimaryKeyConstraint("id", name="dataset_metadata_pkey"),
  1036. sa.Index("dataset_metadata_tenant_idx", "tenant_id"),
  1037. sa.Index("dataset_metadata_dataset_idx", "dataset_id"),
  1038. )
  1039. id = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
  1040. tenant_id = mapped_column(StringUUID, nullable=False)
  1041. dataset_id = mapped_column(StringUUID, nullable=False)
  1042. type: Mapped[str] = mapped_column(String(255), nullable=False)
  1043. name: Mapped[str] = mapped_column(String(255), nullable=False)
  1044. created_at: Mapped[datetime] = mapped_column(
  1045. DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
  1046. )
  1047. updated_at: Mapped[datetime] = mapped_column(
  1048. DateTime, nullable=False, server_default=sa.text("CURRENT_TIMESTAMP(0)")
  1049. )
  1050. created_by = mapped_column(StringUUID, nullable=False)
  1051. updated_by = mapped_column(StringUUID, nullable=True)
  1052. class DatasetMetadataBinding(Base):
  1053. __tablename__ = "dataset_metadata_bindings"
  1054. __table_args__ = (
  1055. sa.PrimaryKeyConstraint("id", name="dataset_metadata_binding_pkey"),
  1056. sa.Index("dataset_metadata_binding_tenant_idx", "tenant_id"),
  1057. sa.Index("dataset_metadata_binding_dataset_idx", "dataset_id"),
  1058. sa.Index("dataset_metadata_binding_metadata_idx", "metadata_id"),
  1059. sa.Index("dataset_metadata_binding_document_idx", "document_id"),
  1060. )
  1061. id = mapped_column(StringUUID, server_default=sa.text("uuid_generate_v4()"))
  1062. tenant_id = mapped_column(StringUUID, nullable=False)
  1063. dataset_id = mapped_column(StringUUID, nullable=False)
  1064. metadata_id = mapped_column(StringUUID, nullable=False)
  1065. document_id = mapped_column(StringUUID, nullable=False)
  1066. created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
  1067. created_by = mapped_column(StringUUID, nullable=False)
  1068. class PipelineBuiltInTemplate(Base): # type: ignore[name-defined]
  1069. __tablename__ = "pipeline_built_in_templates"
  1070. __table_args__ = (db.PrimaryKeyConstraint("id", name="pipeline_built_in_template_pkey"),)
  1071. id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))
  1072. name = db.Column(db.String(255), nullable=False)
  1073. description = db.Column(db.Text, nullable=False)
  1074. chunk_structure = db.Column(db.String(255), nullable=False)
  1075. icon = db.Column(db.JSON, nullable=False)
  1076. yaml_content = db.Column(db.Text, nullable=False)
  1077. copyright = db.Column(db.String(255), nullable=False)
  1078. privacy_policy = db.Column(db.String(255), nullable=False)
  1079. position = db.Column(db.Integer, nullable=False)
  1080. install_count = db.Column(db.Integer, nullable=False, default=0)
  1081. language = db.Column(db.String(255), nullable=False)
  1082. created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
  1083. updated_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
  1084. created_by = db.Column(StringUUID, nullable=False)
  1085. updated_by = db.Column(StringUUID, nullable=True)
  1086. @property
  1087. def created_user_name(self):
  1088. account = db.session.query(Account).filter(Account.id == self.created_by).first()
  1089. if account:
  1090. return account.name
  1091. return ""
  1092. class PipelineCustomizedTemplate(Base): # type: ignore[name-defined]
  1093. __tablename__ = "pipeline_customized_templates"
  1094. __table_args__ = (
  1095. db.PrimaryKeyConstraint("id", name="pipeline_customized_template_pkey"),
  1096. db.Index("pipeline_customized_template_tenant_idx", "tenant_id"),
  1097. )
  1098. id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))
  1099. tenant_id = db.Column(StringUUID, nullable=False)
  1100. name = db.Column(db.String(255), nullable=False)
  1101. description = db.Column(db.Text, nullable=False)
  1102. chunk_structure = db.Column(db.String(255), nullable=False)
  1103. icon = db.Column(db.JSON, nullable=False)
  1104. position = db.Column(db.Integer, nullable=False)
  1105. yaml_content = db.Column(db.Text, nullable=False)
  1106. install_count = db.Column(db.Integer, nullable=False, default=0)
  1107. language = db.Column(db.String(255), nullable=False)
  1108. created_by = db.Column(StringUUID, nullable=False)
  1109. updated_by = db.Column(StringUUID, nullable=True)
  1110. created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
  1111. updated_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
  1112. @property
  1113. def created_user_name(self):
  1114. account = db.session.query(Account).filter(Account.id == self.created_by).first()
  1115. if account:
  1116. return account.name
  1117. return ""
  1118. class Pipeline(Base): # type: ignore[name-defined]
  1119. __tablename__ = "pipelines"
  1120. __table_args__ = (db.PrimaryKeyConstraint("id", name="pipeline_pkey"),)
  1121. id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))
  1122. tenant_id: Mapped[str] = db.Column(StringUUID, nullable=False)
  1123. name = db.Column(db.String(255), nullable=False)
  1124. description = db.Column(db.Text, nullable=False, server_default=db.text("''::character varying"))
  1125. workflow_id = db.Column(StringUUID, nullable=True)
  1126. is_public = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))
  1127. is_published = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))
  1128. created_by = db.Column(StringUUID, nullable=True)
  1129. created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
  1130. updated_by = db.Column(StringUUID, nullable=True)
  1131. updated_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
  1132. @property
  1133. def dataset(self):
  1134. return db.session.query(Dataset).filter(Dataset.pipeline_id == self.id).first()
  1135. class DocumentPipelineExecutionLog(Base):
  1136. __tablename__ = "document_pipeline_execution_logs"
  1137. __table_args__ = (
  1138. db.PrimaryKeyConstraint("id", name="document_pipeline_execution_log_pkey"),
  1139. db.Index("document_pipeline_execution_logs_document_id_idx", "document_id"),
  1140. )
  1141. id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))
  1142. pipeline_id = db.Column(StringUUID, nullable=False)
  1143. document_id = db.Column(StringUUID, nullable=False)
  1144. datasource_type = db.Column(db.String(255), nullable=False)
  1145. datasource_info = db.Column(db.Text, nullable=False)
  1146. datasource_node_id = db.Column(db.String(255), nullable=False)
  1147. input_data = db.Column(db.JSON, nullable=False)
  1148. created_by = db.Column(StringUUID, nullable=True)
  1149. created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
  1150. class PipelineRecommendedPlugin(Base):
  1151. __tablename__ = "pipeline_recommended_plugins"
  1152. __table_args__ = (db.PrimaryKeyConstraint("id", name="pipeline_recommended_plugin_pkey"),)
  1153. id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))
  1154. plugin_id = db.Column(db.Text, nullable=False)
  1155. provider_name = db.Column(db.Text, nullable=False)
  1156. position = db.Column(db.Integer, nullable=False, default=0)
  1157. active = db.Column(db.Boolean, nullable=False, default=True)
  1158. created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
  1159. updated_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())