dataset_service.py 126 KB

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