dataset_service.py 125 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821
  1. import copy
  2. import datetime
  3. import json
  4. import logging
  5. import random
  6. import time
  7. import uuid
  8. from collections import Counter
  9. from typing import Any, Optional
  10. from flask_login import current_user # type: ignore
  11. from sqlalchemy import func, literal, text
  12. from sqlalchemy.orm import Session
  13. from werkzeug.exceptions import NotFound
  14. from configs import dify_config
  15. from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
  16. from core.model_manager import ModelManager
  17. from core.model_runtime.entities.model_entities import ModelType
  18. from core.plugin.entities.plugin import ModelProviderID
  19. from core.rag.index_processor.constant.built_in_field import BuiltInField
  20. from core.rag.index_processor.constant.index_type import IndexType
  21. from core.rag.retrieval.retrieval_methods import RetrievalMethod
  22. from events.dataset_event import dataset_was_deleted
  23. from events.document_event import document_was_deleted
  24. from events.template_event import template_was_deleted
  25. from extensions.ext_database import db
  26. from extensions.ext_redis import redis_client
  27. from libs import helper
  28. from models.account import Account, TenantAccountRole
  29. from models.dataset import (
  30. AppDatasetJoin,
  31. ChildChunk,
  32. Dataset,
  33. DatasetAutoDisableLog,
  34. DatasetCollectionBinding,
  35. DatasetPermission,
  36. DatasetPermissionAll,
  37. DatasetPermissionEnum,
  38. DatasetProcessRule,
  39. DatasetQuery,
  40. Document,
  41. DocumentSegment,
  42. ExternalKnowledgeBindings,
  43. Template,
  44. )
  45. from models.model import UploadFile
  46. from models.source import DataSourceOauthBinding
  47. from services.entities.knowledge_entities.knowledge_entities import (
  48. ChildChunkUpdateArgs,
  49. KnowledgeConfig,
  50. RerankingModel,
  51. RetrievalModel,
  52. SegmentUpdateArgs,
  53. )
  54. from services.errors.account import InvalidActionError, NoPermissionError
  55. from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexingError
  56. from services.errors.dataset import DatasetNameDuplicateError
  57. from services.errors.document import DocumentIndexingError
  58. from services.errors.file import FileNotExistsError
  59. from services.external_knowledge_service import ExternalDatasetService
  60. from services.feature_service import FeatureModel, FeatureService
  61. from services.tag_service import TagService
  62. from services.vector_service import VectorService
  63. from tasks.batch_clean_document_task import batch_clean_document_task
  64. from tasks.batch_clean_template_task import batch_clean_template_task
  65. from tasks.clean_notion_document_task import clean_notion_document_task
  66. from tasks.deal_dataset_vector_index_task import deal_dataset_vector_index_task
  67. from tasks.delete_segment_from_index_task import delete_segment_from_index_task
  68. from tasks.disable_segment_from_index_task import disable_segment_from_index_task
  69. from tasks.disable_segments_from_index_task import disable_segments_from_index_task
  70. from tasks.document_indexing_task import document_indexing_task
  71. from tasks.document_indexing_update_task import document_indexing_update_task
  72. from tasks.duplicate_document_indexing_task import duplicate_document_indexing_task
  73. from tasks.enable_segments_to_index_task import enable_segments_to_index_task
  74. from tasks.recover_document_indexing_task import recover_document_indexing_task
  75. from tasks.retry_document_indexing_task import retry_document_indexing_task
  76. from tasks.sync_website_document_indexing_task import sync_website_document_indexing_task
  77. class DatasetService:
  78. @staticmethod
  79. def get_datasets(
  80. page, per_page, tenant_id=None, user=None, search=None, tag_ids=None, category_ids=None, include_all=False
  81. ):
  82. query = Dataset.query.filter(Dataset.tenant_id == tenant_id).order_by(Dataset.created_at.desc())
  83. if user:
  84. # get permitted dataset ids
  85. dataset_permission = DatasetPermission.query.filter_by(account_id=user.id, tenant_id=tenant_id).all()
  86. permitted_dataset_ids = {dp.dataset_id for dp in dataset_permission} if dataset_permission else None
  87. if user.current_role == TenantAccountRole.DATASET_OPERATOR:
  88. # only show datasets that the user has permission to access
  89. if permitted_dataset_ids:
  90. query = query.filter(Dataset.id.in_(permitted_dataset_ids))
  91. else:
  92. return [], 0
  93. else:
  94. if user.current_role != TenantAccountRole.OWNER or not include_all:
  95. # show all datasets that the user has permission to access
  96. if permitted_dataset_ids:
  97. query = query.filter(
  98. db.or_(
  99. Dataset.permission == DatasetPermissionEnum.ALL_TEAM,
  100. db.and_(
  101. Dataset.permission == DatasetPermissionEnum.ONLY_ME, Dataset.created_by == user.id
  102. ),
  103. db.and_(
  104. Dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM,
  105. Dataset.id.in_(permitted_dataset_ids),
  106. ),
  107. )
  108. )
  109. else:
  110. query = query.filter(
  111. db.or_(
  112. Dataset.permission == DatasetPermissionEnum.ALL_TEAM,
  113. db.and_(
  114. Dataset.permission == DatasetPermissionEnum.ONLY_ME, Dataset.created_by == user.id
  115. ),
  116. )
  117. )
  118. else:
  119. # if no user, only show datasets that are shared with all team members
  120. query = query.filter(Dataset.permission == DatasetPermissionEnum.ALL_TEAM)
  121. if search:
  122. query = query.filter(Dataset.name.ilike(f"%{search}%"))
  123. if tag_ids:
  124. target_ids = TagService.get_target_ids_by_tag_ids("knowledge", tenant_id, tag_ids)
  125. if target_ids:
  126. query = query.filter(Dataset.id.in_(target_ids))
  127. else:
  128. return [], 0
  129. if category_ids:
  130. target_ids_by_category_ids = TagService.get_target_ids_by_tag_ids(
  131. "knowledge_category", tenant_id, category_ids
  132. )
  133. if target_ids_by_category_ids:
  134. query = query.filter(Dataset.id.in_(target_ids_by_category_ids))
  135. else:
  136. return [], 0
  137. datasets = query.paginate(page=page, per_page=per_page, max_per_page=100, error_out=False)
  138. return datasets.items, datasets.total
  139. @staticmethod
  140. def get_process_rules(dataset_id):
  141. # get the latest process rule
  142. dataset_process_rule = (
  143. db.session.query(DatasetProcessRule)
  144. .filter(DatasetProcessRule.dataset_id == dataset_id)
  145. .order_by(DatasetProcessRule.created_at.desc())
  146. .limit(1)
  147. .one_or_none()
  148. )
  149. if dataset_process_rule:
  150. mode = dataset_process_rule.mode
  151. rules = dataset_process_rule.rules_dict
  152. else:
  153. mode = DocumentService.DEFAULT_RULES["mode"]
  154. rules = DocumentService.DEFAULT_RULES["rules"]
  155. return {"mode": mode, "rules": rules}
  156. @staticmethod
  157. def get_datasets_by_ids(ids, tenant_id):
  158. datasets = Dataset.query.filter(Dataset.id.in_(ids), Dataset.tenant_id == tenant_id).paginate(
  159. page=1, per_page=len(ids), max_per_page=len(ids), error_out=False
  160. )
  161. return datasets.items, datasets.total
  162. @staticmethod
  163. def get_datasets_count(tenant_id: str, user) -> int:
  164. query = Dataset.query
  165. if tenant_id:
  166. query = query.filter(Dataset.tenant_id == tenant_id)
  167. if user:
  168. # Get dataset IDs that the user has permission to access
  169. dataset_permission = DatasetPermission.query.filter_by(account_id=user.id, tenant_id=tenant_id).all()
  170. permitted_dataset_ids = {dp.dataset_id for dp in dataset_permission} if dataset_permission else None
  171. if user.current_role == TenantAccountRole.DATASET_OPERATOR:
  172. # Only show datasets that the user has permission to access
  173. if permitted_dataset_ids:
  174. query = query.filter(Dataset.id.in_(permitted_dataset_ids))
  175. else:
  176. return 0
  177. else:
  178. if user.current_role != TenantAccountRole.OWNER:
  179. # Show all datasets that the user has permission to access
  180. if permitted_dataset_ids:
  181. query = query.filter(
  182. db.or_(
  183. Dataset.id.in_(permitted_dataset_ids),
  184. Dataset.permission == DatasetPermissionEnum.ALL_TEAM,
  185. )
  186. )
  187. return query.count()
  188. @staticmethod
  189. def get_tags_count(
  190. tenant_id: str,
  191. ) -> int:
  192. return TagService.get_tags_count(tenant_id)
  193. @staticmethod
  194. def count_datasets(tenant_id=None):
  195. query = Dataset.query
  196. if tenant_id:
  197. query = query.filter(Dataset.tenant_id == tenant_id)
  198. return query.count()
  199. @staticmethod
  200. def get_dataset_update_stats(tenant_id):
  201. """Get dataset update statistics."""
  202. from datetime import datetime, timedelta
  203. now = datetime.utcnow()
  204. base_query = Dataset.query
  205. if tenant_id:
  206. base_query = base_query.filter(Dataset.tenant_id == tenant_id)
  207. # 统计不同时间段内更新的数据集数量
  208. stats = {
  209. "over_180_days": base_query.filter(Dataset.updated_at <= now - timedelta(days=180)).count(),
  210. "within_180_days": base_query.filter(
  211. Dataset.updated_at > now - timedelta(days=180), Dataset.updated_at <= now
  212. ).count(),
  213. "within_90_days": base_query.filter(
  214. Dataset.updated_at > now - timedelta(days=90), Dataset.updated_at <= now
  215. ).count(),
  216. "within_30_days": base_query.filter(
  217. Dataset.updated_at > now - timedelta(days=30), Dataset.updated_at <= now
  218. ).count(),
  219. "within_7_days": base_query.filter(
  220. Dataset.updated_at > now - timedelta(days=7), Dataset.updated_at <= now
  221. ).count(),
  222. "within_3_days": base_query.filter(Dataset.updated_at > now - timedelta(days=3)).count(),
  223. }
  224. return stats
  225. @staticmethod
  226. def get_dataset_type_stats(tenant_id):
  227. sql = """
  228. SELECT t.name, t.id, COUNT(*) as count
  229. FROM datasets d
  230. LEFT JOIN tag_bindings tb ON tb.target_id = d.id
  231. LEFT JOIN tags t ON t.id = tb.tag_id
  232. WHERE t.type = 'knowledge_category'
  233. """
  234. params = {}
  235. if tenant_id:
  236. sql += " AND d.tenant_id = :tenant_id"
  237. params["tenant_id"] = tenant_id
  238. sql += " GROUP BY t.name, t.id"
  239. # 注意这里要用 text(sql)
  240. result = db.session.execute(text(sql), params).fetchall()
  241. total = sum(row.count for row in result)
  242. if total == 0:
  243. return []
  244. # 计算百分比并格式化结果
  245. stats = []
  246. for row in result:
  247. percentage = (row.count / total) * 100
  248. stats.append({"type": row.name, "percentage": round(percentage, 2)})
  249. return stats
  250. @staticmethod
  251. def has_edit_permission(account_id: str, dataset_id: str) -> bool:
  252. result = (
  253. db.session.query(DatasetPermissionAll)
  254. .filter_by(
  255. account_id=account_id,
  256. dataset_id=dataset_id,
  257. has_edit_permission=True
  258. )
  259. .first()
  260. )
  261. return result is not None
  262. @staticmethod
  263. def get_datasets_edit_permission(dataset_id):
  264. results = (
  265. db.session.query(DatasetPermissionAll.account_id, Account.email)
  266. .join(DatasetPermissionAll, Account.id == DatasetPermissionAll.account_id)
  267. .filter(DatasetPermissionAll.dataset_id == dataset_id, DatasetPermissionAll.has_edit_permission == True)
  268. .all()
  269. )
  270. edit_permission_list = []
  271. for row in results:
  272. edit_permission_list.append({"id": row.account_id, "email": row.email})
  273. return edit_permission_list
  274. @staticmethod
  275. def get_datasets_read_permission(dataset_id):
  276. results = (
  277. db.session.query(DatasetPermissionAll.account_id, Account.email)
  278. .join(DatasetPermissionAll, Account.id == DatasetPermissionAll.account_id)
  279. .filter(DatasetPermissionAll.dataset_id == dataset_id, DatasetPermissionAll.has_read_permission == True)
  280. .all()
  281. )
  282. read_permission_list = []
  283. for row in results:
  284. read_permission_list.append({"id": row.account_id, "email": row.email})
  285. return read_permission_list
  286. @staticmethod
  287. def update_dataset_edit_auth(dataset_id, edit_auth):
  288. dataset = DatasetService.get_dataset(dataset_id)
  289. if not dataset:
  290. raise ValueError("Dataset not found")
  291. if dataset.edit_auth != edit_auth:
  292. filtered_data["updated_by"] = user.id
  293. filtered_data["updated_at"] = datetime.datetime.now()
  294. filtered_data["edit_auth"] = edit_auth
  295. dataset.query.filter_by(id=dataset_id).update(filtered_data)
  296. db.session.commit()
  297. # if action:
  298. # deal_dataset_vector_index_task.delay(dataset_id, action)
  299. @staticmethod
  300. def update_dataset_edit_and_read_permission(dataset_id, user, edit_permission_list, read_permission_list):
  301. for edit_permission in edit_permission_list:
  302. print(f" 编辑权限 ID: {edit_permission.get('id')}, Email: {edit_permission.get('email')}")
  303. account_id = edit_permission.get("id")
  304. email = edit_permission.get("email")
  305. dataset_permission_all: Optional[DatasetPermissionAll] = DatasetPermissionAll.query.filter_by(
  306. dataset_id=dataset_id, email=email, account_id=account_id
  307. ).first()
  308. if dataset_permission_all == None:
  309. print(f" 数据表原本无编辑权限 ID: {edit_permission.get('id')}, Email: {edit_permission.get('email')}")
  310. new_dataset_permission_all = DatasetPermissionAll()
  311. new_dataset_permission_all.dataset_id = dataset_id
  312. new_dataset_permission_all.account_id = account_id
  313. new_dataset_permission_all.created_by = user.id
  314. new_dataset_permission_all.created_at = datetime.datetime.now()
  315. new_dataset_permission_all.has_edit_permission = True
  316. new_dataset_permission_all.email = email
  317. db.session.add(new_dataset_permission_all)
  318. db.session.commit()
  319. db.session.flush()
  320. else:
  321. if dataset_permission_all.has_edit_permission == False:
  322. dataset_permission_all.query.filter_by(dataset_id=dataset_id, account_id=account_id).update(
  323. {"has_edit_permission": True, "updated_by": user.id, "updated_at": datetime.datetime.now()}
  324. )
  325. db.session.commit()
  326. for read_permission in read_permission_list:
  327. print(f" 查看权限 ID: {read_permission.get('id')}, Email: {read_permission.get('email')}")
  328. account_id = read_permission.get("id")
  329. email = read_permission.get("email")
  330. dataset_permission_all: Optional[DatasetPermissionAll] = DatasetPermissionAll.query.filter_by(
  331. dataset_id=dataset_id, account_id=account_id
  332. ).first()
  333. if dataset_permission_all == None:
  334. # 插入
  335. print(f" 数据表无有查看权限 ID: {read_permission.get('id')}, Email: {read_permission.get('email')}")
  336. new_dataset_permission_all = DatasetPermissionAll()
  337. new_dataset_permission_all.dataset_id = dataset_id
  338. new_dataset_permission_all.account_id = account_id
  339. new_dataset_permission_all.created_by = user.id
  340. new_dataset_permission_all.updated_by = user.id
  341. new_dataset_permission_all.has_read_permission = True
  342. new_dataset_permission_all.email = email
  343. db.session.add(new_dataset_permission_all)
  344. db.session.commit()
  345. db.session.flush()
  346. else:
  347. if dataset_permission_all.has_read_permission == False:
  348. dataset_permission_all.query.filter_by(dataset_id=dataset_id, account_id=account_id).update(
  349. {"has_read_permission": True, "updated_by": user.id, "updated_at": datetime.datetime.now()}
  350. )
  351. db.session.commit()
  352. @staticmethod
  353. def create_empty_dataset(
  354. tenant_id: str,
  355. name: str,
  356. description: Optional[str],
  357. indexing_technique: Optional[str],
  358. account: Account,
  359. permission: Optional[str] = None,
  360. provider: str = "vendor",
  361. external_knowledge_api_id: Optional[str] = None,
  362. external_knowledge_id: Optional[str] = None,
  363. ):
  364. # check if dataset name already exists
  365. if Dataset.query.filter_by(name=name, tenant_id=tenant_id).first():
  366. raise DatasetNameDuplicateError(f"Dataset with name {name} already exists.")
  367. embedding_model = None
  368. if indexing_technique == "high_quality":
  369. model_manager = ModelManager()
  370. embedding_model = model_manager.get_default_model_instance(
  371. tenant_id=tenant_id, model_type=ModelType.TEXT_EMBEDDING
  372. )
  373. dataset = Dataset(name=name, indexing_technique=indexing_technique)
  374. # dataset = Dataset(name=name, provider=provider, config=config)
  375. dataset.description = description
  376. dataset.created_by = account.id
  377. dataset.updated_by = account.id
  378. dataset.tenant_id = tenant_id
  379. dataset.embedding_model_provider = embedding_model.provider if embedding_model else None
  380. dataset.embedding_model = embedding_model.model if embedding_model else None
  381. dataset.permission = DatasetPermissionEnum.ALL_TEAM
  382. dataset.provider = provider
  383. dataset.dept_id = account.dept_id
  384. db.session.add(dataset)
  385. db.session.flush()
  386. if provider == "external" and external_knowledge_api_id:
  387. external_knowledge_api = ExternalDatasetService.get_external_knowledge_api(external_knowledge_api_id)
  388. if not external_knowledge_api:
  389. raise ValueError("External API template not found.")
  390. external_knowledge_binding = ExternalKnowledgeBindings(
  391. tenant_id=tenant_id,
  392. dataset_id=dataset.id,
  393. external_knowledge_api_id=external_knowledge_api_id,
  394. external_knowledge_id=external_knowledge_id,
  395. created_by=account.id,
  396. )
  397. db.session.add(external_knowledge_binding)
  398. db.session.commit()
  399. return dataset
  400. @staticmethod
  401. def get_dataset(dataset_id) -> Optional[Dataset]:
  402. dataset: Optional[Dataset] = Dataset.query.filter_by(id=dataset_id).first()
  403. return dataset
  404. @staticmethod
  405. def check_dataset_model_setting(dataset):
  406. if dataset.indexing_technique == "high_quality":
  407. try:
  408. model_manager = ModelManager()
  409. model_manager.get_model_instance(
  410. tenant_id=dataset.tenant_id,
  411. provider=dataset.embedding_model_provider,
  412. model_type=ModelType.TEXT_EMBEDDING,
  413. model=dataset.embedding_model,
  414. )
  415. except LLMBadRequestError:
  416. raise ValueError(
  417. "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider."
  418. )
  419. except ProviderTokenNotInitError as ex:
  420. raise ValueError(f"The dataset in unavailable, due to: {ex.description}")
  421. @staticmethod
  422. def check_embedding_model_setting(tenant_id: str, embedding_model_provider: str, embedding_model: str):
  423. try:
  424. model_manager = ModelManager()
  425. model_manager.get_model_instance(
  426. tenant_id=tenant_id,
  427. provider=embedding_model_provider,
  428. model_type=ModelType.TEXT_EMBEDDING,
  429. model=embedding_model,
  430. )
  431. except LLMBadRequestError:
  432. raise ValueError(
  433. "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider."
  434. )
  435. except ProviderTokenNotInitError as ex:
  436. raise ValueError(ex.description)
  437. @staticmethod
  438. def update_dataset(dataset_id, data, user):
  439. dataset = DatasetService.get_dataset(dataset_id)
  440. if not dataset:
  441. raise ValueError("Dataset not found")
  442. DatasetService.check_dataset_permission(dataset, user)
  443. if dataset.provider == "external":
  444. external_retrieval_model = data.get("external_retrieval_model", None)
  445. if external_retrieval_model:
  446. dataset.retrieval_model = external_retrieval_model
  447. dataset.name = data.get("name", dataset.name)
  448. dataset.description = data.get("description", "")
  449. permission = data.get("permission")
  450. if permission:
  451. dataset.permission = permission
  452. external_knowledge_id = data.get("external_knowledge_id", None)
  453. db.session.add(dataset)
  454. if not external_knowledge_id:
  455. raise ValueError("External knowledge id is required.")
  456. external_knowledge_api_id = data.get("external_knowledge_api_id", None)
  457. if not external_knowledge_api_id:
  458. raise ValueError("External knowledge api id is required.")
  459. with Session(db.engine) as session:
  460. external_knowledge_binding = (
  461. session.query(ExternalKnowledgeBindings).filter_by(dataset_id=dataset_id).first()
  462. )
  463. if not external_knowledge_binding:
  464. raise ValueError("External knowledge binding not found.")
  465. if (
  466. external_knowledge_binding.external_knowledge_id != external_knowledge_id
  467. or external_knowledge_binding.external_knowledge_api_id != external_knowledge_api_id
  468. ):
  469. external_knowledge_binding.external_knowledge_id = external_knowledge_id
  470. external_knowledge_binding.external_knowledge_api_id = external_knowledge_api_id
  471. db.session.add(external_knowledge_binding)
  472. db.session.commit()
  473. else:
  474. data.pop("partial_member_list", None)
  475. data.pop("external_knowledge_api_id", None)
  476. data.pop("external_knowledge_id", None)
  477. data.pop("external_retrieval_model", None)
  478. filtered_data = {k: v for k, v in data.items() if v is not None or k == "description"}
  479. action = None
  480. if dataset.indexing_technique != data["indexing_technique"]:
  481. # if update indexing_technique
  482. if data["indexing_technique"] == "economy":
  483. action = "remove"
  484. filtered_data["embedding_model"] = None
  485. filtered_data["embedding_model_provider"] = None
  486. filtered_data["collection_binding_id"] = None
  487. elif data["indexing_technique"] == "high_quality":
  488. action = "add"
  489. # get embedding model setting
  490. try:
  491. model_manager = ModelManager()
  492. embedding_model = model_manager.get_model_instance(
  493. tenant_id=current_user.current_tenant_id,
  494. provider=data["embedding_model_provider"],
  495. model_type=ModelType.TEXT_EMBEDDING,
  496. model=data["embedding_model"],
  497. )
  498. filtered_data["embedding_model"] = embedding_model.model
  499. filtered_data["embedding_model_provider"] = embedding_model.provider
  500. dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
  501. embedding_model.provider, embedding_model.model
  502. )
  503. filtered_data["collection_binding_id"] = dataset_collection_binding.id
  504. except LLMBadRequestError:
  505. raise ValueError(
  506. "No Embedding Model available. Please configure a valid provider "
  507. "in the Settings -> Model Provider."
  508. )
  509. except ProviderTokenNotInitError as ex:
  510. raise ValueError(ex.description)
  511. else:
  512. # add default plugin id to both setting sets, to make sure the plugin model provider is consistent
  513. # Skip embedding model checks if not provided in the update request
  514. if (
  515. "embedding_model_provider" not in data
  516. or "embedding_model" not in data
  517. or not data.get("embedding_model_provider")
  518. or not data.get("embedding_model")
  519. ):
  520. # If the dataset already has embedding model settings, use those
  521. if dataset.embedding_model_provider and dataset.embedding_model:
  522. # Keep existing values
  523. filtered_data["embedding_model_provider"] = dataset.embedding_model_provider
  524. filtered_data["embedding_model"] = dataset.embedding_model
  525. # If collection_binding_id exists, keep it too
  526. if dataset.collection_binding_id:
  527. filtered_data["collection_binding_id"] = dataset.collection_binding_id
  528. # Otherwise, don't try to update embedding model settings at all
  529. # Remove these fields from filtered_data if they exist but are None/empty
  530. if "embedding_model_provider" in filtered_data and not filtered_data["embedding_model_provider"]:
  531. del filtered_data["embedding_model_provider"]
  532. if "embedding_model" in filtered_data and not filtered_data["embedding_model"]:
  533. del filtered_data["embedding_model"]
  534. else:
  535. skip_embedding_update = False
  536. try:
  537. # Handle existing model provider
  538. plugin_model_provider = dataset.embedding_model_provider
  539. plugin_model_provider_str = None
  540. if plugin_model_provider:
  541. plugin_model_provider_str = str(ModelProviderID(plugin_model_provider))
  542. # Handle new model provider from request
  543. new_plugin_model_provider = data["embedding_model_provider"]
  544. new_plugin_model_provider_str = None
  545. if new_plugin_model_provider:
  546. new_plugin_model_provider_str = str(ModelProviderID(new_plugin_model_provider))
  547. # Only update embedding model if both values are provided and different from current
  548. if (
  549. plugin_model_provider_str != new_plugin_model_provider_str
  550. or data["embedding_model"] != dataset.embedding_model
  551. ):
  552. action = "update"
  553. model_manager = ModelManager()
  554. try:
  555. embedding_model = model_manager.get_model_instance(
  556. tenant_id=current_user.current_tenant_id,
  557. provider=data["embedding_model_provider"],
  558. model_type=ModelType.TEXT_EMBEDDING,
  559. model=data["embedding_model"],
  560. )
  561. except ProviderTokenNotInitError:
  562. # If we can't get the embedding model, skip updating it
  563. # and keep the existing settings if available
  564. if dataset.embedding_model_provider and dataset.embedding_model:
  565. filtered_data["embedding_model_provider"] = dataset.embedding_model_provider
  566. filtered_data["embedding_model"] = dataset.embedding_model
  567. if dataset.collection_binding_id:
  568. filtered_data["collection_binding_id"] = dataset.collection_binding_id
  569. # Skip the rest of the embedding model update
  570. skip_embedding_update = True
  571. if not skip_embedding_update:
  572. filtered_data["embedding_model"] = embedding_model.model
  573. filtered_data["embedding_model_provider"] = embedding_model.provider
  574. dataset_collection_binding = (
  575. DatasetCollectionBindingService.get_dataset_collection_binding(
  576. embedding_model.provider, embedding_model.model
  577. )
  578. )
  579. filtered_data["collection_binding_id"] = dataset_collection_binding.id
  580. except LLMBadRequestError:
  581. raise ValueError(
  582. "No Embedding Model available. Please configure a valid provider "
  583. "in the Settings -> Model Provider."
  584. )
  585. except ProviderTokenNotInitError as ex:
  586. raise ValueError(ex.description)
  587. filtered_data["updated_by"] = user.id
  588. filtered_data["updated_at"] = datetime.datetime.now()
  589. # update Retrieval model
  590. filtered_data["retrieval_model"] = data["retrieval_model"]
  591. dataset.query.filter_by(id=dataset_id).update(filtered_data)
  592. db.session.commit()
  593. if action:
  594. deal_dataset_vector_index_task.delay(dataset_id, action)
  595. return dataset
  596. @staticmethod
  597. def delete_dataset(dataset_id, user):
  598. dataset = DatasetService.get_dataset(dataset_id)
  599. if dataset is None:
  600. return False
  601. DatasetService.check_dataset_permission(dataset, user)
  602. dataset_was_deleted.send(dataset)
  603. db.session.delete(dataset)
  604. db.session.commit()
  605. return True
  606. @staticmethod
  607. def dataset_use_check(dataset_id) -> bool:
  608. count = AppDatasetJoin.query.filter_by(dataset_id=dataset_id).count()
  609. if count > 0:
  610. return True
  611. return False
  612. @staticmethod
  613. def check_dataset_permission(dataset, user):
  614. if dataset.tenant_id != user.current_tenant_id:
  615. logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")
  616. raise NoPermissionError("You do not have permission to access this dataset.")
  617. if user.current_role != TenantAccountRole.OWNER:
  618. if dataset.permission == DatasetPermissionEnum.ONLY_ME and dataset.created_by != user.id:
  619. logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")
  620. raise NoPermissionError("You do not have permission to access this dataset.")
  621. if dataset.permission == "partial_members":
  622. user_permission = DatasetPermission.query.filter_by(dataset_id=dataset.id, account_id=user.id).first()
  623. if (
  624. not user_permission
  625. and dataset.tenant_id != user.current_tenant_id
  626. and dataset.created_by != user.id
  627. ):
  628. logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")
  629. raise NoPermissionError("You do not have permission to access this dataset.")
  630. @staticmethod
  631. def check_dataset_permission_new(dataset, user):
  632. if dataset.tenant_id != user.current_tenant_id:
  633. logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")
  634. raise NoPermissionError("You do not have permission to access this dataset.")
  635. if user.current_role != TenantAccountRole.OWNER:
  636. if user.current_role != Acc.ADMIN and dataset.created_by != user.id:
  637. logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")
  638. raise NoPermissionError("You do not have permission to access this dataset.")
  639. @staticmethod
  640. def check_dataset_operator_permission(user: Optional[Account] = None, dataset: Optional[Dataset] = None):
  641. if not dataset:
  642. raise ValueError("Dataset not found")
  643. if not user:
  644. raise ValueError("User not found")
  645. if user.current_role != TenantAccountRole.OWNER:
  646. if dataset.permission == DatasetPermissionEnum.ONLY_ME:
  647. if dataset.created_by != user.id:
  648. raise NoPermissionError("You do not have permission to access this dataset.")
  649. elif dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM:
  650. if not any(
  651. dp.dataset_id == dataset.id for dp in DatasetPermission.query.filter_by(account_id=user.id).all()
  652. ):
  653. raise NoPermissionError("You do not have permission to access this dataset.")
  654. @staticmethod
  655. def get_dataset_queries(dataset_id: str, page: int, per_page: int):
  656. dataset_queries = (
  657. DatasetQuery.query.filter_by(dataset_id=dataset_id)
  658. .order_by(db.desc(DatasetQuery.created_at))
  659. .paginate(page=page, per_page=per_page, max_per_page=100, error_out=False)
  660. )
  661. return dataset_queries.items, dataset_queries.total
  662. @staticmethod
  663. def get_related_apps(dataset_id: str):
  664. return (
  665. AppDatasetJoin.query.filter(AppDatasetJoin.dataset_id == dataset_id)
  666. .order_by(db.desc(AppDatasetJoin.created_at))
  667. .all()
  668. )
  669. @staticmethod
  670. def get_dataset_auto_disable_logs(dataset_id: str) -> dict:
  671. features = FeatureService.get_features(current_user.current_tenant_id)
  672. if not features.billing.enabled or features.billing.subscription.plan == "sandbox":
  673. return {
  674. "document_ids": [],
  675. "count": 0,
  676. }
  677. # get recent 30 days auto disable logs
  678. start_date = datetime.datetime.now() - datetime.timedelta(days=30)
  679. dataset_auto_disable_logs = DatasetAutoDisableLog.query.filter(
  680. DatasetAutoDisableLog.dataset_id == dataset_id,
  681. DatasetAutoDisableLog.created_at >= start_date,
  682. ).all()
  683. if dataset_auto_disable_logs:
  684. return {
  685. "document_ids": [log.document_id for log in dataset_auto_disable_logs],
  686. "count": len(dataset_auto_disable_logs),
  687. }
  688. return {
  689. "document_ids": [],
  690. "count": 0,
  691. }
  692. @staticmethod
  693. def get_datasets2(
  694. page, per_page, tenant_id=None, user=None,
  695. search=None, tag_ids=None, category_ids=None,
  696. auth_type=None, creator_dept=None, creator=None, include_all=False
  697. ):
  698. user_id = user.id
  699. queries = []
  700. # 1. 创建:row.created_by == loginUserId,
  701. query1 = Dataset.query.filter(Dataset.created_by == user_id)
  702. # 2. 编辑:row.deptId == loginDeptId && row.editAuth == 2
  703. query2 = Dataset.query.join(Account, Dataset.dept_id == Account.dept_id)
  704. query2 = query2.filter(Dataset.edit_auth == 2, Account.id == user_id)
  705. # 3.授权编辑:row.editUserIds.includes(loginUserId)
  706. query3 = Dataset.query.join(
  707. DatasetPermissionAll,
  708. Dataset.id == DatasetPermissionAll.dataset_id
  709. ).filter(
  710. DatasetPermissionAll.account_id == user_id,
  711. DatasetPermissionAll.has_edit_permission == True
  712. )
  713. # 4.授权可见:row.lookUserIds.includes(loginUserId)
  714. query4 = Dataset.query.join(
  715. DatasetPermissionAll,
  716. Dataset.id == DatasetPermissionAll.dataset_id
  717. ).filter(
  718. DatasetPermissionAll.account_id == user_id,
  719. DatasetPermissionAll.has_read_permission == True
  720. )
  721. # 根据 auth_type 选择要使用的查询
  722. if auth_type is None:
  723. # 如果 auth_type 为空,使用所有查询
  724. queries = [query1, query2, query3, query4]
  725. elif auth_type == 1:
  726. queries = [query1]
  727. elif auth_type == 2:
  728. queries = [query2]
  729. elif auth_type == 3:
  730. queries = [query3]
  731. elif auth_type == 4:
  732. queries = [query4]
  733. # 合并查询
  734. if not queries:
  735. return [], 0
  736. union_query = queries[0]
  737. for query in queries[1:]:
  738. union_query = union_query.union(query)
  739. # 添加创建人部门过滤
  740. if creator_dept:
  741. union_query = union_query.join(Account, Dataset.created_by == Account.id)
  742. union_query = union_query.filter(Account.dept_id == literal(str(creator_dept)))
  743. # 添加创建人过滤
  744. if creator:
  745. union_query = union_query.filter(Dataset.created_by == creator)
  746. # 其它过滤
  747. if search:
  748. union_query = union_query.filter(Dataset.name.ilike(f"%{search}%"))
  749. if tag_ids:
  750. target_ids = TagService.get_target_ids_by_tag_ids("knowledge", tenant_id, tag_ids)
  751. if target_ids:
  752. union_query = union_query.filter(Dataset.id.in_(target_ids))
  753. else:
  754. return [], 0
  755. if category_ids:
  756. target_ids_by_category_ids = TagService.get_target_ids_by_tag_ids("knowledge_category",
  757. tenant_id, category_ids)
  758. if target_ids_by_category_ids:
  759. union_query = union_query.filter(Dataset.id.in_(target_ids_by_category_ids))
  760. else:
  761. return [], 0
  762. datasets = union_query.order_by(Dataset.created_at.desc()).paginate(page=page, per_page=per_page,
  763. max_per_page=100, error_out=False)
  764. return datasets.items, datasets.total
  765. class TemplateService:
  766. DEFAULT_RULES: dict[str, Any] = {
  767. "mode": "custom",
  768. "rules": {
  769. "pre_processing_rules": [
  770. {"id": "remove_extra_spaces", "enabled": True},
  771. {"id": "remove_urls_emails", "enabled": False},
  772. ],
  773. "segmentation": {"delimiter": "\n", "max_tokens": 500, "chunk_overlap": 50},
  774. },
  775. "limits": {
  776. "indexing_max_segmentation_tokens_length": dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH,
  777. },
  778. }
  779. # 批量删除
  780. @staticmethod
  781. def delete_templates(dataset: Dataset, template_ids: list[str]):
  782. templates = db.session.query(Template).filter(Template.id.in_(template_ids)).all()
  783. file_ids = [
  784. template.data_source_info_dict["upload_file_id"]
  785. for template in templates
  786. if template.data_source_type == "upload_file"
  787. ]
  788. batch_clean_template_task.delay(template_ids, dataset.id, dataset.doc_form, file_ids)
  789. for template in templates:
  790. db.session.delete(template)
  791. db.session.commit()
  792. @staticmethod
  793. def delete_template(template):
  794. file_id = None
  795. if template.data_source_type == "upload_file":
  796. if template.data_source_info:
  797. data_source_info = template.data_source_info_dict
  798. if data_source_info and "upload_file_id" in data_source_info:
  799. file_id = data_source_info["upload_file_id"]
  800. template_was_deleted.send(template.id, file_id=file_id)
  801. db.session.delete(template)
  802. db.session.commit()
  803. @staticmethod
  804. def save_template_with_dataset_id(
  805. upload_file: UploadFile,
  806. dataset: Dataset,
  807. account: Account | Any,
  808. dataset_process_rule: Optional[DatasetProcessRule] = None,
  809. created_from: str = "web",
  810. ):
  811. batch = time.strftime("%Y%m%d%H%M%S") + str(random.randint(100000, 999999))
  812. # save process rule
  813. lock_name = "add_template_lock_dataset_id_{}".format(dataset.id)
  814. with redis_client.lock(lock_name, timeout=600):
  815. position = TemplateService.get_templates_position(dataset.id)
  816. if not upload_file:
  817. raise FileNotExistsError()
  818. file_name = upload_file.name
  819. file_id = upload_file.id
  820. file_url = upload_file.key
  821. data_source_info = {"upload_file_id": upload_file.id}
  822. template = TemplateService.build_template(
  823. dataset,
  824. # dataset_process_rule.id, # type: ignore
  825. # data_source_type,
  826. data_source_info,
  827. created_from,
  828. position,
  829. account,
  830. file_id,
  831. file_name,
  832. batch,
  833. file_url,
  834. )
  835. db.session.add(template)
  836. db.session.flush()
  837. db.session.commit()
  838. templates = []
  839. templates.append(template)
  840. position += 1
  841. return templates, batch
  842. @staticmethod
  843. def get_templates_position(dataset_id):
  844. template = Template.query.filter_by(dataset_id=dataset_id).order_by(Template.position.desc()).first()
  845. if template:
  846. return template.position + 1
  847. else:
  848. return 1
  849. @staticmethod
  850. def build_template(
  851. dataset: Dataset,
  852. # process_rule_id: str,
  853. # data_source_type: 'upload_file',
  854. data_source_info: dict,
  855. created_from: str,
  856. position: int,
  857. account: Account,
  858. file_id: str,
  859. file_name: str,
  860. batch: str,
  861. file_url: str,
  862. ):
  863. template = Template(
  864. tenant_id=dataset.tenant_id,
  865. dataset_id=dataset.id,
  866. position=position,
  867. data_source_type="upload_file",
  868. data_source_info=json.dumps(data_source_info),
  869. # dataset_process_rule_id=process_rule_id,
  870. batch=batch,
  871. name=file_name,
  872. created_from=created_from,
  873. created_by=account.id,
  874. file_id=file_id,
  875. file_url="./storage/" + file_url,
  876. )
  877. return template
  878. @staticmethod
  879. def get_templates(template_id) -> Optional[Template]:
  880. if template_id:
  881. print("模版id" + template_id)
  882. template: Optional[Template] = Template.query.filter_by(id=template_id).first()
  883. return template
  884. else:
  885. return None
  886. class DocumentService:
  887. DEFAULT_RULES: dict[str, Any] = {
  888. "mode": "custom",
  889. "rules": {
  890. "pre_processing_rules": [
  891. {"id": "remove_extra_spaces", "enabled": True},
  892. {"id": "remove_urls_emails", "enabled": False},
  893. ],
  894. "segmentation": {"delimiter": "\n", "max_tokens": 500, "chunk_overlap": 50},
  895. },
  896. "limits": {
  897. "indexing_max_segmentation_tokens_length": dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH,
  898. },
  899. }
  900. DOCUMENT_METADATA_SCHEMA: dict[str, Any] = {
  901. "book": {
  902. "title": str,
  903. "language": str,
  904. "author": str,
  905. "publisher": str,
  906. "publication_date": str,
  907. "isbn": str,
  908. "category": str,
  909. },
  910. "web_page": {
  911. "title": str,
  912. "url": str,
  913. "language": str,
  914. "publish_date": str,
  915. "author/publisher": str,
  916. "topic/keywords": str,
  917. "description": str,
  918. },
  919. "paper": {
  920. "title": str,
  921. "language": str,
  922. "author": str,
  923. "publish_date": str,
  924. "journal/conference_name": str,
  925. "volume/issue/page_numbers": str,
  926. "doi": str,
  927. "topic/keywords": str,
  928. "abstract": str,
  929. },
  930. "social_media_post": {
  931. "platform": str,
  932. "author/username": str,
  933. "publish_date": str,
  934. "post_url": str,
  935. "topic/tags": str,
  936. },
  937. "wikipedia_entry": {
  938. "title": str,
  939. "language": str,
  940. "web_page_url": str,
  941. "last_edit_date": str,
  942. "editor/contributor": str,
  943. "summary/introduction": str,
  944. },
  945. "personal_document": {
  946. "title": str,
  947. "author": str,
  948. "creation_date": str,
  949. "last_modified_date": str,
  950. "document_type": str,
  951. "tags/category": str,
  952. },
  953. "business_document": {
  954. "title": str,
  955. "author": str,
  956. "creation_date": str,
  957. "last_modified_date": str,
  958. "document_type": str,
  959. "department/team": str,
  960. },
  961. "im_chat_log": {
  962. "chat_platform": str,
  963. "chat_participants/group_name": str,
  964. "start_date": str,
  965. "end_date": str,
  966. "summary": str,
  967. },
  968. "synced_from_notion": {
  969. "title": str,
  970. "language": str,
  971. "author/creator": str,
  972. "creation_date": str,
  973. "last_modified_date": str,
  974. "notion_page_link": str,
  975. "category/tags": str,
  976. "description": str,
  977. },
  978. "synced_from_github": {
  979. "repository_name": str,
  980. "repository_description": str,
  981. "repository_owner/organization": str,
  982. "code_filename": str,
  983. "code_file_path": str,
  984. "programming_language": str,
  985. "github_link": str,
  986. "open_source_license": str,
  987. "commit_date": str,
  988. "commit_author": str,
  989. },
  990. "others": dict,
  991. }
  992. @staticmethod
  993. def get_document(dataset_id: str, document_id: Optional[str] = None) -> Optional[Document]:
  994. if document_id:
  995. document = (
  996. db.session.query(Document).filter(Document.id == document_id, Document.dataset_id == dataset_id).first()
  997. )
  998. return document
  999. else:
  1000. return None
  1001. @staticmethod
  1002. def get_document_by_id(document_id: str) -> Optional[Document]:
  1003. document = db.session.query(Document).filter(Document.id == document_id).first()
  1004. return document
  1005. @staticmethod
  1006. def get_document_by_ids(document_ids: list[str]) -> list[Document]:
  1007. documents = (
  1008. db.session.query(Document)
  1009. .filter(
  1010. Document.id.in_(document_ids),
  1011. Document.enabled == True,
  1012. Document.indexing_status == "completed",
  1013. Document.archived == False,
  1014. )
  1015. .all()
  1016. )
  1017. return documents
  1018. @staticmethod
  1019. def get_document_by_dataset_id(dataset_id: str) -> list[Document]:
  1020. documents = (
  1021. db.session.query(Document)
  1022. .filter(
  1023. Document.dataset_id == dataset_id,
  1024. Document.enabled == True,
  1025. )
  1026. .all()
  1027. )
  1028. return documents
  1029. @staticmethod
  1030. def get_working_documents_by_dataset_id(dataset_id: str) -> list[Document]:
  1031. documents = (
  1032. db.session.query(Document)
  1033. .filter(
  1034. Document.dataset_id == dataset_id,
  1035. Document.enabled == True,
  1036. Document.indexing_status == "completed",
  1037. Document.archived == False,
  1038. )
  1039. .all()
  1040. )
  1041. return documents
  1042. @staticmethod
  1043. def get_error_documents_by_dataset_id(dataset_id: str) -> list[Document]:
  1044. documents = (
  1045. db.session.query(Document)
  1046. .filter(Document.dataset_id == dataset_id, Document.indexing_status.in_(["error", "paused"]))
  1047. .all()
  1048. )
  1049. return documents
  1050. @staticmethod
  1051. def get_batch_documents(dataset_id: str, batch: str) -> list[Document]:
  1052. documents = (
  1053. db.session.query(Document)
  1054. .filter(
  1055. Document.batch == batch,
  1056. Document.dataset_id == dataset_id,
  1057. Document.tenant_id == current_user.current_tenant_id,
  1058. )
  1059. .all()
  1060. )
  1061. return documents
  1062. @staticmethod
  1063. def get_document_file_detail(file_id: str):
  1064. file_detail = db.session.query(UploadFile).filter(UploadFile.id == file_id).one_or_none()
  1065. return file_detail
  1066. @staticmethod
  1067. def check_archived(document):
  1068. if document.archived:
  1069. return True
  1070. else:
  1071. return False
  1072. @staticmethod
  1073. def delete_document(document):
  1074. # trigger document_was_deleted signal
  1075. file_id = None
  1076. if document.data_source_type == "upload_file":
  1077. if document.data_source_info:
  1078. data_source_info = document.data_source_info_dict
  1079. if data_source_info and "upload_file_id" in data_source_info:
  1080. file_id = data_source_info["upload_file_id"]
  1081. document_was_deleted.send(
  1082. document.id, dataset_id=document.dataset_id, doc_form=document.doc_form, file_id=file_id
  1083. )
  1084. db.session.delete(document)
  1085. db.session.commit()
  1086. @staticmethod
  1087. def delete_documents(dataset: Dataset, document_ids: list[str]):
  1088. documents = db.session.query(Document).filter(Document.id.in_(document_ids)).all()
  1089. file_ids = [
  1090. document.data_source_info_dict["upload_file_id"]
  1091. for document in documents
  1092. if document.data_source_type == "upload_file"
  1093. ]
  1094. batch_clean_document_task.delay(document_ids, dataset.id, dataset.doc_form, file_ids)
  1095. for document in documents:
  1096. db.session.delete(document)
  1097. db.session.commit()
  1098. @staticmethod
  1099. def rename_document(dataset_id: str, document_id: str, name: str) -> Document:
  1100. dataset = DatasetService.get_dataset(dataset_id)
  1101. if not dataset:
  1102. raise ValueError("Dataset not found.")
  1103. document = DocumentService.get_document(dataset_id, document_id)
  1104. if not document:
  1105. raise ValueError("Document not found.")
  1106. if document.tenant_id != current_user.current_tenant_id:
  1107. raise ValueError("No permission.")
  1108. if dataset.built_in_field_enabled:
  1109. if document.doc_metadata:
  1110. doc_metadata = copy.deepcopy(document.doc_metadata)
  1111. doc_metadata[BuiltInField.document_name.value] = name
  1112. document.doc_metadata = doc_metadata
  1113. document.name = name
  1114. db.session.add(document)
  1115. db.session.commit()
  1116. return document
  1117. @staticmethod
  1118. def pause_document(document):
  1119. if document.indexing_status not in {"waiting", "parsing", "cleaning", "splitting", "indexing"}:
  1120. raise DocumentIndexingError()
  1121. # update document to be paused
  1122. document.is_paused = True
  1123. document.paused_by = current_user.id
  1124. document.paused_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1125. db.session.add(document)
  1126. db.session.commit()
  1127. # set document paused flag
  1128. indexing_cache_key = "document_{}_is_paused".format(document.id)
  1129. redis_client.setnx(indexing_cache_key, "True")
  1130. @staticmethod
  1131. def recover_document(document):
  1132. if not document.is_paused:
  1133. raise DocumentIndexingError()
  1134. # update document to be recover
  1135. document.is_paused = False
  1136. document.paused_by = None
  1137. document.paused_at = None
  1138. db.session.add(document)
  1139. db.session.commit()
  1140. # delete paused flag
  1141. indexing_cache_key = "document_{}_is_paused".format(document.id)
  1142. redis_client.delete(indexing_cache_key)
  1143. # trigger async task
  1144. recover_document_indexing_task.delay(document.dataset_id, document.id)
  1145. @staticmethod
  1146. def retry_document(dataset_id: str, documents: list[Document]):
  1147. for document in documents:
  1148. # add retry flag
  1149. retry_indexing_cache_key = "document_{}_is_retried".format(document.id)
  1150. cache_result = redis_client.get(retry_indexing_cache_key)
  1151. if cache_result is not None:
  1152. raise ValueError("Document is being retried, please try again later")
  1153. # retry document indexing
  1154. document.indexing_status = "waiting"
  1155. db.session.add(document)
  1156. db.session.commit()
  1157. redis_client.setex(retry_indexing_cache_key, 600, 1)
  1158. # trigger async task
  1159. document_ids = [document.id for document in documents]
  1160. retry_document_indexing_task.delay(dataset_id, document_ids)
  1161. @staticmethod
  1162. def sync_website_document(dataset_id: str, document: Document):
  1163. # add sync flag
  1164. sync_indexing_cache_key = "document_{}_is_sync".format(document.id)
  1165. cache_result = redis_client.get(sync_indexing_cache_key)
  1166. if cache_result is not None:
  1167. raise ValueError("Document is being synced, please try again later")
  1168. # sync document indexing
  1169. document.indexing_status = "waiting"
  1170. data_source_info = document.data_source_info_dict
  1171. data_source_info["mode"] = "scrape"
  1172. document.data_source_info = json.dumps(data_source_info, ensure_ascii=False)
  1173. db.session.add(document)
  1174. db.session.commit()
  1175. redis_client.setex(sync_indexing_cache_key, 600, 1)
  1176. sync_website_document_indexing_task.delay(dataset_id, document.id)
  1177. @staticmethod
  1178. def get_documents_position(dataset_id):
  1179. document = Document.query.filter_by(dataset_id=dataset_id).order_by(Document.position.desc()).first()
  1180. if document:
  1181. return document.position + 1
  1182. else:
  1183. return 1
  1184. @staticmethod
  1185. def save_document_with_dataset_id(
  1186. dataset: Dataset,
  1187. knowledge_config: KnowledgeConfig,
  1188. account: Account | Any,
  1189. dataset_process_rule: Optional[DatasetProcessRule] = None,
  1190. created_from: str = "web",
  1191. ):
  1192. # check document limit
  1193. features = FeatureService.get_features(current_user.current_tenant_id)
  1194. if features.billing.enabled:
  1195. if not knowledge_config.original_document_id:
  1196. count = 0
  1197. if knowledge_config.data_source:
  1198. if knowledge_config.data_source.info_list.data_source_type == "upload_file":
  1199. upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids # type: ignore
  1200. count = len(upload_file_list)
  1201. elif knowledge_config.data_source.info_list.data_source_type == "notion_import":
  1202. notion_info_list = knowledge_config.data_source.info_list.notion_info_list
  1203. for notion_info in notion_info_list: # type: ignore
  1204. count = count + len(notion_info.pages)
  1205. elif knowledge_config.data_source.info_list.data_source_type == "website_crawl":
  1206. website_info = knowledge_config.data_source.info_list.website_info_list
  1207. count = len(website_info.urls) # type: ignore
  1208. batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT)
  1209. if count > batch_upload_limit:
  1210. raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")
  1211. DocumentService.check_documents_upload_quota(count, features)
  1212. # if dataset is empty, update dataset data_source_type
  1213. if not dataset.data_source_type:
  1214. dataset.data_source_type = knowledge_config.data_source.info_list.data_source_type # type: ignore
  1215. if not dataset.indexing_technique:
  1216. if knowledge_config.indexing_technique not in Dataset.INDEXING_TECHNIQUE_LIST:
  1217. raise ValueError("Indexing technique is invalid")
  1218. dataset.indexing_technique = knowledge_config.indexing_technique
  1219. if knowledge_config.indexing_technique == "high_quality":
  1220. model_manager = ModelManager()
  1221. if knowledge_config.embedding_model and knowledge_config.embedding_model_provider:
  1222. dataset_embedding_model = knowledge_config.embedding_model
  1223. dataset_embedding_model_provider = knowledge_config.embedding_model_provider
  1224. else:
  1225. embedding_model = model_manager.get_default_model_instance(
  1226. tenant_id=current_user.current_tenant_id, model_type=ModelType.TEXT_EMBEDDING
  1227. )
  1228. dataset_embedding_model = embedding_model.model
  1229. dataset_embedding_model_provider = embedding_model.provider
  1230. dataset.embedding_model = dataset_embedding_model
  1231. dataset.embedding_model_provider = dataset_embedding_model_provider
  1232. dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
  1233. dataset_embedding_model_provider, dataset_embedding_model
  1234. )
  1235. dataset.collection_binding_id = dataset_collection_binding.id
  1236. if not dataset.retrieval_model:
  1237. default_retrieval_model = {
  1238. "search_method": RetrievalMethod.SEMANTIC_SEARCH.value,
  1239. "reranking_enable": False,
  1240. "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},
  1241. "top_k": 2,
  1242. "score_threshold_enabled": False,
  1243. }
  1244. dataset.retrieval_model = (
  1245. knowledge_config.retrieval_model.model_dump()
  1246. if knowledge_config.retrieval_model
  1247. else default_retrieval_model
  1248. ) # type: ignore
  1249. documents = []
  1250. if knowledge_config.original_document_id:
  1251. document = DocumentService.update_document_with_dataset_id(dataset, knowledge_config, account)
  1252. documents.append(document)
  1253. batch = document.batch
  1254. else:
  1255. batch = time.strftime("%Y%m%d%H%M%S") + str(random.randint(100000, 999999))
  1256. # save process rule
  1257. if not dataset_process_rule:
  1258. process_rule = knowledge_config.process_rule
  1259. if process_rule:
  1260. if process_rule.mode in ("custom", "hierarchical"):
  1261. dataset_process_rule = DatasetProcessRule(
  1262. dataset_id=dataset.id,
  1263. mode=process_rule.mode,
  1264. rules=process_rule.rules.model_dump_json() if process_rule.rules else None,
  1265. created_by=account.id,
  1266. )
  1267. elif process_rule.mode == "automatic":
  1268. dataset_process_rule = DatasetProcessRule(
  1269. dataset_id=dataset.id,
  1270. mode=process_rule.mode,
  1271. rules=json.dumps(DatasetProcessRule.AUTOMATIC_RULES),
  1272. created_by=account.id,
  1273. )
  1274. else:
  1275. logging.warn(
  1276. f"Invalid process rule mode: {process_rule.mode}, can not find dataset process rule"
  1277. )
  1278. return
  1279. db.session.add(dataset_process_rule)
  1280. db.session.commit()
  1281. lock_name = "add_document_lock_dataset_id_{}".format(dataset.id)
  1282. with redis_client.lock(lock_name, timeout=600):
  1283. position = DocumentService.get_documents_position(dataset.id)
  1284. document_ids = []
  1285. duplicate_document_ids = []
  1286. if knowledge_config.data_source.info_list.data_source_type == "upload_file": # type: ignore
  1287. upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids # type: ignore
  1288. for file_id in upload_file_list:
  1289. file = (
  1290. db.session.query(UploadFile)
  1291. .filter(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id)
  1292. .first()
  1293. )
  1294. # raise error if file not found
  1295. if not file:
  1296. raise FileNotExistsError()
  1297. file_name = file.name
  1298. data_source_info = {
  1299. "upload_file_id": file_id,
  1300. }
  1301. # check duplicate
  1302. if knowledge_config.duplicate:
  1303. document = Document.query.filter_by(
  1304. dataset_id=dataset.id,
  1305. tenant_id=current_user.current_tenant_id,
  1306. data_source_type="upload_file",
  1307. enabled=True,
  1308. name=file_name,
  1309. ).first()
  1310. if document:
  1311. document.dataset_process_rule_id = dataset_process_rule.id # type: ignore
  1312. document.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1313. document.created_from = created_from
  1314. document.doc_form = knowledge_config.doc_form
  1315. document.doc_language = knowledge_config.doc_language
  1316. document.data_source_info = json.dumps(data_source_info)
  1317. document.batch = batch
  1318. document.indexing_status = "waiting"
  1319. db.session.add(document)
  1320. documents.append(document)
  1321. duplicate_document_ids.append(document.id)
  1322. continue
  1323. document = DocumentService.build_document(
  1324. dataset,
  1325. dataset_process_rule.id, # type: ignore
  1326. knowledge_config.data_source.info_list.data_source_type, # type: ignore
  1327. knowledge_config.doc_form,
  1328. knowledge_config.doc_language,
  1329. data_source_info,
  1330. created_from,
  1331. position,
  1332. account,
  1333. file_name,
  1334. batch,
  1335. )
  1336. db.session.add(document)
  1337. db.session.flush()
  1338. document_ids.append(document.id)
  1339. documents.append(document)
  1340. position += 1
  1341. elif knowledge_config.data_source.info_list.data_source_type == "notion_import": # type: ignore
  1342. notion_info_list = knowledge_config.data_source.info_list.notion_info_list # type: ignore
  1343. if not notion_info_list:
  1344. raise ValueError("No notion info list found.")
  1345. exist_page_ids = []
  1346. exist_document = {}
  1347. documents = Document.query.filter_by(
  1348. dataset_id=dataset.id,
  1349. tenant_id=current_user.current_tenant_id,
  1350. data_source_type="notion_import",
  1351. enabled=True,
  1352. ).all()
  1353. if documents:
  1354. for document in documents:
  1355. data_source_info = json.loads(document.data_source_info)
  1356. exist_page_ids.append(data_source_info["notion_page_id"])
  1357. exist_document[data_source_info["notion_page_id"]] = document.id
  1358. for notion_info in notion_info_list:
  1359. workspace_id = notion_info.workspace_id
  1360. data_source_binding = DataSourceOauthBinding.query.filter(
  1361. db.and_(
  1362. DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
  1363. DataSourceOauthBinding.provider == "notion",
  1364. DataSourceOauthBinding.disabled == False,
  1365. DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
  1366. )
  1367. ).first()
  1368. if not data_source_binding:
  1369. raise ValueError("Data source binding not found.")
  1370. for page in notion_info.pages:
  1371. if page.page_id not in exist_page_ids:
  1372. data_source_info = {
  1373. "notion_workspace_id": workspace_id,
  1374. "notion_page_id": page.page_id,
  1375. "notion_page_icon": page.page_icon.model_dump() if page.page_icon else None,
  1376. "type": page.type,
  1377. }
  1378. # Truncate page name to 255 characters to prevent DB field length errors
  1379. truncated_page_name = page.page_name[:255] if page.page_name else "nopagename"
  1380. document = DocumentService.build_document(
  1381. dataset,
  1382. dataset_process_rule.id, # type: ignore
  1383. knowledge_config.data_source.info_list.data_source_type, # type: ignore
  1384. knowledge_config.doc_form,
  1385. knowledge_config.doc_language,
  1386. data_source_info,
  1387. created_from,
  1388. position,
  1389. account,
  1390. truncated_page_name,
  1391. batch,
  1392. )
  1393. db.session.add(document)
  1394. db.session.flush()
  1395. document_ids.append(document.id)
  1396. documents.append(document)
  1397. position += 1
  1398. else:
  1399. exist_document.pop(page.page_id)
  1400. # delete not selected documents
  1401. if len(exist_document) > 0:
  1402. clean_notion_document_task.delay(list(exist_document.values()), dataset.id)
  1403. elif knowledge_config.data_source.info_list.data_source_type == "website_crawl": # type: ignore
  1404. website_info = knowledge_config.data_source.info_list.website_info_list # type: ignore
  1405. if not website_info:
  1406. raise ValueError("No website info list found.")
  1407. urls = website_info.urls
  1408. for url in urls:
  1409. data_source_info = {
  1410. "url": url,
  1411. "provider": website_info.provider,
  1412. "job_id": website_info.job_id,
  1413. "only_main_content": website_info.only_main_content,
  1414. "mode": "crawl",
  1415. }
  1416. if len(url) > 255:
  1417. document_name = url[:200] + "..."
  1418. else:
  1419. document_name = url
  1420. document = DocumentService.build_document(
  1421. dataset,
  1422. dataset_process_rule.id, # type: ignore
  1423. knowledge_config.data_source.info_list.data_source_type, # type: ignore
  1424. knowledge_config.doc_form,
  1425. knowledge_config.doc_language,
  1426. data_source_info,
  1427. created_from,
  1428. position,
  1429. account,
  1430. document_name,
  1431. batch,
  1432. )
  1433. db.session.add(document)
  1434. db.session.flush()
  1435. document_ids.append(document.id)
  1436. documents.append(document)
  1437. position += 1
  1438. db.session.commit()
  1439. # trigger async task
  1440. if document_ids:
  1441. document_indexing_task.delay(dataset.id, document_ids)
  1442. if duplicate_document_ids:
  1443. duplicate_document_indexing_task.delay(dataset.id, duplicate_document_ids)
  1444. return documents, batch
  1445. @staticmethod
  1446. def check_documents_upload_quota(count: int, features: FeatureModel):
  1447. can_upload_size = features.documents_upload_quota.limit - features.documents_upload_quota.size
  1448. if count > can_upload_size:
  1449. raise ValueError(
  1450. f"You have reached the limit of your subscription. Only {can_upload_size} documents can be uploaded."
  1451. )
  1452. @staticmethod
  1453. def build_document(
  1454. dataset: Dataset,
  1455. process_rule_id: str,
  1456. data_source_type: str,
  1457. document_form: str,
  1458. document_language: str,
  1459. data_source_info: dict,
  1460. created_from: str,
  1461. position: int,
  1462. account: Account,
  1463. name: str,
  1464. batch: str,
  1465. ):
  1466. if account.current_role == TenantAccountRole.EDITOR:
  1467. document = Document(
  1468. tenant_id=dataset.tenant_id,
  1469. dataset_id=dataset.id,
  1470. position=position,
  1471. data_source_type=data_source_type,
  1472. data_source_info=json.dumps(data_source_info),
  1473. dataset_process_rule_id=process_rule_id,
  1474. batch=batch,
  1475. name=name,
  1476. created_from=created_from,
  1477. created_by=account.id,
  1478. doc_form=document_form,
  1479. doc_language=document_language,
  1480. check_status=1,
  1481. enable_applicant=account.id,
  1482. )
  1483. else:
  1484. document = Document(
  1485. tenant_id=dataset.tenant_id,
  1486. dataset_id=dataset.id,
  1487. position=position,
  1488. data_source_type=data_source_type,
  1489. data_source_info=json.dumps(data_source_info),
  1490. dataset_process_rule_id=process_rule_id,
  1491. batch=batch,
  1492. name=name,
  1493. created_from=created_from,
  1494. created_by=account.id,
  1495. doc_form=document_form,
  1496. doc_language=document_language,
  1497. check_status=3,
  1498. check_by=account.id,
  1499. enable_applicant=account.id,
  1500. enabled=True,
  1501. )
  1502. doc_metadata = {}
  1503. if dataset.built_in_field_enabled:
  1504. doc_metadata = {
  1505. BuiltInField.document_name: name,
  1506. BuiltInField.uploader: account.name,
  1507. BuiltInField.upload_date: datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%d %H:%M:%S"),
  1508. BuiltInField.last_update_date: datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%d %H:%M:%S"),
  1509. BuiltInField.source: data_source_type,
  1510. }
  1511. if doc_metadata:
  1512. document.doc_metadata = doc_metadata
  1513. return document
  1514. @staticmethod
  1515. def get_tenant_documents_count():
  1516. documents_count = Document.query.filter(
  1517. Document.completed_at.isnot(None),
  1518. Document.enabled == True,
  1519. Document.archived == False,
  1520. Document.tenant_id == current_user.current_tenant_id,
  1521. ).count()
  1522. return documents_count
  1523. @staticmethod
  1524. def update_document_with_dataset_id(
  1525. dataset: Dataset,
  1526. document_data: KnowledgeConfig,
  1527. account: Account,
  1528. dataset_process_rule: Optional[DatasetProcessRule] = None,
  1529. created_from: str = "web",
  1530. ):
  1531. DatasetService.check_dataset_model_setting(dataset)
  1532. document = DocumentService.get_document(dataset.id, document_data.original_document_id)
  1533. if document is None:
  1534. raise NotFound("Document not found")
  1535. if document.display_status != "available":
  1536. raise ValueError("Document is not available")
  1537. # save process rule
  1538. if document_data.process_rule:
  1539. process_rule = document_data.process_rule
  1540. if process_rule.mode in {"custom", "hierarchical"}:
  1541. dataset_process_rule = DatasetProcessRule(
  1542. dataset_id=dataset.id,
  1543. mode=process_rule.mode,
  1544. rules=process_rule.rules.model_dump_json() if process_rule.rules else None,
  1545. created_by=account.id,
  1546. )
  1547. elif process_rule.mode == "automatic":
  1548. dataset_process_rule = DatasetProcessRule(
  1549. dataset_id=dataset.id,
  1550. mode=process_rule.mode,
  1551. rules=json.dumps(DatasetProcessRule.AUTOMATIC_RULES),
  1552. created_by=account.id,
  1553. )
  1554. if dataset_process_rule is not None:
  1555. db.session.add(dataset_process_rule)
  1556. db.session.commit()
  1557. document.dataset_process_rule_id = dataset_process_rule.id
  1558. # update document data source
  1559. if document_data.data_source:
  1560. file_name = ""
  1561. data_source_info = {}
  1562. if document_data.data_source.info_list.data_source_type == "upload_file":
  1563. if not document_data.data_source.info_list.file_info_list:
  1564. raise ValueError("No file info list found.")
  1565. upload_file_list = document_data.data_source.info_list.file_info_list.file_ids
  1566. for file_id in upload_file_list:
  1567. file = (
  1568. db.session.query(UploadFile)
  1569. .filter(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id)
  1570. .first()
  1571. )
  1572. # raise error if file not found
  1573. if not file:
  1574. raise FileNotExistsError()
  1575. file_name = file.name
  1576. data_source_info = {
  1577. "upload_file_id": file_id,
  1578. }
  1579. elif document_data.data_source.info_list.data_source_type == "notion_import":
  1580. if not document_data.data_source.info_list.notion_info_list:
  1581. raise ValueError("No notion info list found.")
  1582. notion_info_list = document_data.data_source.info_list.notion_info_list
  1583. for notion_info in notion_info_list:
  1584. workspace_id = notion_info.workspace_id
  1585. data_source_binding = DataSourceOauthBinding.query.filter(
  1586. db.and_(
  1587. DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
  1588. DataSourceOauthBinding.provider == "notion",
  1589. DataSourceOauthBinding.disabled == False,
  1590. DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
  1591. )
  1592. ).first()
  1593. if not data_source_binding:
  1594. raise ValueError("Data source binding not found.")
  1595. for page in notion_info.pages:
  1596. data_source_info = {
  1597. "notion_workspace_id": workspace_id,
  1598. "notion_page_id": page.page_id,
  1599. "notion_page_icon": page.page_icon.model_dump() if page.page_icon else None, # type: ignore
  1600. "type": page.type,
  1601. }
  1602. elif document_data.data_source.info_list.data_source_type == "website_crawl":
  1603. website_info = document_data.data_source.info_list.website_info_list
  1604. if website_info:
  1605. urls = website_info.urls
  1606. for url in urls:
  1607. data_source_info = {
  1608. "url": url,
  1609. "provider": website_info.provider,
  1610. "job_id": website_info.job_id,
  1611. "only_main_content": website_info.only_main_content, # type: ignore
  1612. "mode": "crawl",
  1613. }
  1614. document.data_source_type = document_data.data_source.info_list.data_source_type
  1615. document.data_source_info = json.dumps(data_source_info)
  1616. document.name = file_name
  1617. # update document name
  1618. if document_data.name:
  1619. document.name = document_data.name
  1620. # update document to be waiting
  1621. document.indexing_status = "waiting"
  1622. document.completed_at = None
  1623. document.processing_started_at = None
  1624. document.parsing_completed_at = None
  1625. document.cleaning_completed_at = None
  1626. document.splitting_completed_at = None
  1627. document.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1628. document.created_from = created_from
  1629. document.doc_form = document_data.doc_form
  1630. db.session.add(document)
  1631. db.session.commit()
  1632. # update document segment
  1633. update_params = {DocumentSegment.status: "re_segment"}
  1634. DocumentSegment.query.filter_by(document_id=document.id).update(update_params)
  1635. db.session.commit()
  1636. # trigger async task
  1637. document_indexing_update_task.delay(document.dataset_id, document.id)
  1638. return document
  1639. @staticmethod
  1640. def save_document_without_dataset_id(tenant_id: str, knowledge_config: KnowledgeConfig, account: Account):
  1641. features = FeatureService.get_features(current_user.current_tenant_id)
  1642. if features.billing.enabled:
  1643. count = 0
  1644. if knowledge_config.data_source.info_list.data_source_type == "upload_file": # type: ignore
  1645. upload_file_list = (
  1646. knowledge_config.data_source.info_list.file_info_list.file_ids # type: ignore
  1647. if knowledge_config.data_source.info_list.file_info_list # type: ignore
  1648. else []
  1649. )
  1650. count = len(upload_file_list)
  1651. elif knowledge_config.data_source.info_list.data_source_type == "notion_import": # type: ignore
  1652. notion_info_list = knowledge_config.data_source.info_list.notion_info_list # type: ignore
  1653. if notion_info_list:
  1654. for notion_info in notion_info_list:
  1655. count = count + len(notion_info.pages)
  1656. elif knowledge_config.data_source.info_list.data_source_type == "website_crawl": # type: ignore
  1657. website_info = knowledge_config.data_source.info_list.website_info_list # type: ignore
  1658. if website_info:
  1659. count = len(website_info.urls)
  1660. batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT)
  1661. if count > batch_upload_limit:
  1662. raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")
  1663. DocumentService.check_documents_upload_quota(count, features)
  1664. dataset_collection_binding_id = None
  1665. retrieval_model = None
  1666. if knowledge_config.indexing_technique == "high_quality":
  1667. dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
  1668. knowledge_config.embedding_model_provider, # type: ignore
  1669. knowledge_config.embedding_model, # type: ignore
  1670. )
  1671. dataset_collection_binding_id = dataset_collection_binding.id
  1672. if knowledge_config.retrieval_model:
  1673. retrieval_model = knowledge_config.retrieval_model
  1674. else:
  1675. retrieval_model = RetrievalModel(
  1676. search_method=RetrievalMethod.SEMANTIC_SEARCH.value,
  1677. reranking_enable=False,
  1678. reranking_model=RerankingModel(reranking_provider_name="", reranking_model_name=""),
  1679. top_k=2,
  1680. score_threshold_enabled=False,
  1681. )
  1682. # save dataset
  1683. dataset = Dataset(
  1684. tenant_id=tenant_id,
  1685. name="",
  1686. data_source_type=knowledge_config.data_source.info_list.data_source_type, # type: ignore
  1687. indexing_technique=knowledge_config.indexing_technique,
  1688. created_by=account.id,
  1689. embedding_model=knowledge_config.embedding_model,
  1690. embedding_model_provider=knowledge_config.embedding_model_provider,
  1691. collection_binding_id=dataset_collection_binding_id,
  1692. retrieval_model=retrieval_model.model_dump() if retrieval_model else None,
  1693. )
  1694. db.session.add(dataset) # type: ignore
  1695. db.session.flush()
  1696. documents, batch = DocumentService.save_document_with_dataset_id(dataset, knowledge_config, account)
  1697. cut_length = 18
  1698. cut_name = documents[0].name[:cut_length]
  1699. dataset.name = cut_name + "..."
  1700. dataset.description = "useful for when you want to answer queries about the " + documents[0].name
  1701. db.session.commit()
  1702. return dataset, documents, batch
  1703. @classmethod
  1704. def document_create_args_validate(cls, knowledge_config: KnowledgeConfig):
  1705. if not knowledge_config.data_source and not knowledge_config.process_rule:
  1706. raise ValueError("Data source or Process rule is required")
  1707. else:
  1708. if knowledge_config.data_source:
  1709. DocumentService.data_source_args_validate(knowledge_config)
  1710. if knowledge_config.process_rule:
  1711. DocumentService.process_rule_args_validate(knowledge_config)
  1712. @classmethod
  1713. def data_source_args_validate(cls, knowledge_config: KnowledgeConfig):
  1714. if not knowledge_config.data_source:
  1715. raise ValueError("Data source is required")
  1716. if knowledge_config.data_source.info_list.data_source_type not in Document.DATA_SOURCES:
  1717. raise ValueError("Data source type is invalid")
  1718. if not knowledge_config.data_source.info_list:
  1719. raise ValueError("Data source info is required")
  1720. if knowledge_config.data_source.info_list.data_source_type == "upload_file":
  1721. if not knowledge_config.data_source.info_list.file_info_list:
  1722. raise ValueError("File source info is required")
  1723. if knowledge_config.data_source.info_list.data_source_type == "notion_import":
  1724. if not knowledge_config.data_source.info_list.notion_info_list:
  1725. raise ValueError("Notion source info is required")
  1726. if knowledge_config.data_source.info_list.data_source_type == "website_crawl":
  1727. if not knowledge_config.data_source.info_list.website_info_list:
  1728. raise ValueError("Website source info is required")
  1729. @classmethod
  1730. def process_rule_args_validate(cls, knowledge_config: KnowledgeConfig):
  1731. if not knowledge_config.process_rule:
  1732. raise ValueError("Process rule is required")
  1733. if not knowledge_config.process_rule.mode:
  1734. raise ValueError("Process rule mode is required")
  1735. if knowledge_config.process_rule.mode not in DatasetProcessRule.MODES:
  1736. raise ValueError("Process rule mode is invalid")
  1737. if knowledge_config.process_rule.mode == "automatic":
  1738. knowledge_config.process_rule.rules = None
  1739. else:
  1740. if not knowledge_config.process_rule.rules:
  1741. raise ValueError("Process rule rules is required")
  1742. if knowledge_config.process_rule.rules.pre_processing_rules is None:
  1743. raise ValueError("Process rule pre_processing_rules is required")
  1744. unique_pre_processing_rule_dicts = {}
  1745. for pre_processing_rule in knowledge_config.process_rule.rules.pre_processing_rules:
  1746. if not pre_processing_rule.id:
  1747. raise ValueError("Process rule pre_processing_rules id is required")
  1748. if not isinstance(pre_processing_rule.enabled, bool):
  1749. raise ValueError("Process rule pre_processing_rules enabled is invalid")
  1750. unique_pre_processing_rule_dicts[pre_processing_rule.id] = pre_processing_rule
  1751. knowledge_config.process_rule.rules.pre_processing_rules = list(unique_pre_processing_rule_dicts.values())
  1752. if not knowledge_config.process_rule.rules.segmentation:
  1753. raise ValueError("Process rule segmentation is required")
  1754. if not knowledge_config.process_rule.rules.segmentation.separator:
  1755. raise ValueError("Process rule segmentation separator is required")
  1756. if not isinstance(knowledge_config.process_rule.rules.segmentation.separator, str):
  1757. raise ValueError("Process rule segmentation separator is invalid")
  1758. if not (
  1759. knowledge_config.process_rule.mode == "hierarchical"
  1760. and knowledge_config.process_rule.rules.parent_mode == "full-doc"
  1761. ):
  1762. if not knowledge_config.process_rule.rules.segmentation.max_tokens:
  1763. raise ValueError("Process rule segmentation max_tokens is required")
  1764. if not isinstance(knowledge_config.process_rule.rules.segmentation.max_tokens, int):
  1765. raise ValueError("Process rule segmentation max_tokens is invalid")
  1766. @classmethod
  1767. def estimate_args_validate(cls, args: dict):
  1768. if "info_list" not in args or not args["info_list"]:
  1769. raise ValueError("Data source info is required")
  1770. if not isinstance(args["info_list"], dict):
  1771. raise ValueError("Data info is invalid")
  1772. if "process_rule" not in args or not args["process_rule"]:
  1773. raise ValueError("Process rule is required")
  1774. if not isinstance(args["process_rule"], dict):
  1775. raise ValueError("Process rule is invalid")
  1776. if "mode" not in args["process_rule"] or not args["process_rule"]["mode"]:
  1777. raise ValueError("Process rule mode is required")
  1778. if args["process_rule"]["mode"] not in DatasetProcessRule.MODES:
  1779. raise ValueError("Process rule mode is invalid")
  1780. if args["process_rule"]["mode"] == "automatic":
  1781. args["process_rule"]["rules"] = {}
  1782. else:
  1783. if "rules" not in args["process_rule"] or not args["process_rule"]["rules"]:
  1784. raise ValueError("Process rule rules is required")
  1785. if not isinstance(args["process_rule"]["rules"], dict):
  1786. raise ValueError("Process rule rules is invalid")
  1787. if (
  1788. "pre_processing_rules" not in args["process_rule"]["rules"]
  1789. or args["process_rule"]["rules"]["pre_processing_rules"] is None
  1790. ):
  1791. raise ValueError("Process rule pre_processing_rules is required")
  1792. if not isinstance(args["process_rule"]["rules"]["pre_processing_rules"], list):
  1793. raise ValueError("Process rule pre_processing_rules is invalid")
  1794. unique_pre_processing_rule_dicts = {}
  1795. for pre_processing_rule in args["process_rule"]["rules"]["pre_processing_rules"]:
  1796. if "id" not in pre_processing_rule or not pre_processing_rule["id"]:
  1797. raise ValueError("Process rule pre_processing_rules id is required")
  1798. if pre_processing_rule["id"] not in DatasetProcessRule.PRE_PROCESSING_RULES:
  1799. raise ValueError("Process rule pre_processing_rules id is invalid")
  1800. if "enabled" not in pre_processing_rule or pre_processing_rule["enabled"] is None:
  1801. raise ValueError("Process rule pre_processing_rules enabled is required")
  1802. if not isinstance(pre_processing_rule["enabled"], bool):
  1803. raise ValueError("Process rule pre_processing_rules enabled is invalid")
  1804. unique_pre_processing_rule_dicts[pre_processing_rule["id"]] = pre_processing_rule
  1805. args["process_rule"]["rules"]["pre_processing_rules"] = list(unique_pre_processing_rule_dicts.values())
  1806. if (
  1807. "segmentation" not in args["process_rule"]["rules"]
  1808. or args["process_rule"]["rules"]["segmentation"] is None
  1809. ):
  1810. raise ValueError("Process rule segmentation is required")
  1811. if not isinstance(args["process_rule"]["rules"]["segmentation"], dict):
  1812. raise ValueError("Process rule segmentation is invalid")
  1813. if (
  1814. "separator" not in args["process_rule"]["rules"]["segmentation"]
  1815. or not args["process_rule"]["rules"]["segmentation"]["separator"]
  1816. ):
  1817. raise ValueError("Process rule segmentation separator is required")
  1818. if not isinstance(args["process_rule"]["rules"]["segmentation"]["separator"], str):
  1819. raise ValueError("Process rule segmentation separator is invalid")
  1820. if (
  1821. "max_tokens" not in args["process_rule"]["rules"]["segmentation"]
  1822. or not args["process_rule"]["rules"]["segmentation"]["max_tokens"]
  1823. ):
  1824. raise ValueError("Process rule segmentation max_tokens is required")
  1825. if not isinstance(args["process_rule"]["rules"]["segmentation"]["max_tokens"], int):
  1826. raise ValueError("Process rule segmentation max_tokens is invalid")
  1827. class SegmentService:
  1828. @classmethod
  1829. def segment_create_args_validate(cls, args: dict, document: Document):
  1830. if document.doc_form == "qa_model":
  1831. if "answer" not in args or not args["answer"]:
  1832. raise ValueError("Answer is required")
  1833. if not args["answer"].strip():
  1834. raise ValueError("Answer is empty")
  1835. if "content" not in args or not args["content"] or not args["content"].strip():
  1836. raise ValueError("Content is empty")
  1837. @classmethod
  1838. def create_segment(cls, args: dict, document: Document, dataset: Dataset):
  1839. content = args["content"]
  1840. doc_id = str(uuid.uuid4())
  1841. segment_hash = helper.generate_text_hash(content)
  1842. tokens = 0
  1843. if dataset.indexing_technique == "high_quality":
  1844. model_manager = ModelManager()
  1845. embedding_model = model_manager.get_model_instance(
  1846. tenant_id=current_user.current_tenant_id,
  1847. provider=dataset.embedding_model_provider,
  1848. model_type=ModelType.TEXT_EMBEDDING,
  1849. model=dataset.embedding_model,
  1850. )
  1851. # calc embedding use tokens
  1852. tokens = embedding_model.get_text_embedding_num_tokens(texts=[content])[0]
  1853. lock_name = "add_segment_lock_document_id_{}".format(document.id)
  1854. with redis_client.lock(lock_name, timeout=600):
  1855. max_position = (
  1856. db.session.query(func.max(DocumentSegment.position))
  1857. .filter(DocumentSegment.document_id == document.id)
  1858. .scalar()
  1859. )
  1860. segment_document = DocumentSegment(
  1861. tenant_id=current_user.current_tenant_id,
  1862. dataset_id=document.dataset_id,
  1863. document_id=document.id,
  1864. index_node_id=doc_id,
  1865. index_node_hash=segment_hash,
  1866. position=max_position + 1 if max_position else 1,
  1867. content=content,
  1868. word_count=len(content),
  1869. tokens=tokens,
  1870. status="completed",
  1871. indexing_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
  1872. completed_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
  1873. created_by=current_user.id,
  1874. )
  1875. if document.doc_form == "qa_model":
  1876. segment_document.word_count += len(args["answer"])
  1877. segment_document.answer = args["answer"]
  1878. db.session.add(segment_document)
  1879. # update document word count
  1880. document.word_count += segment_document.word_count
  1881. db.session.add(document)
  1882. db.session.commit()
  1883. # save vector index
  1884. try:
  1885. VectorService.create_segments_vector([args["keywords"]], [segment_document], dataset, document.doc_form)
  1886. except Exception as e:
  1887. logging.exception("create segment index failed")
  1888. segment_document.enabled = False
  1889. segment_document.disabled_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1890. segment_document.status = "error"
  1891. segment_document.error = str(e)
  1892. db.session.commit()
  1893. segment = db.session.query(DocumentSegment).filter(DocumentSegment.id == segment_document.id).first()
  1894. return segment
  1895. @classmethod
  1896. def multi_create_segment(cls, segments: list, document: Document, dataset: Dataset):
  1897. lock_name = "multi_add_segment_lock_document_id_{}".format(document.id)
  1898. increment_word_count = 0
  1899. with redis_client.lock(lock_name, timeout=600):
  1900. embedding_model = None
  1901. if dataset.indexing_technique == "high_quality":
  1902. model_manager = ModelManager()
  1903. embedding_model = model_manager.get_model_instance(
  1904. tenant_id=current_user.current_tenant_id,
  1905. provider=dataset.embedding_model_provider,
  1906. model_type=ModelType.TEXT_EMBEDDING,
  1907. model=dataset.embedding_model,
  1908. )
  1909. max_position = (
  1910. db.session.query(func.max(DocumentSegment.position))
  1911. .filter(DocumentSegment.document_id == document.id)
  1912. .scalar()
  1913. )
  1914. pre_segment_data_list = []
  1915. segment_data_list = []
  1916. keywords_list = []
  1917. position = max_position + 1 if max_position else 1
  1918. for segment_item in segments:
  1919. content = segment_item["content"]
  1920. doc_id = str(uuid.uuid4())
  1921. segment_hash = helper.generate_text_hash(content)
  1922. tokens = 0
  1923. if dataset.indexing_technique == "high_quality" and embedding_model:
  1924. # calc embedding use tokens
  1925. if document.doc_form == "qa_model":
  1926. tokens = embedding_model.get_text_embedding_num_tokens(
  1927. texts=[content + segment_item["answer"]]
  1928. )[0]
  1929. else:
  1930. tokens = embedding_model.get_text_embedding_num_tokens(texts=[content])[0]
  1931. segment_document = DocumentSegment(
  1932. tenant_id=current_user.current_tenant_id,
  1933. dataset_id=document.dataset_id,
  1934. document_id=document.id,
  1935. index_node_id=doc_id,
  1936. index_node_hash=segment_hash,
  1937. position=position,
  1938. content=content,
  1939. word_count=len(content),
  1940. tokens=tokens,
  1941. status="completed",
  1942. indexing_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
  1943. completed_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
  1944. created_by=current_user.id,
  1945. )
  1946. if document.doc_form == "qa_model":
  1947. segment_document.answer = segment_item["answer"]
  1948. segment_document.word_count += len(segment_item["answer"])
  1949. increment_word_count += segment_document.word_count
  1950. db.session.add(segment_document)
  1951. segment_data_list.append(segment_document)
  1952. position += 1
  1953. pre_segment_data_list.append(segment_document)
  1954. if "keywords" in segment_item:
  1955. keywords_list.append(segment_item["keywords"])
  1956. else:
  1957. keywords_list.append(None)
  1958. # update document word count
  1959. document.word_count += increment_word_count
  1960. db.session.add(document)
  1961. try:
  1962. # save vector index
  1963. VectorService.create_segments_vector(keywords_list, pre_segment_data_list, dataset, document.doc_form)
  1964. except Exception as e:
  1965. logging.exception("create segment index failed")
  1966. for segment_document in segment_data_list:
  1967. segment_document.enabled = False
  1968. segment_document.disabled_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1969. segment_document.status = "error"
  1970. segment_document.error = str(e)
  1971. db.session.commit()
  1972. return segment_data_list
  1973. @classmethod
  1974. def update_segment(cls, args: SegmentUpdateArgs, segment: DocumentSegment, document: Document, dataset: Dataset):
  1975. indexing_cache_key = "segment_{}_indexing".format(segment.id)
  1976. cache_result = redis_client.get(indexing_cache_key)
  1977. if cache_result is not None:
  1978. raise ValueError("Segment is indexing, please try again later")
  1979. if args.enabled is not None:
  1980. action = args.enabled
  1981. if segment.enabled != action:
  1982. if not action:
  1983. segment.enabled = action
  1984. segment.disabled_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1985. segment.disabled_by = current_user.id
  1986. db.session.add(segment)
  1987. db.session.commit()
  1988. # Set cache to prevent indexing the same segment multiple times
  1989. redis_client.setex(indexing_cache_key, 600, 1)
  1990. disable_segment_from_index_task.delay(segment.id)
  1991. return segment
  1992. if not segment.enabled:
  1993. if args.enabled is not None:
  1994. if not args.enabled:
  1995. raise ValueError("Can't update disabled segment")
  1996. else:
  1997. raise ValueError("Can't update disabled segment")
  1998. try:
  1999. word_count_change = segment.word_count
  2000. content = args.content or segment.content
  2001. if segment.content == content:
  2002. segment.word_count = len(content)
  2003. if document.doc_form == "qa_model":
  2004. segment.answer = args.answer
  2005. segment.word_count += len(args.answer) if args.answer else 0
  2006. word_count_change = segment.word_count - word_count_change
  2007. keyword_changed = False
  2008. if args.keywords:
  2009. if Counter(segment.keywords) != Counter(args.keywords):
  2010. segment.keywords = args.keywords
  2011. keyword_changed = True
  2012. segment.enabled = True
  2013. segment.disabled_at = None
  2014. segment.disabled_by = None
  2015. db.session.add(segment)
  2016. db.session.commit()
  2017. # update document word count
  2018. if word_count_change != 0:
  2019. document.word_count = max(0, document.word_count + word_count_change)
  2020. db.session.add(document)
  2021. # update segment index task
  2022. if document.doc_form == IndexType.PARENT_CHILD_INDEX and args.regenerate_child_chunks:
  2023. # regenerate child chunks
  2024. # get embedding model instance
  2025. if dataset.indexing_technique == "high_quality":
  2026. # check embedding model setting
  2027. model_manager = ModelManager()
  2028. if dataset.embedding_model_provider:
  2029. embedding_model_instance = model_manager.get_model_instance(
  2030. tenant_id=dataset.tenant_id,
  2031. provider=dataset.embedding_model_provider,
  2032. model_type=ModelType.TEXT_EMBEDDING,
  2033. model=dataset.embedding_model,
  2034. )
  2035. else:
  2036. embedding_model_instance = model_manager.get_default_model_instance(
  2037. tenant_id=dataset.tenant_id,
  2038. model_type=ModelType.TEXT_EMBEDDING,
  2039. )
  2040. else:
  2041. raise ValueError("The knowledge base index technique is not high quality!")
  2042. # get the process rule
  2043. processing_rule = (
  2044. db.session.query(DatasetProcessRule)
  2045. .filter(DatasetProcessRule.id == document.dataset_process_rule_id)
  2046. .first()
  2047. )
  2048. if not processing_rule:
  2049. raise ValueError("No processing rule found.")
  2050. VectorService.generate_child_chunks(
  2051. segment, document, dataset, embedding_model_instance, processing_rule, True
  2052. )
  2053. elif document.doc_form in (IndexType.PARAGRAPH_INDEX, IndexType.QA_INDEX):
  2054. if args.enabled or keyword_changed:
  2055. VectorService.create_segments_vector(
  2056. [args.keywords] if args.keywords else None,
  2057. [segment],
  2058. dataset,
  2059. document.doc_form,
  2060. )
  2061. else:
  2062. segment_hash = helper.generate_text_hash(content)
  2063. tokens = 0
  2064. if dataset.indexing_technique == "high_quality":
  2065. model_manager = ModelManager()
  2066. embedding_model = model_manager.get_model_instance(
  2067. tenant_id=current_user.current_tenant_id,
  2068. provider=dataset.embedding_model_provider,
  2069. model_type=ModelType.TEXT_EMBEDDING,
  2070. model=dataset.embedding_model,
  2071. )
  2072. # calc embedding use tokens
  2073. if document.doc_form == "qa_model":
  2074. tokens = embedding_model.get_text_embedding_num_tokens(texts=[content + segment.answer])[0]
  2075. else:
  2076. tokens = embedding_model.get_text_embedding_num_tokens(texts=[content])[0]
  2077. segment.content = content
  2078. segment.index_node_hash = segment_hash
  2079. segment.word_count = len(content)
  2080. segment.tokens = tokens
  2081. segment.status = "completed"
  2082. segment.indexing_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  2083. segment.completed_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  2084. segment.updated_by = current_user.id
  2085. segment.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  2086. segment.enabled = True
  2087. segment.disabled_at = None
  2088. segment.disabled_by = None
  2089. if document.doc_form == "qa_model":
  2090. segment.answer = args.answer
  2091. segment.word_count += len(args.answer) if args.answer else 0
  2092. word_count_change = segment.word_count - word_count_change
  2093. # update document word count
  2094. if word_count_change != 0:
  2095. document.word_count = max(0, document.word_count + word_count_change)
  2096. db.session.add(document)
  2097. db.session.add(segment)
  2098. db.session.commit()
  2099. if document.doc_form == IndexType.PARENT_CHILD_INDEX and args.regenerate_child_chunks:
  2100. # get embedding model instance
  2101. if dataset.indexing_technique == "high_quality":
  2102. # check embedding model setting
  2103. model_manager = ModelManager()
  2104. if dataset.embedding_model_provider:
  2105. embedding_model_instance = model_manager.get_model_instance(
  2106. tenant_id=dataset.tenant_id,
  2107. provider=dataset.embedding_model_provider,
  2108. model_type=ModelType.TEXT_EMBEDDING,
  2109. model=dataset.embedding_model,
  2110. )
  2111. else:
  2112. embedding_model_instance = model_manager.get_default_model_instance(
  2113. tenant_id=dataset.tenant_id,
  2114. model_type=ModelType.TEXT_EMBEDDING,
  2115. )
  2116. else:
  2117. raise ValueError("The knowledge base index technique is not high quality!")
  2118. # get the process rule
  2119. processing_rule = (
  2120. db.session.query(DatasetProcessRule)
  2121. .filter(DatasetProcessRule.id == document.dataset_process_rule_id)
  2122. .first()
  2123. )
  2124. if not processing_rule:
  2125. raise ValueError("No processing rule found.")
  2126. VectorService.generate_child_chunks(
  2127. segment, document, dataset, embedding_model_instance, processing_rule, True
  2128. )
  2129. elif document.doc_form in (IndexType.PARAGRAPH_INDEX, IndexType.QA_INDEX):
  2130. # update segment vector index
  2131. VectorService.update_segment_vector(args.keywords, segment, dataset)
  2132. except Exception as e:
  2133. logging.exception("update segment index failed")
  2134. segment.enabled = False
  2135. segment.disabled_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  2136. segment.status = "error"
  2137. segment.error = str(e)
  2138. db.session.commit()
  2139. new_segment = db.session.query(DocumentSegment).filter(DocumentSegment.id == segment.id).first()
  2140. return new_segment
  2141. @classmethod
  2142. def delete_segment(cls, segment: DocumentSegment, document: Document, dataset: Dataset):
  2143. indexing_cache_key = "segment_{}_delete_indexing".format(segment.id)
  2144. cache_result = redis_client.get(indexing_cache_key)
  2145. if cache_result is not None:
  2146. raise ValueError("Segment is deleting.")
  2147. # enabled segment need to delete index
  2148. if segment.enabled:
  2149. # send delete segment index task
  2150. redis_client.setex(indexing_cache_key, 600, 1)
  2151. delete_segment_from_index_task.delay([segment.index_node_id], dataset.id, document.id)
  2152. db.session.delete(segment)
  2153. # update document word count
  2154. document.word_count -= segment.word_count
  2155. db.session.add(document)
  2156. db.session.commit()
  2157. @classmethod
  2158. def delete_segments(cls, segment_ids: list, document: Document, dataset: Dataset):
  2159. index_node_ids = (
  2160. DocumentSegment.query.with_entities(DocumentSegment.index_node_id)
  2161. .filter(
  2162. DocumentSegment.id.in_(segment_ids),
  2163. DocumentSegment.dataset_id == dataset.id,
  2164. DocumentSegment.document_id == document.id,
  2165. DocumentSegment.tenant_id == current_user.current_tenant_id,
  2166. )
  2167. .all()
  2168. )
  2169. index_node_ids = [index_node_id[0] for index_node_id in index_node_ids]
  2170. delete_segment_from_index_task.delay(index_node_ids, dataset.id, document.id)
  2171. db.session.query(DocumentSegment).filter(DocumentSegment.id.in_(segment_ids)).delete()
  2172. db.session.commit()
  2173. @classmethod
  2174. def update_segments_status(cls, segment_ids: list, action: str, dataset: Dataset, document: Document):
  2175. if action == "enable":
  2176. segments = (
  2177. db.session.query(DocumentSegment)
  2178. .filter(
  2179. DocumentSegment.id.in_(segment_ids),
  2180. DocumentSegment.dataset_id == dataset.id,
  2181. DocumentSegment.document_id == document.id,
  2182. DocumentSegment.enabled == False,
  2183. )
  2184. .all()
  2185. )
  2186. if not segments:
  2187. return
  2188. real_deal_segmment_ids = []
  2189. for segment in segments:
  2190. indexing_cache_key = "segment_{}_indexing".format(segment.id)
  2191. cache_result = redis_client.get(indexing_cache_key)
  2192. if cache_result is not None:
  2193. continue
  2194. segment.enabled = True
  2195. segment.disabled_at = None
  2196. segment.disabled_by = None
  2197. db.session.add(segment)
  2198. real_deal_segmment_ids.append(segment.id)
  2199. db.session.commit()
  2200. enable_segments_to_index_task.delay(real_deal_segmment_ids, dataset.id, document.id)
  2201. elif action == "disable":
  2202. segments = (
  2203. db.session.query(DocumentSegment)
  2204. .filter(
  2205. DocumentSegment.id.in_(segment_ids),
  2206. DocumentSegment.dataset_id == dataset.id,
  2207. DocumentSegment.document_id == document.id,
  2208. DocumentSegment.enabled == True,
  2209. )
  2210. .all()
  2211. )
  2212. if not segments:
  2213. return
  2214. real_deal_segmment_ids = []
  2215. for segment in segments:
  2216. indexing_cache_key = "segment_{}_indexing".format(segment.id)
  2217. cache_result = redis_client.get(indexing_cache_key)
  2218. if cache_result is not None:
  2219. continue
  2220. segment.enabled = False
  2221. segment.disabled_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  2222. segment.disabled_by = current_user.id
  2223. db.session.add(segment)
  2224. real_deal_segmment_ids.append(segment.id)
  2225. db.session.commit()
  2226. disable_segments_from_index_task.delay(real_deal_segmment_ids, dataset.id, document.id)
  2227. else:
  2228. raise InvalidActionError()
  2229. @classmethod
  2230. def create_child_chunk(
  2231. cls, content: str, segment: DocumentSegment, document: Document, dataset: Dataset
  2232. ) -> ChildChunk:
  2233. lock_name = "add_child_lock_{}".format(segment.id)
  2234. with redis_client.lock(lock_name, timeout=20):
  2235. index_node_id = str(uuid.uuid4())
  2236. index_node_hash = helper.generate_text_hash(content)
  2237. child_chunk_count = (
  2238. db.session.query(ChildChunk)
  2239. .filter(
  2240. ChildChunk.tenant_id == current_user.current_tenant_id,
  2241. ChildChunk.dataset_id == dataset.id,
  2242. ChildChunk.document_id == document.id,
  2243. ChildChunk.segment_id == segment.id,
  2244. )
  2245. .count()
  2246. )
  2247. max_position = (
  2248. db.session.query(func.max(ChildChunk.position))
  2249. .filter(
  2250. ChildChunk.tenant_id == current_user.current_tenant_id,
  2251. ChildChunk.dataset_id == dataset.id,
  2252. ChildChunk.document_id == document.id,
  2253. ChildChunk.segment_id == segment.id,
  2254. )
  2255. .scalar()
  2256. )
  2257. child_chunk = ChildChunk(
  2258. tenant_id=current_user.current_tenant_id,
  2259. dataset_id=dataset.id,
  2260. document_id=document.id,
  2261. segment_id=segment.id,
  2262. position=max_position + 1,
  2263. index_node_id=index_node_id,
  2264. index_node_hash=index_node_hash,
  2265. content=content,
  2266. word_count=len(content),
  2267. type="customized",
  2268. created_by=current_user.id,
  2269. )
  2270. db.session.add(child_chunk)
  2271. # save vector index
  2272. try:
  2273. VectorService.create_child_chunk_vector(child_chunk, dataset)
  2274. except Exception as e:
  2275. logging.exception("create child chunk index failed")
  2276. db.session.rollback()
  2277. raise ChildChunkIndexingError(str(e))
  2278. db.session.commit()
  2279. return child_chunk
  2280. @classmethod
  2281. def update_child_chunks(
  2282. cls,
  2283. child_chunks_update_args: list[ChildChunkUpdateArgs],
  2284. segment: DocumentSegment,
  2285. document: Document,
  2286. dataset: Dataset,
  2287. ) -> list[ChildChunk]:
  2288. child_chunks = (
  2289. db.session.query(ChildChunk)
  2290. .filter(
  2291. ChildChunk.dataset_id == dataset.id,
  2292. ChildChunk.document_id == document.id,
  2293. ChildChunk.segment_id == segment.id,
  2294. )
  2295. .all()
  2296. )
  2297. child_chunks_map = {chunk.id: chunk for chunk in child_chunks}
  2298. new_child_chunks, update_child_chunks, delete_child_chunks, new_child_chunks_args = [], [], [], []
  2299. for child_chunk_update_args in child_chunks_update_args:
  2300. if child_chunk_update_args.id:
  2301. child_chunk = child_chunks_map.pop(child_chunk_update_args.id, None)
  2302. if child_chunk:
  2303. if child_chunk.content != child_chunk_update_args.content:
  2304. child_chunk.content = child_chunk_update_args.content
  2305. child_chunk.word_count = len(child_chunk.content)
  2306. child_chunk.updated_by = current_user.id
  2307. child_chunk.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  2308. child_chunk.type = "customized"
  2309. update_child_chunks.append(child_chunk)
  2310. else:
  2311. new_child_chunks_args.append(child_chunk_update_args)
  2312. if child_chunks_map:
  2313. delete_child_chunks = list(child_chunks_map.values())
  2314. try:
  2315. if update_child_chunks:
  2316. db.session.bulk_save_objects(update_child_chunks)
  2317. if delete_child_chunks:
  2318. for child_chunk in delete_child_chunks:
  2319. db.session.delete(child_chunk)
  2320. if new_child_chunks_args:
  2321. child_chunk_count = len(child_chunks)
  2322. for position, args in enumerate(new_child_chunks_args, start=child_chunk_count + 1):
  2323. index_node_id = str(uuid.uuid4())
  2324. index_node_hash = helper.generate_text_hash(args.content)
  2325. child_chunk = ChildChunk(
  2326. tenant_id=current_user.current_tenant_id,
  2327. dataset_id=dataset.id,
  2328. document_id=document.id,
  2329. segment_id=segment.id,
  2330. position=position,
  2331. index_node_id=index_node_id,
  2332. index_node_hash=index_node_hash,
  2333. content=args.content,
  2334. word_count=len(args.content),
  2335. type="customized",
  2336. created_by=current_user.id,
  2337. )
  2338. db.session.add(child_chunk)
  2339. db.session.flush()
  2340. new_child_chunks.append(child_chunk)
  2341. VectorService.update_child_chunk_vector(new_child_chunks, update_child_chunks, delete_child_chunks, dataset)
  2342. db.session.commit()
  2343. except Exception as e:
  2344. logging.exception("update child chunk index failed")
  2345. db.session.rollback()
  2346. raise ChildChunkIndexingError(str(e))
  2347. return sorted(new_child_chunks + update_child_chunks, key=lambda x: x.position)
  2348. @classmethod
  2349. def update_child_chunk(
  2350. cls,
  2351. content: str,
  2352. child_chunk: ChildChunk,
  2353. segment: DocumentSegment,
  2354. document: Document,
  2355. dataset: Dataset,
  2356. ) -> ChildChunk:
  2357. try:
  2358. child_chunk.content = content
  2359. child_chunk.word_count = len(content)
  2360. child_chunk.updated_by = current_user.id
  2361. child_chunk.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  2362. child_chunk.type = "customized"
  2363. db.session.add(child_chunk)
  2364. VectorService.update_child_chunk_vector([], [child_chunk], [], dataset)
  2365. db.session.commit()
  2366. except Exception as e:
  2367. logging.exception("update child chunk index failed")
  2368. db.session.rollback()
  2369. raise ChildChunkIndexingError(str(e))
  2370. return child_chunk
  2371. @classmethod
  2372. def delete_child_chunk(cls, child_chunk: ChildChunk, dataset: Dataset):
  2373. db.session.delete(child_chunk)
  2374. try:
  2375. VectorService.delete_child_chunk_vector(child_chunk, dataset)
  2376. except Exception as e:
  2377. logging.exception("delete child chunk index failed")
  2378. db.session.rollback()
  2379. raise ChildChunkDeleteIndexError(str(e))
  2380. db.session.commit()
  2381. @classmethod
  2382. def get_child_chunks(
  2383. cls, segment_id: str, document_id: str, dataset_id: str, page: int, limit: int, keyword: Optional[str] = None
  2384. ):
  2385. query = ChildChunk.query.filter_by(
  2386. tenant_id=current_user.current_tenant_id,
  2387. dataset_id=dataset_id,
  2388. document_id=document_id,
  2389. segment_id=segment_id,
  2390. ).order_by(ChildChunk.position.asc())
  2391. if keyword:
  2392. query = query.where(ChildChunk.content.ilike(f"%{keyword}%"))
  2393. return query.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
  2394. @classmethod
  2395. def get_child_chunk_by_id(cls, child_chunk_id: str, tenant_id: str) -> Optional[ChildChunk]:
  2396. """Get a child chunk by its ID."""
  2397. result = ChildChunk.query.filter(ChildChunk.id == child_chunk_id, ChildChunk.tenant_id == tenant_id).first()
  2398. return result if isinstance(result, ChildChunk) else None
  2399. @classmethod
  2400. def get_segments(
  2401. cls, document_id: str, tenant_id: str, status_list: list[str] | None = None, keyword: str | None = None
  2402. ):
  2403. """Get segments for a document with optional filtering."""
  2404. query = DocumentSegment.query.filter(
  2405. DocumentSegment.document_id == document_id, DocumentSegment.tenant_id == tenant_id
  2406. )
  2407. if status_list:
  2408. query = query.filter(DocumentSegment.status.in_(status_list))
  2409. if keyword:
  2410. query = query.filter(DocumentSegment.content.ilike(f"%{keyword}%"))
  2411. segments = query.order_by(DocumentSegment.position.asc()).all()
  2412. total = len(segments)
  2413. return segments, total
  2414. @classmethod
  2415. def update_segment_by_id(
  2416. cls, tenant_id: str, dataset_id: str, document_id: str, segment_id: str, segment_data: dict, user_id: str
  2417. ) -> tuple[DocumentSegment, Document]:
  2418. """Update a segment by its ID with validation and checks."""
  2419. # check dataset
  2420. dataset = db.session.query(Dataset).filter(Dataset.tenant_id == tenant_id, Dataset.id == dataset_id).first()
  2421. if not dataset:
  2422. raise NotFound("Dataset not found.")
  2423. # check user's model setting
  2424. DatasetService.check_dataset_model_setting(dataset)
  2425. # check document
  2426. document = DocumentService.get_document(dataset_id, document_id)
  2427. if not document:
  2428. raise NotFound("Document not found.")
  2429. # check embedding model setting if high quality
  2430. if dataset.indexing_technique == "high_quality":
  2431. try:
  2432. model_manager = ModelManager()
  2433. model_manager.get_model_instance(
  2434. tenant_id=user_id,
  2435. provider=dataset.embedding_model_provider,
  2436. model_type=ModelType.TEXT_EMBEDDING,
  2437. model=dataset.embedding_model,
  2438. )
  2439. except LLMBadRequestError:
  2440. raise ValueError(
  2441. "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider."
  2442. )
  2443. except ProviderTokenNotInitError as ex:
  2444. raise ValueError(ex.description)
  2445. # check segment
  2446. segment = DocumentSegment.query.filter(
  2447. DocumentSegment.id == segment_id, DocumentSegment.tenant_id == user_id
  2448. ).first()
  2449. if not segment:
  2450. raise NotFound("Segment not found.")
  2451. # validate and update segment
  2452. cls.segment_create_args_validate(segment_data, document)
  2453. updated_segment = cls.update_segment(SegmentUpdateArgs(**segment_data), segment, document, dataset)
  2454. return updated_segment, document
  2455. @classmethod
  2456. def get_segment_by_id(cls, segment_id: str, tenant_id: str) -> Optional[DocumentSegment]:
  2457. """Get a segment by its ID."""
  2458. result = DocumentSegment.query.filter(
  2459. DocumentSegment.id == segment_id, DocumentSegment.tenant_id == tenant_id
  2460. ).first()
  2461. return result if isinstance(result, DocumentSegment) else None
  2462. class DatasetCollectionBindingService:
  2463. @classmethod
  2464. def get_dataset_collection_binding(
  2465. cls, provider_name: str, model_name: str, collection_type: str = "dataset"
  2466. ) -> DatasetCollectionBinding:
  2467. dataset_collection_binding = (
  2468. db.session.query(DatasetCollectionBinding)
  2469. .filter(
  2470. DatasetCollectionBinding.provider_name == provider_name,
  2471. DatasetCollectionBinding.model_name == model_name,
  2472. DatasetCollectionBinding.type == collection_type,
  2473. )
  2474. .order_by(DatasetCollectionBinding.created_at)
  2475. .first()
  2476. )
  2477. if not dataset_collection_binding:
  2478. dataset_collection_binding = DatasetCollectionBinding(
  2479. provider_name=provider_name,
  2480. model_name=model_name,
  2481. collection_name=Dataset.gen_collection_name_by_id(str(uuid.uuid4())),
  2482. type=collection_type,
  2483. )
  2484. db.session.add(dataset_collection_binding)
  2485. db.session.commit()
  2486. return dataset_collection_binding
  2487. @classmethod
  2488. def get_dataset_collection_binding_by_id_and_type(
  2489. cls, collection_binding_id: str, collection_type: str = "dataset"
  2490. ) -> DatasetCollectionBinding:
  2491. dataset_collection_binding = (
  2492. db.session.query(DatasetCollectionBinding)
  2493. .filter(
  2494. DatasetCollectionBinding.id == collection_binding_id, DatasetCollectionBinding.type == collection_type
  2495. )
  2496. .order_by(DatasetCollectionBinding.created_at)
  2497. .first()
  2498. )
  2499. if not dataset_collection_binding:
  2500. raise ValueError("Dataset collection binding not found")
  2501. return dataset_collection_binding
  2502. class DatasetPermissionService:
  2503. @classmethod
  2504. def get_dataset_partial_member_list(cls, dataset_id):
  2505. user_list_query = (
  2506. db.session.query(
  2507. DatasetPermission.account_id,
  2508. )
  2509. .filter(DatasetPermission.dataset_id == dataset_id)
  2510. .all()
  2511. )
  2512. user_list = []
  2513. for user in user_list_query:
  2514. user_list.append(user.account_id)
  2515. return user_list
  2516. @classmethod
  2517. def update_partial_member_list(cls, tenant_id, dataset_id, user_list):
  2518. try:
  2519. db.session.query(DatasetPermission).filter(DatasetPermission.dataset_id == dataset_id).delete()
  2520. permissions = []
  2521. for user in user_list:
  2522. permission = DatasetPermission(
  2523. tenant_id=tenant_id,
  2524. dataset_id=dataset_id,
  2525. account_id=user["user_id"],
  2526. )
  2527. permissions.append(permission)
  2528. db.session.add_all(permissions)
  2529. db.session.commit()
  2530. except Exception as e:
  2531. db.session.rollback()
  2532. raise e
  2533. @classmethod
  2534. def check_permission(cls, user, dataset, requested_permission, requested_partial_member_list):
  2535. if not user.is_dataset_editor:
  2536. raise NoPermissionError("User does not have permission to edit this dataset.")
  2537. if user.is_dataset_operator and dataset.permission != requested_permission:
  2538. raise NoPermissionError("Dataset operators cannot change the dataset permissions.")
  2539. if user.is_dataset_operator and requested_permission == "partial_members":
  2540. if not requested_partial_member_list:
  2541. raise ValueError("Partial member list is required when setting to partial members.")
  2542. local_member_list = cls.get_dataset_partial_member_list(dataset.id)
  2543. request_member_list = [user["user_id"] for user in requested_partial_member_list]
  2544. if set(local_member_list) != set(request_member_list):
  2545. raise ValueError("Dataset operators cannot change the dataset permissions.")
  2546. @classmethod
  2547. def clear_partial_member_list(cls, dataset_id):
  2548. try:
  2549. db.session.query(DatasetPermission).filter(DatasetPermission.dataset_id == dataset_id).delete()
  2550. db.session.commit()
  2551. except Exception as e:
  2552. db.session.rollback()
  2553. raise e
  2554. @staticmethod
  2555. def save_template_with_dataset_id(
  2556. dataset: Dataset,
  2557. knowledge_config: KnowledgeConfig,
  2558. account: Account | Any,
  2559. dataset_process_rule: Optional[DatasetProcessRule] = None,
  2560. created_from: str = "web",
  2561. ):
  2562. return 1