Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

dataset_service.py 96KB

2 лет назад
10 месяцев назад
10 месяцев назад
1 год назад
2 лет назад
10 месяцев назад
10 месяцев назад
2 лет назад
10 месяцев назад
10 месяцев назад
2 лет назад
10 месяцев назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
10 месяцев назад
2 лет назад
2 лет назад
10 месяцев назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
2 лет назад
2 лет назад
2 лет назад
10 месяцев назад
10 месяцев назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
10 месяцев назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
10 месяцев назад
10 месяцев назад
2 лет назад
2 лет назад
2 лет назад
10 месяцев назад
2 лет назад
10 месяцев назад
9 месяцев назад
9 месяцев назад
2 лет назад
9 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
2 лет назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
2 лет назад
2 лет назад
2 лет назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
2 лет назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
2 лет назад
2 лет назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
10 месяцев назад
2 лет назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
2 лет назад
10 месяцев назад
1 год назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
10 месяцев назад
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125
  1. import datetime
  2. import json
  3. import logging
  4. import random
  5. import time
  6. import uuid
  7. from typing import Any, Optional
  8. from flask_login import current_user # type: ignore
  9. from sqlalchemy import func
  10. from werkzeug.exceptions import NotFound
  11. from configs import dify_config
  12. from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
  13. from core.model_manager import ModelManager
  14. from core.model_runtime.entities.model_entities import ModelType
  15. from core.rag.index_processor.constant.index_type import IndexType
  16. from core.rag.retrieval.retrieval_methods import RetrievalMethod
  17. from events.dataset_event import dataset_was_deleted
  18. from events.document_event import document_was_deleted
  19. from extensions.ext_database import db
  20. from extensions.ext_redis import redis_client
  21. from libs import helper
  22. from models.account import Account, TenantAccountRole
  23. from models.dataset import (
  24. AppDatasetJoin,
  25. ChildChunk,
  26. Dataset,
  27. DatasetAutoDisableLog,
  28. DatasetCollectionBinding,
  29. DatasetPermission,
  30. DatasetPermissionEnum,
  31. DatasetProcessRule,
  32. DatasetQuery,
  33. Document,
  34. DocumentSegment,
  35. ExternalKnowledgeBindings,
  36. )
  37. from models.model import UploadFile
  38. from models.source import DataSourceOauthBinding
  39. from services.entities.knowledge_entities.knowledge_entities import (
  40. ChildChunkUpdateArgs,
  41. KnowledgeConfig,
  42. RerankingModel,
  43. RetrievalModel,
  44. SegmentUpdateArgs,
  45. )
  46. from services.errors.account import InvalidActionError, NoPermissionError
  47. from services.errors.chunk import ChildChunkDeleteIndexError, ChildChunkIndexingError
  48. from services.errors.dataset import DatasetNameDuplicateError
  49. from services.errors.document import DocumentIndexingError
  50. from services.errors.file import FileNotExistsError
  51. from services.external_knowledge_service import ExternalDatasetService
  52. from services.feature_service import FeatureModel, FeatureService
  53. from services.tag_service import TagService
  54. from services.vector_service import VectorService
  55. from tasks.batch_clean_document_task import batch_clean_document_task
  56. from tasks.clean_notion_document_task import clean_notion_document_task
  57. from tasks.deal_dataset_vector_index_task import deal_dataset_vector_index_task
  58. from tasks.delete_segment_from_index_task import delete_segment_from_index_task
  59. from tasks.disable_segment_from_index_task import disable_segment_from_index_task
  60. from tasks.disable_segments_from_index_task import disable_segments_from_index_task
  61. from tasks.document_indexing_task import document_indexing_task
  62. from tasks.document_indexing_update_task import document_indexing_update_task
  63. from tasks.duplicate_document_indexing_task import duplicate_document_indexing_task
  64. from tasks.enable_segments_to_index_task import enable_segments_to_index_task
  65. from tasks.recover_document_indexing_task import recover_document_indexing_task
  66. from tasks.retry_document_indexing_task import retry_document_indexing_task
  67. from tasks.sync_website_document_indexing_task import sync_website_document_indexing_task
  68. class DatasetService:
  69. @staticmethod
  70. def get_datasets(page, per_page, tenant_id=None, user=None, search=None, tag_ids=None, include_all=False):
  71. query = Dataset.query.filter(Dataset.tenant_id == tenant_id).order_by(Dataset.created_at.desc())
  72. if user:
  73. # get permitted dataset ids
  74. dataset_permission = DatasetPermission.query.filter_by(account_id=user.id, tenant_id=tenant_id).all()
  75. permitted_dataset_ids = {dp.dataset_id for dp in dataset_permission} if dataset_permission else None
  76. if user.current_role == TenantAccountRole.DATASET_OPERATOR:
  77. # only show datasets that the user has permission to access
  78. if permitted_dataset_ids:
  79. query = query.filter(Dataset.id.in_(permitted_dataset_ids))
  80. else:
  81. return [], 0
  82. else:
  83. if user.current_role != TenantAccountRole.OWNER or not include_all:
  84. # show all datasets that the user has permission to access
  85. if permitted_dataset_ids:
  86. query = query.filter(
  87. db.or_(
  88. Dataset.permission == DatasetPermissionEnum.ALL_TEAM,
  89. db.and_(
  90. Dataset.permission == DatasetPermissionEnum.ONLY_ME, Dataset.created_by == user.id
  91. ),
  92. db.and_(
  93. Dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM,
  94. Dataset.id.in_(permitted_dataset_ids),
  95. ),
  96. )
  97. )
  98. else:
  99. query = query.filter(
  100. db.or_(
  101. Dataset.permission == DatasetPermissionEnum.ALL_TEAM,
  102. db.and_(
  103. Dataset.permission == DatasetPermissionEnum.ONLY_ME, Dataset.created_by == user.id
  104. ),
  105. )
  106. )
  107. else:
  108. # if no user, only show datasets that are shared with all team members
  109. query = query.filter(Dataset.permission == DatasetPermissionEnum.ALL_TEAM)
  110. if search:
  111. query = query.filter(Dataset.name.ilike(f"%{search}%"))
  112. if tag_ids:
  113. target_ids = TagService.get_target_ids_by_tag_ids("knowledge", tenant_id, tag_ids)
  114. if target_ids:
  115. query = query.filter(Dataset.id.in_(target_ids))
  116. else:
  117. return [], 0
  118. datasets = query.paginate(page=page, per_page=per_page, max_per_page=100, error_out=False)
  119. return datasets.items, datasets.total
  120. @staticmethod
  121. def get_process_rules(dataset_id):
  122. # get the latest process rule
  123. dataset_process_rule = (
  124. db.session.query(DatasetProcessRule)
  125. .filter(DatasetProcessRule.dataset_id == dataset_id)
  126. .order_by(DatasetProcessRule.created_at.desc())
  127. .limit(1)
  128. .one_or_none()
  129. )
  130. if dataset_process_rule:
  131. mode = dataset_process_rule.mode
  132. rules = dataset_process_rule.rules_dict
  133. else:
  134. mode = DocumentService.DEFAULT_RULES["mode"]
  135. rules = DocumentService.DEFAULT_RULES["rules"]
  136. return {"mode": mode, "rules": rules}
  137. @staticmethod
  138. def get_datasets_by_ids(ids, tenant_id):
  139. datasets = Dataset.query.filter(Dataset.id.in_(ids), Dataset.tenant_id == tenant_id).paginate(
  140. page=1, per_page=len(ids), max_per_page=len(ids), error_out=False
  141. )
  142. return datasets.items, datasets.total
  143. @staticmethod
  144. def create_empty_dataset(
  145. tenant_id: str,
  146. name: str,
  147. description: Optional[str],
  148. indexing_technique: Optional[str],
  149. account: Account,
  150. permission: Optional[str] = None,
  151. provider: str = "vendor",
  152. external_knowledge_api_id: Optional[str] = None,
  153. external_knowledge_id: Optional[str] = None,
  154. ):
  155. # check if dataset name already exists
  156. if Dataset.query.filter_by(name=name, tenant_id=tenant_id).first():
  157. raise DatasetNameDuplicateError(f"Dataset with name {name} already exists.")
  158. embedding_model = None
  159. if indexing_technique == "high_quality":
  160. model_manager = ModelManager()
  161. embedding_model = model_manager.get_default_model_instance(
  162. tenant_id=tenant_id, model_type=ModelType.TEXT_EMBEDDING
  163. )
  164. dataset = Dataset(name=name, indexing_technique=indexing_technique)
  165. # dataset = Dataset(name=name, provider=provider, config=config)
  166. dataset.description = description
  167. dataset.created_by = account.id
  168. dataset.updated_by = account.id
  169. dataset.tenant_id = tenant_id
  170. dataset.embedding_model_provider = embedding_model.provider if embedding_model else None
  171. dataset.embedding_model = embedding_model.model if embedding_model else None
  172. dataset.permission = permission or DatasetPermissionEnum.ONLY_ME
  173. dataset.provider = provider
  174. db.session.add(dataset)
  175. db.session.flush()
  176. if provider == "external" and external_knowledge_api_id:
  177. external_knowledge_api = ExternalDatasetService.get_external_knowledge_api(external_knowledge_api_id)
  178. if not external_knowledge_api:
  179. raise ValueError("External API template not found.")
  180. external_knowledge_binding = ExternalKnowledgeBindings(
  181. tenant_id=tenant_id,
  182. dataset_id=dataset.id,
  183. external_knowledge_api_id=external_knowledge_api_id,
  184. external_knowledge_id=external_knowledge_id,
  185. created_by=account.id,
  186. )
  187. db.session.add(external_knowledge_binding)
  188. db.session.commit()
  189. return dataset
  190. @staticmethod
  191. def get_dataset(dataset_id) -> Optional[Dataset]:
  192. dataset: Optional[Dataset] = Dataset.query.filter_by(id=dataset_id).first()
  193. return dataset
  194. @staticmethod
  195. def check_dataset_model_setting(dataset):
  196. if dataset.indexing_technique == "high_quality":
  197. try:
  198. model_manager = ModelManager()
  199. model_manager.get_model_instance(
  200. tenant_id=dataset.tenant_id,
  201. provider=dataset.embedding_model_provider,
  202. model_type=ModelType.TEXT_EMBEDDING,
  203. model=dataset.embedding_model,
  204. )
  205. except LLMBadRequestError:
  206. raise ValueError(
  207. "No Embedding Model available. Please configure a valid provider "
  208. "in the Settings -> Model Provider."
  209. )
  210. except ProviderTokenNotInitError as ex:
  211. raise ValueError(f"The dataset in unavailable, due to: {ex.description}")
  212. @staticmethod
  213. def check_embedding_model_setting(tenant_id: str, embedding_model_provider: str, embedding_model: str):
  214. try:
  215. model_manager = ModelManager()
  216. model_manager.get_model_instance(
  217. tenant_id=tenant_id,
  218. provider=embedding_model_provider,
  219. model_type=ModelType.TEXT_EMBEDDING,
  220. model=embedding_model,
  221. )
  222. except LLMBadRequestError:
  223. raise ValueError(
  224. "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider."
  225. )
  226. except ProviderTokenNotInitError as ex:
  227. raise ValueError(f"The dataset in unavailable, due to: {ex.description}")
  228. @staticmethod
  229. def update_dataset(dataset_id, data, user):
  230. dataset = DatasetService.get_dataset(dataset_id)
  231. if not dataset:
  232. raise ValueError("Dataset not found")
  233. DatasetService.check_dataset_permission(dataset, user)
  234. if dataset.provider == "external":
  235. external_retrieval_model = data.get("external_retrieval_model", None)
  236. if external_retrieval_model:
  237. dataset.retrieval_model = external_retrieval_model
  238. dataset.name = data.get("name", dataset.name)
  239. dataset.description = data.get("description", "")
  240. permission = data.get("permission")
  241. if permission:
  242. dataset.permission = permission
  243. external_knowledge_id = data.get("external_knowledge_id", None)
  244. db.session.add(dataset)
  245. if not external_knowledge_id:
  246. raise ValueError("External knowledge id is required.")
  247. external_knowledge_api_id = data.get("external_knowledge_api_id", None)
  248. if not external_knowledge_api_id:
  249. raise ValueError("External knowledge api id is required.")
  250. external_knowledge_binding = ExternalKnowledgeBindings.query.filter_by(dataset_id=dataset_id).first()
  251. if (
  252. external_knowledge_binding.external_knowledge_id != external_knowledge_id
  253. or external_knowledge_binding.external_knowledge_api_id != external_knowledge_api_id
  254. ):
  255. external_knowledge_binding.external_knowledge_id = external_knowledge_id
  256. external_knowledge_binding.external_knowledge_api_id = external_knowledge_api_id
  257. db.session.add(external_knowledge_binding)
  258. db.session.commit()
  259. else:
  260. data.pop("partial_member_list", None)
  261. data.pop("external_knowledge_api_id", None)
  262. data.pop("external_knowledge_id", None)
  263. data.pop("external_retrieval_model", None)
  264. filtered_data = {k: v for k, v in data.items() if v is not None or k == "description"}
  265. action = None
  266. if dataset.indexing_technique != data["indexing_technique"]:
  267. # if update indexing_technique
  268. if data["indexing_technique"] == "economy":
  269. action = "remove"
  270. filtered_data["embedding_model"] = None
  271. filtered_data["embedding_model_provider"] = None
  272. filtered_data["collection_binding_id"] = None
  273. elif data["indexing_technique"] == "high_quality":
  274. action = "add"
  275. # get embedding model setting
  276. try:
  277. model_manager = ModelManager()
  278. embedding_model = model_manager.get_model_instance(
  279. tenant_id=current_user.current_tenant_id,
  280. provider=data["embedding_model_provider"],
  281. model_type=ModelType.TEXT_EMBEDDING,
  282. model=data["embedding_model"],
  283. )
  284. filtered_data["embedding_model"] = embedding_model.model
  285. filtered_data["embedding_model_provider"] = embedding_model.provider
  286. dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
  287. embedding_model.provider, embedding_model.model
  288. )
  289. filtered_data["collection_binding_id"] = dataset_collection_binding.id
  290. except LLMBadRequestError:
  291. raise ValueError(
  292. "No Embedding Model available. Please configure a valid provider "
  293. "in the Settings -> Model Provider."
  294. )
  295. except ProviderTokenNotInitError as ex:
  296. raise ValueError(ex.description)
  297. else:
  298. if (
  299. data["embedding_model_provider"] != dataset.embedding_model_provider
  300. or data["embedding_model"] != dataset.embedding_model
  301. ):
  302. action = "update"
  303. try:
  304. model_manager = ModelManager()
  305. embedding_model = model_manager.get_model_instance(
  306. tenant_id=current_user.current_tenant_id,
  307. provider=data["embedding_model_provider"],
  308. model_type=ModelType.TEXT_EMBEDDING,
  309. model=data["embedding_model"],
  310. )
  311. filtered_data["embedding_model"] = embedding_model.model
  312. filtered_data["embedding_model_provider"] = embedding_model.provider
  313. dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
  314. embedding_model.provider, embedding_model.model
  315. )
  316. filtered_data["collection_binding_id"] = dataset_collection_binding.id
  317. except LLMBadRequestError:
  318. raise ValueError(
  319. "No Embedding Model available. Please configure a valid provider "
  320. "in the Settings -> Model Provider."
  321. )
  322. except ProviderTokenNotInitError as ex:
  323. raise ValueError(ex.description)
  324. filtered_data["updated_by"] = user.id
  325. filtered_data["updated_at"] = datetime.datetime.now()
  326. # update Retrieval model
  327. filtered_data["retrieval_model"] = data["retrieval_model"]
  328. dataset.query.filter_by(id=dataset_id).update(filtered_data)
  329. db.session.commit()
  330. if action:
  331. deal_dataset_vector_index_task.delay(dataset_id, action)
  332. return dataset
  333. @staticmethod
  334. def delete_dataset(dataset_id, user):
  335. dataset = DatasetService.get_dataset(dataset_id)
  336. if dataset is None:
  337. return False
  338. DatasetService.check_dataset_permission(dataset, user)
  339. dataset_was_deleted.send(dataset)
  340. db.session.delete(dataset)
  341. db.session.commit()
  342. return True
  343. @staticmethod
  344. def dataset_use_check(dataset_id) -> bool:
  345. count = AppDatasetJoin.query.filter_by(dataset_id=dataset_id).count()
  346. if count > 0:
  347. return True
  348. return False
  349. @staticmethod
  350. def check_dataset_permission(dataset, user):
  351. if dataset.tenant_id != user.current_tenant_id:
  352. logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")
  353. raise NoPermissionError("You do not have permission to access this dataset.")
  354. if user.current_role != TenantAccountRole.OWNER:
  355. if dataset.permission == DatasetPermissionEnum.ONLY_ME and dataset.created_by != user.id:
  356. logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")
  357. raise NoPermissionError("You do not have permission to access this dataset.")
  358. if dataset.permission == "partial_members":
  359. user_permission = DatasetPermission.query.filter_by(dataset_id=dataset.id, account_id=user.id).first()
  360. if (
  361. not user_permission
  362. and dataset.tenant_id != user.current_tenant_id
  363. and dataset.created_by != user.id
  364. ):
  365. logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")
  366. raise NoPermissionError("You do not have permission to access this dataset.")
  367. @staticmethod
  368. def check_dataset_operator_permission(user: Optional[Account] = None, dataset: Optional[Dataset] = None):
  369. if not dataset:
  370. raise ValueError("Dataset not found")
  371. if not user:
  372. raise ValueError("User not found")
  373. if user.current_role != TenantAccountRole.OWNER:
  374. if dataset.permission == DatasetPermissionEnum.ONLY_ME:
  375. if dataset.created_by != user.id:
  376. raise NoPermissionError("You do not have permission to access this dataset.")
  377. elif dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM:
  378. if not any(
  379. dp.dataset_id == dataset.id for dp in DatasetPermission.query.filter_by(account_id=user.id).all()
  380. ):
  381. raise NoPermissionError("You do not have permission to access this dataset.")
  382. @staticmethod
  383. def get_dataset_queries(dataset_id: str, page: int, per_page: int):
  384. dataset_queries = (
  385. DatasetQuery.query.filter_by(dataset_id=dataset_id)
  386. .order_by(db.desc(DatasetQuery.created_at))
  387. .paginate(page=page, per_page=per_page, max_per_page=100, error_out=False)
  388. )
  389. return dataset_queries.items, dataset_queries.total
  390. @staticmethod
  391. def get_related_apps(dataset_id: str):
  392. return (
  393. AppDatasetJoin.query.filter(AppDatasetJoin.dataset_id == dataset_id)
  394. .order_by(db.desc(AppDatasetJoin.created_at))
  395. .all()
  396. )
  397. @staticmethod
  398. def get_dataset_auto_disable_logs(dataset_id: str) -> dict:
  399. features = FeatureService.get_features(current_user.current_tenant_id)
  400. if not features.billing.enabled or features.billing.subscription.plan == "sandbox":
  401. return {
  402. "document_ids": [],
  403. "count": 0,
  404. }
  405. # get recent 30 days auto disable logs
  406. start_date = datetime.datetime.now() - datetime.timedelta(days=30)
  407. dataset_auto_disable_logs = DatasetAutoDisableLog.query.filter(
  408. DatasetAutoDisableLog.dataset_id == dataset_id,
  409. DatasetAutoDisableLog.created_at >= start_date,
  410. ).all()
  411. if dataset_auto_disable_logs:
  412. return {
  413. "document_ids": [log.document_id for log in dataset_auto_disable_logs],
  414. "count": len(dataset_auto_disable_logs),
  415. }
  416. return {
  417. "document_ids": [],
  418. "count": 0,
  419. }
  420. class DocumentService:
  421. DEFAULT_RULES: dict[str, Any] = {
  422. "mode": "custom",
  423. "rules": {
  424. "pre_processing_rules": [
  425. {"id": "remove_extra_spaces", "enabled": True},
  426. {"id": "remove_urls_emails", "enabled": False},
  427. ],
  428. "segmentation": {"delimiter": "\n", "max_tokens": 500, "chunk_overlap": 50},
  429. },
  430. "limits": {
  431. "indexing_max_segmentation_tokens_length": dify_config.INDEXING_MAX_SEGMENTATION_TOKENS_LENGTH,
  432. },
  433. }
  434. DOCUMENT_METADATA_SCHEMA: dict[str, Any] = {
  435. "book": {
  436. "title": str,
  437. "language": str,
  438. "author": str,
  439. "publisher": str,
  440. "publication_date": str,
  441. "isbn": str,
  442. "category": str,
  443. },
  444. "web_page": {
  445. "title": str,
  446. "url": str,
  447. "language": str,
  448. "publish_date": str,
  449. "author/publisher": str,
  450. "topic/keywords": str,
  451. "description": str,
  452. },
  453. "paper": {
  454. "title": str,
  455. "language": str,
  456. "author": str,
  457. "publish_date": str,
  458. "journal/conference_name": str,
  459. "volume/issue/page_numbers": str,
  460. "doi": str,
  461. "topic/keywords": str,
  462. "abstract": str,
  463. },
  464. "social_media_post": {
  465. "platform": str,
  466. "author/username": str,
  467. "publish_date": str,
  468. "post_url": str,
  469. "topic/tags": str,
  470. },
  471. "wikipedia_entry": {
  472. "title": str,
  473. "language": str,
  474. "web_page_url": str,
  475. "last_edit_date": str,
  476. "editor/contributor": str,
  477. "summary/introduction": str,
  478. },
  479. "personal_document": {
  480. "title": str,
  481. "author": str,
  482. "creation_date": str,
  483. "last_modified_date": str,
  484. "document_type": str,
  485. "tags/category": str,
  486. },
  487. "business_document": {
  488. "title": str,
  489. "author": str,
  490. "creation_date": str,
  491. "last_modified_date": str,
  492. "document_type": str,
  493. "department/team": str,
  494. },
  495. "im_chat_log": {
  496. "chat_platform": str,
  497. "chat_participants/group_name": str,
  498. "start_date": str,
  499. "end_date": str,
  500. "summary": str,
  501. },
  502. "synced_from_notion": {
  503. "title": str,
  504. "language": str,
  505. "author/creator": str,
  506. "creation_date": str,
  507. "last_modified_date": str,
  508. "notion_page_link": str,
  509. "category/tags": str,
  510. "description": str,
  511. },
  512. "synced_from_github": {
  513. "repository_name": str,
  514. "repository_description": str,
  515. "repository_owner/organization": str,
  516. "code_filename": str,
  517. "code_file_path": str,
  518. "programming_language": str,
  519. "github_link": str,
  520. "open_source_license": str,
  521. "commit_date": str,
  522. "commit_author": str,
  523. },
  524. "others": dict,
  525. }
  526. @staticmethod
  527. def get_document(dataset_id: str, document_id: Optional[str] = None) -> Optional[Document]:
  528. if document_id:
  529. document = (
  530. db.session.query(Document).filter(Document.id == document_id, Document.dataset_id == dataset_id).first()
  531. )
  532. return document
  533. else:
  534. return None
  535. @staticmethod
  536. def get_document_by_id(document_id: str) -> Optional[Document]:
  537. document = db.session.query(Document).filter(Document.id == document_id).first()
  538. return document
  539. @staticmethod
  540. def get_document_by_dataset_id(dataset_id: str) -> list[Document]:
  541. documents = db.session.query(Document).filter(Document.dataset_id == dataset_id, Document.enabled == True).all()
  542. return documents
  543. @staticmethod
  544. def get_error_documents_by_dataset_id(dataset_id: str) -> list[Document]:
  545. documents = (
  546. db.session.query(Document)
  547. .filter(Document.dataset_id == dataset_id, Document.indexing_status.in_(["error", "paused"]))
  548. .all()
  549. )
  550. return documents
  551. @staticmethod
  552. def get_batch_documents(dataset_id: str, batch: str) -> list[Document]:
  553. documents = (
  554. db.session.query(Document)
  555. .filter(
  556. Document.batch == batch,
  557. Document.dataset_id == dataset_id,
  558. Document.tenant_id == current_user.current_tenant_id,
  559. )
  560. .all()
  561. )
  562. return documents
  563. @staticmethod
  564. def get_document_file_detail(file_id: str):
  565. file_detail = db.session.query(UploadFile).filter(UploadFile.id == file_id).one_or_none()
  566. return file_detail
  567. @staticmethod
  568. def check_archived(document):
  569. if document.archived:
  570. return True
  571. else:
  572. return False
  573. @staticmethod
  574. def delete_document(document):
  575. # trigger document_was_deleted signal
  576. file_id = None
  577. if document.data_source_type == "upload_file":
  578. if document.data_source_info:
  579. data_source_info = document.data_source_info_dict
  580. if data_source_info and "upload_file_id" in data_source_info:
  581. file_id = data_source_info["upload_file_id"]
  582. document_was_deleted.send(
  583. document.id, dataset_id=document.dataset_id, doc_form=document.doc_form, file_id=file_id
  584. )
  585. db.session.delete(document)
  586. db.session.commit()
  587. @staticmethod
  588. def delete_documents(dataset: Dataset, document_ids: list[str]):
  589. documents = db.session.query(Document).filter(Document.id.in_(document_ids)).all()
  590. file_ids = [
  591. document.data_source_info_dict["upload_file_id"]
  592. for document in documents
  593. if document.data_source_type == "upload_file"
  594. ]
  595. batch_clean_document_task.delay(document_ids, dataset.id, dataset.doc_form, file_ids)
  596. for document in documents:
  597. db.session.delete(document)
  598. db.session.commit()
  599. @staticmethod
  600. def rename_document(dataset_id: str, document_id: str, name: str) -> Document:
  601. dataset = DatasetService.get_dataset(dataset_id)
  602. if not dataset:
  603. raise ValueError("Dataset not found.")
  604. document = DocumentService.get_document(dataset_id, document_id)
  605. if not document:
  606. raise ValueError("Document not found.")
  607. if document.tenant_id != current_user.current_tenant_id:
  608. raise ValueError("No permission.")
  609. document.name = name
  610. db.session.add(document)
  611. db.session.commit()
  612. return document
  613. @staticmethod
  614. def pause_document(document):
  615. if document.indexing_status not in {"waiting", "parsing", "cleaning", "splitting", "indexing"}:
  616. raise DocumentIndexingError()
  617. # update document to be paused
  618. document.is_paused = True
  619. document.paused_by = current_user.id
  620. document.paused_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  621. db.session.add(document)
  622. db.session.commit()
  623. # set document paused flag
  624. indexing_cache_key = "document_{}_is_paused".format(document.id)
  625. redis_client.setnx(indexing_cache_key, "True")
  626. @staticmethod
  627. def recover_document(document):
  628. if not document.is_paused:
  629. raise DocumentIndexingError()
  630. # update document to be recover
  631. document.is_paused = False
  632. document.paused_by = None
  633. document.paused_at = None
  634. db.session.add(document)
  635. db.session.commit()
  636. # delete paused flag
  637. indexing_cache_key = "document_{}_is_paused".format(document.id)
  638. redis_client.delete(indexing_cache_key)
  639. # trigger async task
  640. recover_document_indexing_task.delay(document.dataset_id, document.id)
  641. @staticmethod
  642. def retry_document(dataset_id: str, documents: list[Document]):
  643. for document in documents:
  644. # add retry flag
  645. retry_indexing_cache_key = "document_{}_is_retried".format(document.id)
  646. cache_result = redis_client.get(retry_indexing_cache_key)
  647. if cache_result is not None:
  648. raise ValueError("Document is being retried, please try again later")
  649. # retry document indexing
  650. document.indexing_status = "waiting"
  651. db.session.add(document)
  652. db.session.commit()
  653. redis_client.setex(retry_indexing_cache_key, 600, 1)
  654. # trigger async task
  655. document_ids = [document.id for document in documents]
  656. retry_document_indexing_task.delay(dataset_id, document_ids)
  657. @staticmethod
  658. def sync_website_document(dataset_id: str, document: Document):
  659. # add sync flag
  660. sync_indexing_cache_key = "document_{}_is_sync".format(document.id)
  661. cache_result = redis_client.get(sync_indexing_cache_key)
  662. if cache_result is not None:
  663. raise ValueError("Document is being synced, please try again later")
  664. # sync document indexing
  665. document.indexing_status = "waiting"
  666. data_source_info = document.data_source_info_dict
  667. data_source_info["mode"] = "scrape"
  668. document.data_source_info = json.dumps(data_source_info, ensure_ascii=False)
  669. db.session.add(document)
  670. db.session.commit()
  671. redis_client.setex(sync_indexing_cache_key, 600, 1)
  672. sync_website_document_indexing_task.delay(dataset_id, document.id)
  673. @staticmethod
  674. def get_documents_position(dataset_id):
  675. document = Document.query.filter_by(dataset_id=dataset_id).order_by(Document.position.desc()).first()
  676. if document:
  677. return document.position + 1
  678. else:
  679. return 1
  680. @staticmethod
  681. def save_document_with_dataset_id(
  682. dataset: Dataset,
  683. knowledge_config: KnowledgeConfig,
  684. account: Account | Any,
  685. dataset_process_rule: Optional[DatasetProcessRule] = None,
  686. created_from: str = "web",
  687. ):
  688. # check document limit
  689. features = FeatureService.get_features(current_user.current_tenant_id)
  690. if features.billing.enabled:
  691. if not knowledge_config.original_document_id:
  692. count = 0
  693. if knowledge_config.data_source:
  694. if knowledge_config.data_source.info_list.data_source_type == "upload_file":
  695. upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids # type: ignore
  696. count = len(upload_file_list)
  697. elif knowledge_config.data_source.info_list.data_source_type == "notion_import":
  698. notion_info_list = knowledge_config.data_source.info_list.notion_info_list
  699. for notion_info in notion_info_list: # type: ignore
  700. count = count + len(notion_info.pages)
  701. elif knowledge_config.data_source.info_list.data_source_type == "website_crawl":
  702. website_info = knowledge_config.data_source.info_list.website_info_list
  703. count = len(website_info.urls) # type: ignore
  704. batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT)
  705. if count > batch_upload_limit:
  706. raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")
  707. DocumentService.check_documents_upload_quota(count, features)
  708. # if dataset is empty, update dataset data_source_type
  709. if not dataset.data_source_type:
  710. dataset.data_source_type = knowledge_config.data_source.info_list.data_source_type # type: ignore
  711. if not dataset.indexing_technique:
  712. if knowledge_config.indexing_technique not in Dataset.INDEXING_TECHNIQUE_LIST:
  713. raise ValueError("Indexing technique is invalid")
  714. dataset.indexing_technique = knowledge_config.indexing_technique
  715. if knowledge_config.indexing_technique == "high_quality":
  716. model_manager = ModelManager()
  717. if knowledge_config.embedding_model and knowledge_config.embedding_model_provider:
  718. dataset_embedding_model = knowledge_config.embedding_model
  719. dataset_embedding_model_provider = knowledge_config.embedding_model_provider
  720. else:
  721. embedding_model = model_manager.get_default_model_instance(
  722. tenant_id=current_user.current_tenant_id, model_type=ModelType.TEXT_EMBEDDING
  723. )
  724. dataset_embedding_model = embedding_model.model
  725. dataset_embedding_model_provider = embedding_model.provider
  726. dataset.embedding_model = dataset_embedding_model
  727. dataset.embedding_model_provider = dataset_embedding_model_provider
  728. dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
  729. dataset_embedding_model_provider, dataset_embedding_model
  730. )
  731. dataset.collection_binding_id = dataset_collection_binding.id
  732. if not dataset.retrieval_model:
  733. default_retrieval_model = {
  734. "search_method": RetrievalMethod.SEMANTIC_SEARCH.value,
  735. "reranking_enable": False,
  736. "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},
  737. "top_k": 2,
  738. "score_threshold_enabled": False,
  739. }
  740. dataset.retrieval_model = (
  741. knowledge_config.retrieval_model.model_dump()
  742. if knowledge_config.retrieval_model
  743. else default_retrieval_model
  744. ) # type: ignore
  745. documents = []
  746. if knowledge_config.original_document_id:
  747. document = DocumentService.update_document_with_dataset_id(dataset, knowledge_config, account)
  748. documents.append(document)
  749. batch = document.batch
  750. else:
  751. batch = time.strftime("%Y%m%d%H%M%S") + str(random.randint(100000, 999999))
  752. # save process rule
  753. if not dataset_process_rule:
  754. process_rule = knowledge_config.process_rule
  755. if process_rule:
  756. if process_rule.mode in ("custom", "hierarchical"):
  757. dataset_process_rule = DatasetProcessRule(
  758. dataset_id=dataset.id,
  759. mode=process_rule.mode,
  760. rules=process_rule.rules.model_dump_json() if process_rule.rules else None,
  761. created_by=account.id,
  762. )
  763. elif process_rule.mode == "automatic":
  764. dataset_process_rule = DatasetProcessRule(
  765. dataset_id=dataset.id,
  766. mode=process_rule.mode,
  767. rules=json.dumps(DatasetProcessRule.AUTOMATIC_RULES),
  768. created_by=account.id,
  769. )
  770. else:
  771. logging.warn(
  772. f"Invalid process rule mode: {process_rule.mode}, can not find dataset process rule"
  773. )
  774. return
  775. db.session.add(dataset_process_rule)
  776. db.session.commit()
  777. lock_name = "add_document_lock_dataset_id_{}".format(dataset.id)
  778. with redis_client.lock(lock_name, timeout=600):
  779. position = DocumentService.get_documents_position(dataset.id)
  780. document_ids = []
  781. duplicate_document_ids = []
  782. if knowledge_config.data_source.info_list.data_source_type == "upload_file": # type: ignore
  783. upload_file_list = knowledge_config.data_source.info_list.file_info_list.file_ids # type: ignore
  784. for file_id in upload_file_list:
  785. file = (
  786. db.session.query(UploadFile)
  787. .filter(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id)
  788. .first()
  789. )
  790. # raise error if file not found
  791. if not file:
  792. raise FileNotExistsError()
  793. file_name = file.name
  794. data_source_info = {
  795. "upload_file_id": file_id,
  796. }
  797. # check duplicate
  798. if knowledge_config.duplicate:
  799. document = Document.query.filter_by(
  800. dataset_id=dataset.id,
  801. tenant_id=current_user.current_tenant_id,
  802. data_source_type="upload_file",
  803. enabled=True,
  804. name=file_name,
  805. ).first()
  806. if document:
  807. document.dataset_process_rule_id = dataset_process_rule.id # type: ignore
  808. document.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
  809. document.created_from = created_from
  810. document.doc_form = knowledge_config.doc_form
  811. document.doc_language = knowledge_config.doc_language
  812. document.data_source_info = json.dumps(data_source_info)
  813. document.batch = batch
  814. document.indexing_status = "waiting"
  815. db.session.add(document)
  816. documents.append(document)
  817. duplicate_document_ids.append(document.id)
  818. continue
  819. document = DocumentService.build_document(
  820. dataset,
  821. dataset_process_rule.id, # type: ignore
  822. knowledge_config.data_source.info_list.data_source_type, # type: ignore
  823. knowledge_config.doc_form,
  824. knowledge_config.doc_language,
  825. data_source_info,
  826. created_from,
  827. position,
  828. account,
  829. file_name,
  830. batch,
  831. )
  832. db.session.add(document)
  833. db.session.flush()
  834. document_ids.append(document.id)
  835. documents.append(document)
  836. position += 1
  837. elif knowledge_config.data_source.info_list.data_source_type == "notion_import": # type: ignore
  838. notion_info_list = knowledge_config.data_source.info_list.notion_info_list # type: ignore
  839. if not notion_info_list:
  840. raise ValueError("No notion info list found.")
  841. exist_page_ids = []
  842. exist_document = {}
  843. documents = Document.query.filter_by(
  844. dataset_id=dataset.id,
  845. tenant_id=current_user.current_tenant_id,
  846. data_source_type="notion_import",
  847. enabled=True,
  848. ).all()
  849. if documents:
  850. for document in documents:
  851. data_source_info = json.loads(document.data_source_info)
  852. exist_page_ids.append(data_source_info["notion_page_id"])
  853. exist_document[data_source_info["notion_page_id"]] = document.id
  854. for notion_info in notion_info_list:
  855. workspace_id = notion_info.workspace_id
  856. data_source_binding = DataSourceOauthBinding.query.filter(
  857. db.and_(
  858. DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
  859. DataSourceOauthBinding.provider == "notion",
  860. DataSourceOauthBinding.disabled == False,
  861. DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
  862. )
  863. ).first()
  864. if not data_source_binding:
  865. raise ValueError("Data source binding not found.")
  866. for page in notion_info.pages:
  867. if page.page_id not in exist_page_ids:
  868. data_source_info = {
  869. "notion_workspace_id": workspace_id,
  870. "notion_page_id": page.page_id,
  871. "notion_page_icon": page.page_icon.model_dump() if page.page_icon else None,
  872. "type": page.type,
  873. }
  874. document = DocumentService.build_document(
  875. dataset,
  876. dataset_process_rule.id, # type: ignore
  877. knowledge_config.data_source.info_list.data_source_type, # type: ignore
  878. knowledge_config.doc_form,
  879. knowledge_config.doc_language,
  880. data_source_info,
  881. created_from,
  882. position,
  883. account,
  884. page.page_name,
  885. batch,
  886. )
  887. db.session.add(document)
  888. db.session.flush()
  889. document_ids.append(document.id)
  890. documents.append(document)
  891. position += 1
  892. else:
  893. exist_document.pop(page.page_id)
  894. # delete not selected documents
  895. if len(exist_document) > 0:
  896. clean_notion_document_task.delay(list(exist_document.values()), dataset.id)
  897. elif knowledge_config.data_source.info_list.data_source_type == "website_crawl": # type: ignore
  898. website_info = knowledge_config.data_source.info_list.website_info_list # type: ignore
  899. if not website_info:
  900. raise ValueError("No website info list found.")
  901. urls = website_info.urls
  902. for url in urls:
  903. data_source_info = {
  904. "url": url,
  905. "provider": website_info.provider,
  906. "job_id": website_info.job_id,
  907. "only_main_content": website_info.only_main_content,
  908. "mode": "crawl",
  909. }
  910. if len(url) > 255:
  911. document_name = url[:200] + "..."
  912. else:
  913. document_name = url
  914. document = DocumentService.build_document(
  915. dataset,
  916. dataset_process_rule.id, # type: ignore
  917. knowledge_config.data_source.info_list.data_source_type, # type: ignore
  918. knowledge_config.doc_form,
  919. knowledge_config.doc_language,
  920. data_source_info,
  921. created_from,
  922. position,
  923. account,
  924. document_name,
  925. batch,
  926. )
  927. db.session.add(document)
  928. db.session.flush()
  929. document_ids.append(document.id)
  930. documents.append(document)
  931. position += 1
  932. db.session.commit()
  933. # trigger async task
  934. if document_ids:
  935. document_indexing_task.delay(dataset.id, document_ids)
  936. if duplicate_document_ids:
  937. duplicate_document_indexing_task.delay(dataset.id, duplicate_document_ids)
  938. return documents, batch
  939. @staticmethod
  940. def check_documents_upload_quota(count: int, features: FeatureModel):
  941. can_upload_size = features.documents_upload_quota.limit - features.documents_upload_quota.size
  942. if count > can_upload_size:
  943. raise ValueError(
  944. f"You have reached the limit of your subscription. Only {can_upload_size} documents can be uploaded."
  945. )
  946. @staticmethod
  947. def build_document(
  948. dataset: Dataset,
  949. process_rule_id: str,
  950. data_source_type: str,
  951. document_form: str,
  952. document_language: str,
  953. data_source_info: dict,
  954. created_from: str,
  955. position: int,
  956. account: Account,
  957. name: str,
  958. batch: str,
  959. ):
  960. document = Document(
  961. tenant_id=dataset.tenant_id,
  962. dataset_id=dataset.id,
  963. position=position,
  964. data_source_type=data_source_type,
  965. data_source_info=json.dumps(data_source_info),
  966. dataset_process_rule_id=process_rule_id,
  967. batch=batch,
  968. name=name,
  969. created_from=created_from,
  970. created_by=account.id,
  971. doc_form=document_form,
  972. doc_language=document_language,
  973. )
  974. return document
  975. @staticmethod
  976. def get_tenant_documents_count():
  977. documents_count = Document.query.filter(
  978. Document.completed_at.isnot(None),
  979. Document.enabled == True,
  980. Document.archived == False,
  981. Document.tenant_id == current_user.current_tenant_id,
  982. ).count()
  983. return documents_count
  984. @staticmethod
  985. def update_document_with_dataset_id(
  986. dataset: Dataset,
  987. document_data: KnowledgeConfig,
  988. account: Account,
  989. dataset_process_rule: Optional[DatasetProcessRule] = None,
  990. created_from: str = "web",
  991. ):
  992. DatasetService.check_dataset_model_setting(dataset)
  993. document = DocumentService.get_document(dataset.id, document_data.original_document_id)
  994. if document is None:
  995. raise NotFound("Document not found")
  996. if document.display_status != "available":
  997. raise ValueError("Document is not available")
  998. # save process rule
  999. if document_data.process_rule:
  1000. process_rule = document_data.process_rule
  1001. if process_rule.mode in {"custom", "hierarchical"}:
  1002. dataset_process_rule = DatasetProcessRule(
  1003. dataset_id=dataset.id,
  1004. mode=process_rule.mode,
  1005. rules=process_rule.rules.model_dump_json() if process_rule.rules else None,
  1006. created_by=account.id,
  1007. )
  1008. elif process_rule.mode == "automatic":
  1009. dataset_process_rule = DatasetProcessRule(
  1010. dataset_id=dataset.id,
  1011. mode=process_rule.mode,
  1012. rules=json.dumps(DatasetProcessRule.AUTOMATIC_RULES),
  1013. created_by=account.id,
  1014. )
  1015. if dataset_process_rule is not None:
  1016. db.session.add(dataset_process_rule)
  1017. db.session.commit()
  1018. document.dataset_process_rule_id = dataset_process_rule.id
  1019. # update document data source
  1020. if document_data.data_source:
  1021. file_name = ""
  1022. data_source_info = {}
  1023. if document_data.data_source.info_list.data_source_type == "upload_file":
  1024. if not document_data.data_source.info_list.file_info_list:
  1025. raise ValueError("No file info list found.")
  1026. upload_file_list = document_data.data_source.info_list.file_info_list.file_ids
  1027. for file_id in upload_file_list:
  1028. file = (
  1029. db.session.query(UploadFile)
  1030. .filter(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id)
  1031. .first()
  1032. )
  1033. # raise error if file not found
  1034. if not file:
  1035. raise FileNotExistsError()
  1036. file_name = file.name
  1037. data_source_info = {
  1038. "upload_file_id": file_id,
  1039. }
  1040. elif document_data.data_source.info_list.data_source_type == "notion_import":
  1041. if not document_data.data_source.info_list.notion_info_list:
  1042. raise ValueError("No notion info list found.")
  1043. notion_info_list = document_data.data_source.info_list.notion_info_list
  1044. for notion_info in notion_info_list:
  1045. workspace_id = notion_info.workspace_id
  1046. data_source_binding = DataSourceOauthBinding.query.filter(
  1047. db.and_(
  1048. DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
  1049. DataSourceOauthBinding.provider == "notion",
  1050. DataSourceOauthBinding.disabled == False,
  1051. DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',
  1052. )
  1053. ).first()
  1054. if not data_source_binding:
  1055. raise ValueError("Data source binding not found.")
  1056. for page in notion_info.pages:
  1057. data_source_info = {
  1058. "notion_workspace_id": workspace_id,
  1059. "notion_page_id": page.page_id,
  1060. "notion_page_icon": page.page_icon.model_dump() if page.page_icon else None, # type: ignore
  1061. "type": page.type,
  1062. }
  1063. elif document_data.data_source.info_list.data_source_type == "website_crawl":
  1064. website_info = document_data.data_source.info_list.website_info_list
  1065. if website_info:
  1066. urls = website_info.urls
  1067. for url in urls:
  1068. data_source_info = {
  1069. "url": url,
  1070. "provider": website_info.provider,
  1071. "job_id": website_info.job_id,
  1072. "only_main_content": website_info.only_main_content, # type: ignore
  1073. "mode": "crawl",
  1074. }
  1075. document.data_source_type = document_data.data_source.info_list.data_source_type
  1076. document.data_source_info = json.dumps(data_source_info)
  1077. document.name = file_name
  1078. # update document name
  1079. if document_data.name:
  1080. document.name = document_data.name
  1081. # update document to be waiting
  1082. document.indexing_status = "waiting"
  1083. document.completed_at = None
  1084. document.processing_started_at = None
  1085. document.parsing_completed_at = None
  1086. document.cleaning_completed_at = None
  1087. document.splitting_completed_at = None
  1088. document.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1089. document.created_from = created_from
  1090. document.doc_form = document_data.doc_form
  1091. db.session.add(document)
  1092. db.session.commit()
  1093. # update document segment
  1094. update_params = {DocumentSegment.status: "re_segment"}
  1095. DocumentSegment.query.filter_by(document_id=document.id).update(update_params)
  1096. db.session.commit()
  1097. # trigger async task
  1098. document_indexing_update_task.delay(document.dataset_id, document.id)
  1099. return document
  1100. @staticmethod
  1101. def save_document_without_dataset_id(tenant_id: str, knowledge_config: KnowledgeConfig, account: Account):
  1102. features = FeatureService.get_features(current_user.current_tenant_id)
  1103. if features.billing.enabled:
  1104. count = 0
  1105. if knowledge_config.data_source.info_list.data_source_type == "upload_file": # type: ignore
  1106. upload_file_list = (
  1107. knowledge_config.data_source.info_list.file_info_list.file_ids # type: ignore
  1108. if knowledge_config.data_source.info_list.file_info_list # type: ignore
  1109. else []
  1110. )
  1111. count = len(upload_file_list)
  1112. elif knowledge_config.data_source.info_list.data_source_type == "notion_import": # type: ignore
  1113. notion_info_list = knowledge_config.data_source.info_list.notion_info_list # type: ignore
  1114. if notion_info_list:
  1115. for notion_info in notion_info_list:
  1116. count = count + len(notion_info.pages)
  1117. elif knowledge_config.data_source.info_list.data_source_type == "website_crawl": # type: ignore
  1118. website_info = knowledge_config.data_source.info_list.website_info_list # type: ignore
  1119. if website_info:
  1120. count = len(website_info.urls)
  1121. batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT)
  1122. if count > batch_upload_limit:
  1123. raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")
  1124. DocumentService.check_documents_upload_quota(count, features)
  1125. dataset_collection_binding_id = None
  1126. retrieval_model = None
  1127. if knowledge_config.indexing_technique == "high_quality":
  1128. dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
  1129. knowledge_config.embedding_model_provider, # type: ignore
  1130. knowledge_config.embedding_model, # type: ignore
  1131. )
  1132. dataset_collection_binding_id = dataset_collection_binding.id
  1133. if knowledge_config.retrieval_model:
  1134. retrieval_model = knowledge_config.retrieval_model
  1135. else:
  1136. retrieval_model = RetrievalModel(
  1137. search_method=RetrievalMethod.SEMANTIC_SEARCH.value,
  1138. reranking_enable=False,
  1139. reranking_model=RerankingModel(reranking_provider_name="", reranking_model_name=""),
  1140. top_k=2,
  1141. score_threshold_enabled=False,
  1142. )
  1143. # save dataset
  1144. dataset = Dataset(
  1145. tenant_id=tenant_id,
  1146. name="",
  1147. data_source_type=knowledge_config.data_source.info_list.data_source_type, # type: ignore
  1148. indexing_technique=knowledge_config.indexing_technique,
  1149. created_by=account.id,
  1150. embedding_model=knowledge_config.embedding_model,
  1151. embedding_model_provider=knowledge_config.embedding_model_provider,
  1152. collection_binding_id=dataset_collection_binding_id,
  1153. retrieval_model=retrieval_model.model_dump() if retrieval_model else None,
  1154. )
  1155. db.session.add(dataset) # type: ignore
  1156. db.session.flush()
  1157. documents, batch = DocumentService.save_document_with_dataset_id(dataset, knowledge_config, account)
  1158. cut_length = 18
  1159. cut_name = documents[0].name[:cut_length]
  1160. dataset.name = cut_name + "..."
  1161. dataset.description = "useful for when you want to answer queries about the " + documents[0].name
  1162. db.session.commit()
  1163. return dataset, documents, batch
  1164. @classmethod
  1165. def document_create_args_validate(cls, knowledge_config: KnowledgeConfig):
  1166. if not knowledge_config.data_source and not knowledge_config.process_rule:
  1167. raise ValueError("Data source or Process rule is required")
  1168. else:
  1169. if knowledge_config.data_source:
  1170. DocumentService.data_source_args_validate(knowledge_config)
  1171. if knowledge_config.process_rule:
  1172. DocumentService.process_rule_args_validate(knowledge_config)
  1173. @classmethod
  1174. def data_source_args_validate(cls, knowledge_config: KnowledgeConfig):
  1175. if not knowledge_config.data_source:
  1176. raise ValueError("Data source is required")
  1177. if knowledge_config.data_source.info_list.data_source_type not in Document.DATA_SOURCES:
  1178. raise ValueError("Data source type is invalid")
  1179. if not knowledge_config.data_source.info_list:
  1180. raise ValueError("Data source info is required")
  1181. if knowledge_config.data_source.info_list.data_source_type == "upload_file":
  1182. if not knowledge_config.data_source.info_list.file_info_list:
  1183. raise ValueError("File source info is required")
  1184. if knowledge_config.data_source.info_list.data_source_type == "notion_import":
  1185. if not knowledge_config.data_source.info_list.notion_info_list:
  1186. raise ValueError("Notion source info is required")
  1187. if knowledge_config.data_source.info_list.data_source_type == "website_crawl":
  1188. if not knowledge_config.data_source.info_list.website_info_list:
  1189. raise ValueError("Website source info is required")
  1190. @classmethod
  1191. def process_rule_args_validate(cls, knowledge_config: KnowledgeConfig):
  1192. if not knowledge_config.process_rule:
  1193. raise ValueError("Process rule is required")
  1194. if not knowledge_config.process_rule.mode:
  1195. raise ValueError("Process rule mode is required")
  1196. if knowledge_config.process_rule.mode not in DatasetProcessRule.MODES:
  1197. raise ValueError("Process rule mode is invalid")
  1198. if knowledge_config.process_rule.mode == "automatic":
  1199. knowledge_config.process_rule.rules = None
  1200. else:
  1201. if not knowledge_config.process_rule.rules:
  1202. raise ValueError("Process rule rules is required")
  1203. if knowledge_config.process_rule.rules.pre_processing_rules is None:
  1204. raise ValueError("Process rule pre_processing_rules is required")
  1205. unique_pre_processing_rule_dicts = {}
  1206. for pre_processing_rule in knowledge_config.process_rule.rules.pre_processing_rules:
  1207. if not pre_processing_rule.id:
  1208. raise ValueError("Process rule pre_processing_rules id is required")
  1209. if not isinstance(pre_processing_rule.enabled, bool):
  1210. raise ValueError("Process rule pre_processing_rules enabled is invalid")
  1211. unique_pre_processing_rule_dicts[pre_processing_rule.id] = pre_processing_rule
  1212. knowledge_config.process_rule.rules.pre_processing_rules = list(unique_pre_processing_rule_dicts.values())
  1213. if not knowledge_config.process_rule.rules.segmentation:
  1214. raise ValueError("Process rule segmentation is required")
  1215. if not knowledge_config.process_rule.rules.segmentation.separator:
  1216. raise ValueError("Process rule segmentation separator is required")
  1217. if not isinstance(knowledge_config.process_rule.rules.segmentation.separator, str):
  1218. raise ValueError("Process rule segmentation separator is invalid")
  1219. if not (
  1220. knowledge_config.process_rule.mode == "hierarchical"
  1221. and knowledge_config.process_rule.rules.parent_mode == "full-doc"
  1222. ):
  1223. if not knowledge_config.process_rule.rules.segmentation.max_tokens:
  1224. raise ValueError("Process rule segmentation max_tokens is required")
  1225. if not isinstance(knowledge_config.process_rule.rules.segmentation.max_tokens, int):
  1226. raise ValueError("Process rule segmentation max_tokens is invalid")
  1227. @classmethod
  1228. def estimate_args_validate(cls, args: dict):
  1229. if "info_list" not in args or not args["info_list"]:
  1230. raise ValueError("Data source info is required")
  1231. if not isinstance(args["info_list"], dict):
  1232. raise ValueError("Data info is invalid")
  1233. if "process_rule" not in args or not args["process_rule"]:
  1234. raise ValueError("Process rule is required")
  1235. if not isinstance(args["process_rule"], dict):
  1236. raise ValueError("Process rule is invalid")
  1237. if "mode" not in args["process_rule"] or not args["process_rule"]["mode"]:
  1238. raise ValueError("Process rule mode is required")
  1239. if args["process_rule"]["mode"] not in DatasetProcessRule.MODES:
  1240. raise ValueError("Process rule mode is invalid")
  1241. if args["process_rule"]["mode"] == "automatic":
  1242. args["process_rule"]["rules"] = {}
  1243. else:
  1244. if "rules" not in args["process_rule"] or not args["process_rule"]["rules"]:
  1245. raise ValueError("Process rule rules is required")
  1246. if not isinstance(args["process_rule"]["rules"], dict):
  1247. raise ValueError("Process rule rules is invalid")
  1248. if (
  1249. "pre_processing_rules" not in args["process_rule"]["rules"]
  1250. or args["process_rule"]["rules"]["pre_processing_rules"] is None
  1251. ):
  1252. raise ValueError("Process rule pre_processing_rules is required")
  1253. if not isinstance(args["process_rule"]["rules"]["pre_processing_rules"], list):
  1254. raise ValueError("Process rule pre_processing_rules is invalid")
  1255. unique_pre_processing_rule_dicts = {}
  1256. for pre_processing_rule in args["process_rule"]["rules"]["pre_processing_rules"]:
  1257. if "id" not in pre_processing_rule or not pre_processing_rule["id"]:
  1258. raise ValueError("Process rule pre_processing_rules id is required")
  1259. if pre_processing_rule["id"] not in DatasetProcessRule.PRE_PROCESSING_RULES:
  1260. raise ValueError("Process rule pre_processing_rules id is invalid")
  1261. if "enabled" not in pre_processing_rule or pre_processing_rule["enabled"] is None:
  1262. raise ValueError("Process rule pre_processing_rules enabled is required")
  1263. if not isinstance(pre_processing_rule["enabled"], bool):
  1264. raise ValueError("Process rule pre_processing_rules enabled is invalid")
  1265. unique_pre_processing_rule_dicts[pre_processing_rule["id"]] = pre_processing_rule
  1266. args["process_rule"]["rules"]["pre_processing_rules"] = list(unique_pre_processing_rule_dicts.values())
  1267. if (
  1268. "segmentation" not in args["process_rule"]["rules"]
  1269. or args["process_rule"]["rules"]["segmentation"] is None
  1270. ):
  1271. raise ValueError("Process rule segmentation is required")
  1272. if not isinstance(args["process_rule"]["rules"]["segmentation"], dict):
  1273. raise ValueError("Process rule segmentation is invalid")
  1274. if (
  1275. "separator" not in args["process_rule"]["rules"]["segmentation"]
  1276. or not args["process_rule"]["rules"]["segmentation"]["separator"]
  1277. ):
  1278. raise ValueError("Process rule segmentation separator is required")
  1279. if not isinstance(args["process_rule"]["rules"]["segmentation"]["separator"], str):
  1280. raise ValueError("Process rule segmentation separator is invalid")
  1281. if (
  1282. "max_tokens" not in args["process_rule"]["rules"]["segmentation"]
  1283. or not args["process_rule"]["rules"]["segmentation"]["max_tokens"]
  1284. ):
  1285. raise ValueError("Process rule segmentation max_tokens is required")
  1286. if not isinstance(args["process_rule"]["rules"]["segmentation"]["max_tokens"], int):
  1287. raise ValueError("Process rule segmentation max_tokens is invalid")
  1288. class SegmentService:
  1289. @classmethod
  1290. def segment_create_args_validate(cls, args: dict, document: Document):
  1291. if document.doc_form == "qa_model":
  1292. if "answer" not in args or not args["answer"]:
  1293. raise ValueError("Answer is required")
  1294. if not args["answer"].strip():
  1295. raise ValueError("Answer is empty")
  1296. if "content" not in args or not args["content"] or not args["content"].strip():
  1297. raise ValueError("Content is empty")
  1298. @classmethod
  1299. def create_segment(cls, args: dict, document: Document, dataset: Dataset):
  1300. content = args["content"]
  1301. doc_id = str(uuid.uuid4())
  1302. segment_hash = helper.generate_text_hash(content)
  1303. tokens = 0
  1304. if dataset.indexing_technique == "high_quality":
  1305. model_manager = ModelManager()
  1306. embedding_model = model_manager.get_model_instance(
  1307. tenant_id=current_user.current_tenant_id,
  1308. provider=dataset.embedding_model_provider,
  1309. model_type=ModelType.TEXT_EMBEDDING,
  1310. model=dataset.embedding_model,
  1311. )
  1312. # calc embedding use tokens
  1313. tokens = embedding_model.get_text_embedding_num_tokens(texts=[content])
  1314. lock_name = "add_segment_lock_document_id_{}".format(document.id)
  1315. with redis_client.lock(lock_name, timeout=600):
  1316. max_position = (
  1317. db.session.query(func.max(DocumentSegment.position))
  1318. .filter(DocumentSegment.document_id == document.id)
  1319. .scalar()
  1320. )
  1321. segment_document = DocumentSegment(
  1322. tenant_id=current_user.current_tenant_id,
  1323. dataset_id=document.dataset_id,
  1324. document_id=document.id,
  1325. index_node_id=doc_id,
  1326. index_node_hash=segment_hash,
  1327. position=max_position + 1 if max_position else 1,
  1328. content=content,
  1329. word_count=len(content),
  1330. tokens=tokens,
  1331. status="completed",
  1332. indexing_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
  1333. completed_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
  1334. created_by=current_user.id,
  1335. )
  1336. if document.doc_form == "qa_model":
  1337. segment_document.word_count += len(args["answer"])
  1338. segment_document.answer = args["answer"]
  1339. db.session.add(segment_document)
  1340. # update document word count
  1341. document.word_count += segment_document.word_count
  1342. db.session.add(document)
  1343. db.session.commit()
  1344. # save vector index
  1345. try:
  1346. VectorService.create_segments_vector([args["keywords"]], [segment_document], dataset, document.doc_form)
  1347. except Exception as e:
  1348. logging.exception("create segment index failed")
  1349. segment_document.enabled = False
  1350. segment_document.disabled_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1351. segment_document.status = "error"
  1352. segment_document.error = str(e)
  1353. db.session.commit()
  1354. segment = db.session.query(DocumentSegment).filter(DocumentSegment.id == segment_document.id).first()
  1355. return segment
  1356. @classmethod
  1357. def multi_create_segment(cls, segments: list, document: Document, dataset: Dataset):
  1358. lock_name = "multi_add_segment_lock_document_id_{}".format(document.id)
  1359. increment_word_count = 0
  1360. with redis_client.lock(lock_name, timeout=600):
  1361. embedding_model = None
  1362. if dataset.indexing_technique == "high_quality":
  1363. model_manager = ModelManager()
  1364. embedding_model = model_manager.get_model_instance(
  1365. tenant_id=current_user.current_tenant_id,
  1366. provider=dataset.embedding_model_provider,
  1367. model_type=ModelType.TEXT_EMBEDDING,
  1368. model=dataset.embedding_model,
  1369. )
  1370. max_position = (
  1371. db.session.query(func.max(DocumentSegment.position))
  1372. .filter(DocumentSegment.document_id == document.id)
  1373. .scalar()
  1374. )
  1375. pre_segment_data_list = []
  1376. segment_data_list = []
  1377. keywords_list = []
  1378. position = max_position + 1 if max_position else 1
  1379. for segment_item in segments:
  1380. content = segment_item["content"]
  1381. doc_id = str(uuid.uuid4())
  1382. segment_hash = helper.generate_text_hash(content)
  1383. tokens = 0
  1384. if dataset.indexing_technique == "high_quality" and embedding_model:
  1385. # calc embedding use tokens
  1386. if document.doc_form == "qa_model":
  1387. tokens = embedding_model.get_text_embedding_num_tokens(texts=[content + segment_item["answer"]])
  1388. else:
  1389. tokens = embedding_model.get_text_embedding_num_tokens(texts=[content])
  1390. segment_document = DocumentSegment(
  1391. tenant_id=current_user.current_tenant_id,
  1392. dataset_id=document.dataset_id,
  1393. document_id=document.id,
  1394. index_node_id=doc_id,
  1395. index_node_hash=segment_hash,
  1396. position=position,
  1397. content=content,
  1398. word_count=len(content),
  1399. tokens=tokens,
  1400. status="completed",
  1401. indexing_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
  1402. completed_at=datetime.datetime.now(datetime.UTC).replace(tzinfo=None),
  1403. created_by=current_user.id,
  1404. )
  1405. if document.doc_form == "qa_model":
  1406. segment_document.answer = segment_item["answer"]
  1407. segment_document.word_count += len(segment_item["answer"])
  1408. increment_word_count += segment_document.word_count
  1409. db.session.add(segment_document)
  1410. segment_data_list.append(segment_document)
  1411. position += 1
  1412. pre_segment_data_list.append(segment_document)
  1413. if "keywords" in segment_item:
  1414. keywords_list.append(segment_item["keywords"])
  1415. else:
  1416. keywords_list.append(None)
  1417. # update document word count
  1418. document.word_count += increment_word_count
  1419. db.session.add(document)
  1420. try:
  1421. # save vector index
  1422. VectorService.create_segments_vector(keywords_list, pre_segment_data_list, dataset, document.doc_form)
  1423. except Exception as e:
  1424. logging.exception("create segment index failed")
  1425. for segment_document in segment_data_list:
  1426. segment_document.enabled = False
  1427. segment_document.disabled_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1428. segment_document.status = "error"
  1429. segment_document.error = str(e)
  1430. db.session.commit()
  1431. return segment_data_list
  1432. @classmethod
  1433. def update_segment(cls, args: SegmentUpdateArgs, segment: DocumentSegment, document: Document, dataset: Dataset):
  1434. indexing_cache_key = "segment_{}_indexing".format(segment.id)
  1435. cache_result = redis_client.get(indexing_cache_key)
  1436. if cache_result is not None:
  1437. raise ValueError("Segment is indexing, please try again later")
  1438. if args.enabled is not None:
  1439. action = args.enabled
  1440. if segment.enabled != action:
  1441. if not action:
  1442. segment.enabled = action
  1443. segment.disabled_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1444. segment.disabled_by = current_user.id
  1445. db.session.add(segment)
  1446. db.session.commit()
  1447. # Set cache to prevent indexing the same segment multiple times
  1448. redis_client.setex(indexing_cache_key, 600, 1)
  1449. disable_segment_from_index_task.delay(segment.id)
  1450. return segment
  1451. if not segment.enabled:
  1452. if args.enabled is not None:
  1453. if not args.enabled:
  1454. raise ValueError("Can't update disabled segment")
  1455. else:
  1456. raise ValueError("Can't update disabled segment")
  1457. try:
  1458. word_count_change = segment.word_count
  1459. content = args.content or segment.content
  1460. if segment.content == content:
  1461. segment.word_count = len(content)
  1462. if document.doc_form == "qa_model":
  1463. segment.answer = args.answer
  1464. segment.word_count += len(args.answer) if args.answer else 0
  1465. word_count_change = segment.word_count - word_count_change
  1466. if args.keywords:
  1467. segment.keywords = args.keywords
  1468. segment.enabled = True
  1469. segment.disabled_at = None
  1470. segment.disabled_by = None
  1471. db.session.add(segment)
  1472. db.session.commit()
  1473. # update document word count
  1474. if word_count_change != 0:
  1475. document.word_count = max(0, document.word_count + word_count_change)
  1476. db.session.add(document)
  1477. # update segment index task
  1478. if args.enabled:
  1479. VectorService.create_segments_vector(
  1480. [args.keywords] if args.keywords else None,
  1481. [segment],
  1482. dataset,
  1483. document.doc_form,
  1484. )
  1485. if document.doc_form == IndexType.PARENT_CHILD_INDEX and args.regenerate_child_chunks:
  1486. # regenerate child chunks
  1487. # get embedding model instance
  1488. if dataset.indexing_technique == "high_quality":
  1489. # check embedding model setting
  1490. model_manager = ModelManager()
  1491. if dataset.embedding_model_provider:
  1492. embedding_model_instance = model_manager.get_model_instance(
  1493. tenant_id=dataset.tenant_id,
  1494. provider=dataset.embedding_model_provider,
  1495. model_type=ModelType.TEXT_EMBEDDING,
  1496. model=dataset.embedding_model,
  1497. )
  1498. else:
  1499. embedding_model_instance = model_manager.get_default_model_instance(
  1500. tenant_id=dataset.tenant_id,
  1501. model_type=ModelType.TEXT_EMBEDDING,
  1502. )
  1503. else:
  1504. raise ValueError("The knowledge base index technique is not high quality!")
  1505. # get the process rule
  1506. processing_rule = (
  1507. db.session.query(DatasetProcessRule)
  1508. .filter(DatasetProcessRule.id == document.dataset_process_rule_id)
  1509. .first()
  1510. )
  1511. if not processing_rule:
  1512. raise ValueError("No processing rule found.")
  1513. VectorService.generate_child_chunks(
  1514. segment, document, dataset, embedding_model_instance, processing_rule, True
  1515. )
  1516. else:
  1517. segment_hash = helper.generate_text_hash(content)
  1518. tokens = 0
  1519. if dataset.indexing_technique == "high_quality":
  1520. model_manager = ModelManager()
  1521. embedding_model = model_manager.get_model_instance(
  1522. tenant_id=current_user.current_tenant_id,
  1523. provider=dataset.embedding_model_provider,
  1524. model_type=ModelType.TEXT_EMBEDDING,
  1525. model=dataset.embedding_model,
  1526. )
  1527. # calc embedding use tokens
  1528. if document.doc_form == "qa_model":
  1529. tokens = embedding_model.get_text_embedding_num_tokens(texts=[content + segment.answer])
  1530. else:
  1531. tokens = embedding_model.get_text_embedding_num_tokens(texts=[content])
  1532. segment.content = content
  1533. segment.index_node_hash = segment_hash
  1534. segment.word_count = len(content)
  1535. segment.tokens = tokens
  1536. segment.status = "completed"
  1537. segment.indexing_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1538. segment.completed_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1539. segment.updated_by = current_user.id
  1540. segment.updated_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1541. segment.enabled = True
  1542. segment.disabled_at = None
  1543. segment.disabled_by = None
  1544. if document.doc_form == "qa_model":
  1545. segment.answer = args.answer
  1546. segment.word_count += len(args.answer) if args.answer else 0
  1547. word_count_change = segment.word_count - word_count_change
  1548. # update document word count
  1549. if word_count_change != 0:
  1550. document.word_count = max(0, document.word_count + word_count_change)
  1551. db.session.add(document)
  1552. db.session.add(segment)
  1553. db.session.commit()
  1554. if document.doc_form == IndexType.PARENT_CHILD_INDEX and args.regenerate_child_chunks:
  1555. # get embedding model instance
  1556. if dataset.indexing_technique == "high_quality":
  1557. # check embedding model setting
  1558. model_manager = ModelManager()
  1559. if dataset.embedding_model_provider:
  1560. embedding_model_instance = model_manager.get_model_instance(
  1561. tenant_id=dataset.tenant_id,
  1562. provider=dataset.embedding_model_provider,
  1563. model_type=ModelType.TEXT_EMBEDDING,
  1564. model=dataset.embedding_model,
  1565. )
  1566. else:
  1567. embedding_model_instance = model_manager.get_default_model_instance(
  1568. tenant_id=dataset.tenant_id,
  1569. model_type=ModelType.TEXT_EMBEDDING,
  1570. )
  1571. else:
  1572. raise ValueError("The knowledge base index technique is not high quality!")
  1573. # get the process rule
  1574. processing_rule = (
  1575. db.session.query(DatasetProcessRule)
  1576. .filter(DatasetProcessRule.id == document.dataset_process_rule_id)
  1577. .first()
  1578. )
  1579. if not processing_rule:
  1580. raise ValueError("No processing rule found.")
  1581. VectorService.generate_child_chunks(
  1582. segment, document, dataset, embedding_model_instance, processing_rule, True
  1583. )
  1584. elif document.doc_form in (IndexType.PARAGRAPH_INDEX, IndexType.QA_INDEX):
  1585. # update segment vector index
  1586. VectorService.update_segment_vector(args.keywords, segment, dataset)
  1587. except Exception as e:
  1588. logging.exception("update segment index failed")
  1589. segment.enabled = False
  1590. segment.disabled_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
  1591. segment.status = "error"
  1592. segment.error = str(e)
  1593. db.session.commit()
  1594. new_segment = db.session.query(DocumentSegment).filter(DocumentSegment.id == segment.id).first()
  1595. return new_segment
  1596. @classmethod
  1597. def delete_segment(cls, segment: DocumentSegment, document: Document, dataset: Dataset):
  1598. indexing_cache_key = "segment_{}_delete_indexing".format(segment.id)
  1599. cache_result = redis_client.get(indexing_cache_key)
  1600. if cache_result is not None:
  1601. raise ValueError("Segment is deleting.")
  1602. # enabled segment need to delete index
  1603. if segment.enabled:
  1604. # send delete segment index task
  1605. redis_client.setex(indexing_cache_key, 600, 1)
  1606. delete_segment_from_index_task.delay([segment.index_node_id], dataset.id, document.id)
  1607. db.session.delete(segment)
  1608. # update document word count
  1609. document.word_count -= segment.word_count
  1610. db.session.add(document)
  1611. db.session.commit()
  1612. @classmethod
  1613. def delete_segments(cls, segment_ids: list, document: Document, dataset: Dataset):
  1614. index_node_ids = (
  1615. DocumentSegment.query.with_entities(DocumentSegment.index_node_id)
  1616. .filter(
  1617. DocumentSegment.id.in_(segment_ids),
  1618. DocumentSegment.dataset_id == dataset.id,
  1619. DocumentSegment.document_id == document.id,
  1620. DocumentSegment.tenant_id == current_user.current_tenant_id,
  1621. )
  1622. .all()
  1623. )
  1624. index_node_ids = [index_node_id[0] for index_node_id in index_node_ids]
  1625. delete_segment_from_index_task.delay(index_node_ids, dataset.id, document.id)
  1626. db.session.query(DocumentSegment).filter(DocumentSegment.id.in_(segment_ids)).delete()
  1627. db.session.commit()
  1628. @classmethod
  1629. def update_segments_status(cls, segment_ids: list, action: str, dataset: Dataset, document: Document):
  1630. if action == "enable":
  1631. segments = (
  1632. db.session.query(DocumentSegment)
  1633. .filter(
  1634. DocumentSegment.id.in_(segment_ids),
  1635. DocumentSegment.dataset_id == dataset.id,
  1636. DocumentSegment.document_id == document.id,
  1637. DocumentSegment.enabled == False,
  1638. )
  1639. .all()
  1640. )
  1641. if not segments:
  1642. return
  1643. real_deal_segmment_ids = []
  1644. for segment in segments:
  1645. indexing_cache_key = "segment_{}_indexing".format(segment.id)
  1646. cache_result = redis_client.get(indexing_cache_key)
  1647. if cache_result is not None:
  1648. continue
  1649. segment.enabled = True
  1650. segment.disabled_at = None
  1651. segment.disabled_by = None
  1652. db.session.add(segment)
  1653. real_deal_segmment_ids.append(segment.id)
  1654. db.session.commit()
  1655. enable_segments_to_index_task.delay(real_deal_segmment_ids, dataset.id, document.id)
  1656. elif action == "disable":
  1657. segments = (
  1658. db.session.query(DocumentSegment)
  1659. .filter(
  1660. DocumentSegment.id.in_(segment_ids),
  1661. DocumentSegment.dataset_id == dataset.id,
  1662. DocumentSegment.document_id == document.id,
  1663. DocumentSegment.enabled == True,
  1664. )
  1665. .all()
  1666. )
  1667. if not segments:
  1668. return
  1669. real_deal_segmment_ids = []
  1670. for segment in segments:
  1671. indexing_cache_key = "segment_{}_indexing".format(segment.id)
  1672. cache_result = redis_client.get(indexing_cache_key)
  1673. if cache_result is not None:
  1674. continue
  1675. segment.enabled = False
  1676. segment.disabled_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
  1677. segment.disabled_by = current_user.id
  1678. db.session.add(segment)
  1679. real_deal_segmment_ids.append(segment.id)
  1680. db.session.commit()
  1681. disable_segments_from_index_task.delay(real_deal_segmment_ids, dataset.id, document.id)
  1682. else:
  1683. raise InvalidActionError()
  1684. @classmethod
  1685. def create_child_chunk(
  1686. cls, content: str, segment: DocumentSegment, document: Document, dataset: Dataset
  1687. ) -> ChildChunk:
  1688. lock_name = "add_child_lock_{}".format(segment.id)
  1689. with redis_client.lock(lock_name, timeout=20):
  1690. index_node_id = str(uuid.uuid4())
  1691. index_node_hash = helper.generate_text_hash(content)
  1692. child_chunk_count = (
  1693. db.session.query(ChildChunk)
  1694. .filter(
  1695. ChildChunk.tenant_id == current_user.current_tenant_id,
  1696. ChildChunk.dataset_id == dataset.id,
  1697. ChildChunk.document_id == document.id,
  1698. ChildChunk.segment_id == segment.id,
  1699. )
  1700. .count()
  1701. )
  1702. max_position = (
  1703. db.session.query(func.max(ChildChunk.position))
  1704. .filter(
  1705. ChildChunk.tenant_id == current_user.current_tenant_id,
  1706. ChildChunk.dataset_id == dataset.id,
  1707. ChildChunk.document_id == document.id,
  1708. ChildChunk.segment_id == segment.id,
  1709. )
  1710. .scalar()
  1711. )
  1712. child_chunk = ChildChunk(
  1713. tenant_id=current_user.current_tenant_id,
  1714. dataset_id=dataset.id,
  1715. document_id=document.id,
  1716. segment_id=segment.id,
  1717. position=max_position + 1,
  1718. index_node_id=index_node_id,
  1719. index_node_hash=index_node_hash,
  1720. content=content,
  1721. word_count=len(content),
  1722. type="customized",
  1723. created_by=current_user.id,
  1724. )
  1725. db.session.add(child_chunk)
  1726. # save vector index
  1727. try:
  1728. VectorService.create_child_chunk_vector(child_chunk, dataset)
  1729. except Exception as e:
  1730. logging.exception("create child chunk index failed")
  1731. db.session.rollback()
  1732. raise ChildChunkIndexingError(str(e))
  1733. db.session.commit()
  1734. return child_chunk
  1735. @classmethod
  1736. def update_child_chunks(
  1737. cls,
  1738. child_chunks_update_args: list[ChildChunkUpdateArgs],
  1739. segment: DocumentSegment,
  1740. document: Document,
  1741. dataset: Dataset,
  1742. ) -> list[ChildChunk]:
  1743. child_chunks = (
  1744. db.session.query(ChildChunk)
  1745. .filter(
  1746. ChildChunk.dataset_id == dataset.id,
  1747. ChildChunk.document_id == document.id,
  1748. ChildChunk.segment_id == segment.id,
  1749. )
  1750. .all()
  1751. )
  1752. child_chunks_map = {chunk.id: chunk for chunk in child_chunks}
  1753. new_child_chunks, update_child_chunks, delete_child_chunks, new_child_chunks_args = [], [], [], []
  1754. for child_chunk_update_args in child_chunks_update_args:
  1755. if child_chunk_update_args.id:
  1756. child_chunk = child_chunks_map.pop(child_chunk_update_args.id, None)
  1757. if child_chunk:
  1758. if child_chunk.content != child_chunk_update_args.content:
  1759. child_chunk.content = child_chunk_update_args.content
  1760. child_chunk.word_count = len(child_chunk.content)
  1761. child_chunk.updated_by = current_user.id
  1762. child_chunk.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
  1763. child_chunk.type = "customized"
  1764. update_child_chunks.append(child_chunk)
  1765. else:
  1766. new_child_chunks_args.append(child_chunk_update_args)
  1767. if child_chunks_map:
  1768. delete_child_chunks = list(child_chunks_map.values())
  1769. try:
  1770. if update_child_chunks:
  1771. db.session.bulk_save_objects(update_child_chunks)
  1772. if delete_child_chunks:
  1773. for child_chunk in delete_child_chunks:
  1774. db.session.delete(child_chunk)
  1775. if new_child_chunks_args:
  1776. child_chunk_count = len(child_chunks)
  1777. for position, args in enumerate(new_child_chunks_args, start=child_chunk_count + 1):
  1778. index_node_id = str(uuid.uuid4())
  1779. index_node_hash = helper.generate_text_hash(args.content)
  1780. child_chunk = ChildChunk(
  1781. tenant_id=current_user.current_tenant_id,
  1782. dataset_id=dataset.id,
  1783. document_id=document.id,
  1784. segment_id=segment.id,
  1785. position=position,
  1786. index_node_id=index_node_id,
  1787. index_node_hash=index_node_hash,
  1788. content=args.content,
  1789. word_count=len(args.content),
  1790. type="customized",
  1791. created_by=current_user.id,
  1792. )
  1793. db.session.add(child_chunk)
  1794. db.session.flush()
  1795. new_child_chunks.append(child_chunk)
  1796. VectorService.update_child_chunk_vector(new_child_chunks, update_child_chunks, delete_child_chunks, dataset)
  1797. db.session.commit()
  1798. except Exception as e:
  1799. logging.exception("update child chunk index failed")
  1800. db.session.rollback()
  1801. raise ChildChunkIndexingError(str(e))
  1802. return sorted(new_child_chunks + update_child_chunks, key=lambda x: x.position)
  1803. @classmethod
  1804. def update_child_chunk(
  1805. cls,
  1806. content: str,
  1807. child_chunk: ChildChunk,
  1808. segment: DocumentSegment,
  1809. document: Document,
  1810. dataset: Dataset,
  1811. ) -> ChildChunk:
  1812. try:
  1813. child_chunk.content = content
  1814. child_chunk.word_count = len(content)
  1815. child_chunk.updated_by = current_user.id
  1816. child_chunk.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)
  1817. child_chunk.type = "customized"
  1818. db.session.add(child_chunk)
  1819. VectorService.update_child_chunk_vector([], [child_chunk], [], dataset)
  1820. db.session.commit()
  1821. except Exception as e:
  1822. logging.exception("update child chunk index failed")
  1823. db.session.rollback()
  1824. raise ChildChunkIndexingError(str(e))
  1825. return child_chunk
  1826. @classmethod
  1827. def delete_child_chunk(cls, child_chunk: ChildChunk, dataset: Dataset):
  1828. db.session.delete(child_chunk)
  1829. try:
  1830. VectorService.delete_child_chunk_vector(child_chunk, dataset)
  1831. except Exception as e:
  1832. logging.exception("delete child chunk index failed")
  1833. db.session.rollback()
  1834. raise ChildChunkDeleteIndexError(str(e))
  1835. db.session.commit()
  1836. @classmethod
  1837. def get_child_chunks(
  1838. cls, segment_id: str, document_id: str, dataset_id: str, page: int, limit: int, keyword: Optional[str] = None
  1839. ):
  1840. query = ChildChunk.query.filter_by(
  1841. tenant_id=current_user.current_tenant_id,
  1842. dataset_id=dataset_id,
  1843. document_id=document_id,
  1844. segment_id=segment_id,
  1845. ).order_by(ChildChunk.position.asc())
  1846. if keyword:
  1847. query = query.where(ChildChunk.content.ilike(f"%{keyword}%"))
  1848. return query.paginate(page=page, per_page=limit, max_per_page=100, error_out=False)
  1849. class DatasetCollectionBindingService:
  1850. @classmethod
  1851. def get_dataset_collection_binding(
  1852. cls, provider_name: str, model_name: str, collection_type: str = "dataset"
  1853. ) -> DatasetCollectionBinding:
  1854. dataset_collection_binding = (
  1855. db.session.query(DatasetCollectionBinding)
  1856. .filter(
  1857. DatasetCollectionBinding.provider_name == provider_name,
  1858. DatasetCollectionBinding.model_name == model_name,
  1859. DatasetCollectionBinding.type == collection_type,
  1860. )
  1861. .order_by(DatasetCollectionBinding.created_at)
  1862. .first()
  1863. )
  1864. if not dataset_collection_binding:
  1865. dataset_collection_binding = DatasetCollectionBinding(
  1866. provider_name=provider_name,
  1867. model_name=model_name,
  1868. collection_name=Dataset.gen_collection_name_by_id(str(uuid.uuid4())),
  1869. type=collection_type,
  1870. )
  1871. db.session.add(dataset_collection_binding)
  1872. db.session.commit()
  1873. return dataset_collection_binding
  1874. @classmethod
  1875. def get_dataset_collection_binding_by_id_and_type(
  1876. cls, collection_binding_id: str, collection_type: str = "dataset"
  1877. ) -> DatasetCollectionBinding:
  1878. dataset_collection_binding = (
  1879. db.session.query(DatasetCollectionBinding)
  1880. .filter(
  1881. DatasetCollectionBinding.id == collection_binding_id, DatasetCollectionBinding.type == collection_type
  1882. )
  1883. .order_by(DatasetCollectionBinding.created_at)
  1884. .first()
  1885. )
  1886. if not dataset_collection_binding:
  1887. raise ValueError("Dataset collection binding not found")
  1888. return dataset_collection_binding
  1889. class DatasetPermissionService:
  1890. @classmethod
  1891. def get_dataset_partial_member_list(cls, dataset_id):
  1892. user_list_query = (
  1893. db.session.query(
  1894. DatasetPermission.account_id,
  1895. )
  1896. .filter(DatasetPermission.dataset_id == dataset_id)
  1897. .all()
  1898. )
  1899. user_list = []
  1900. for user in user_list_query:
  1901. user_list.append(user.account_id)
  1902. return user_list
  1903. @classmethod
  1904. def update_partial_member_list(cls, tenant_id, dataset_id, user_list):
  1905. try:
  1906. db.session.query(DatasetPermission).filter(DatasetPermission.dataset_id == dataset_id).delete()
  1907. permissions = []
  1908. for user in user_list:
  1909. permission = DatasetPermission(
  1910. tenant_id=tenant_id,
  1911. dataset_id=dataset_id,
  1912. account_id=user["user_id"],
  1913. )
  1914. permissions.append(permission)
  1915. db.session.add_all(permissions)
  1916. db.session.commit()
  1917. except Exception as e:
  1918. db.session.rollback()
  1919. raise e
  1920. @classmethod
  1921. def check_permission(cls, user, dataset, requested_permission, requested_partial_member_list):
  1922. if not user.is_dataset_editor:
  1923. raise NoPermissionError("User does not have permission to edit this dataset.")
  1924. if user.is_dataset_operator and dataset.permission != requested_permission:
  1925. raise NoPermissionError("Dataset operators cannot change the dataset permissions.")
  1926. if user.is_dataset_operator and requested_permission == "partial_members":
  1927. if not requested_partial_member_list:
  1928. raise ValueError("Partial member list is required when setting to partial members.")
  1929. local_member_list = cls.get_dataset_partial_member_list(dataset.id)
  1930. request_member_list = [user["user_id"] for user in requested_partial_member_list]
  1931. if set(local_member_list) != set(request_member_list):
  1932. raise ValueError("Dataset operators cannot change the dataset permissions.")
  1933. @classmethod
  1934. def clear_partial_member_list(cls, dataset_id):
  1935. try:
  1936. db.session.query(DatasetPermission).filter(DatasetPermission.dataset_id == dataset_id).delete()
  1937. db.session.commit()
  1938. except Exception as e:
  1939. db.session.rollback()
  1940. raise e