Underground-Digital/Workflow-Engine
0
1import datetime2import json3import logging4import random5import time6import uuid7from typing import Any, Optional8 9from flask_login import current_user10from sqlalchemy import func11from werkzeug.exceptions import NotFound12 13from configs import dify_config14from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError15from core.model_manager import ModelManager16from core.model_runtime.entities.model_entities import ModelType17from core.rag.datasource.keyword.keyword_factory import Keyword18from core.rag.models.document import Document as RAGDocument19from core.rag.retrieval.retrieval_methods import RetrievalMethod20from events.dataset_event import dataset_was_deleted21from events.document_event import document_was_deleted22from extensions.ext_database import db23from extensions.ext_redis import redis_client24from libs import helper25from models.account import Account, TenantAccountRole26from models.dataset import (27 AppDatasetJoin,28 Dataset,29 DatasetCollectionBinding,30 DatasetPermission,31 DatasetPermissionEnum,32 DatasetProcessRule,33 DatasetQuery,34 Document,35 DocumentSegment,36 ExternalKnowledgeBindings,37)38from models.model import UploadFile39from models.source import DataSourceOauthBinding40from services.errors.account import NoPermissionError41from services.errors.dataset import DatasetNameDuplicateError42from services.errors.document import DocumentIndexingError43from services.errors.file import FileNotExistsError44from services.external_knowledge_service import ExternalDatasetService45from services.feature_service import FeatureModel, FeatureService46from services.tag_service import TagService47from services.vector_service import VectorService48from tasks.clean_notion_document_task import clean_notion_document_task49from tasks.deal_dataset_vector_index_task import deal_dataset_vector_index_task50from tasks.delete_segment_from_index_task import delete_segment_from_index_task51from tasks.disable_segment_from_index_task import disable_segment_from_index_task52from tasks.document_indexing_task import document_indexing_task53from tasks.document_indexing_update_task import document_indexing_update_task54from tasks.duplicate_document_indexing_task import duplicate_document_indexing_task55from tasks.recover_document_indexing_task import recover_document_indexing_task56from tasks.retry_document_indexing_task import retry_document_indexing_task57from tasks.sync_website_document_indexing_task import sync_website_document_indexing_task58 59 60class DatasetService:61 @staticmethod62 def get_datasets(page, per_page, tenant_id=None, user=None, search=None, tag_ids=None):63 query = Dataset.query.filter(Dataset.tenant_id == tenant_id).order_by(Dataset.created_at.desc())64 65 if user:66 # get permitted dataset ids67 dataset_permission = DatasetPermission.query.filter_by(account_id=user.id, tenant_id=tenant_id).all()68 permitted_dataset_ids = {dp.dataset_id for dp in dataset_permission} if dataset_permission else None69 70 if user.current_role == TenantAccountRole.DATASET_OPERATOR:71 # only show datasets that the user has permission to access72 if permitted_dataset_ids:73 query = query.filter(Dataset.id.in_(permitted_dataset_ids))74 else:75 return [], 076 else:77 # show all datasets that the user has permission to access78 if permitted_dataset_ids:79 query = query.filter(80 db.or_(81 Dataset.permission == DatasetPermissionEnum.ALL_TEAM,82 db.and_(Dataset.permission == DatasetPermissionEnum.ONLY_ME, Dataset.created_by == user.id),83 db.and_(84 Dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM,85 Dataset.id.in_(permitted_dataset_ids),86 ),87 )88 )89 else:90 query = query.filter(91 db.or_(92 Dataset.permission == DatasetPermissionEnum.ALL_TEAM,93 db.and_(Dataset.permission == DatasetPermissionEnum.ONLY_ME, Dataset.created_by == user.id),94 )95 )96 else:97 # if no user, only show datasets that are shared with all team members98 query = query.filter(Dataset.permission == DatasetPermissionEnum.ALL_TEAM)99 100 if search:101 query = query.filter(Dataset.name.ilike(f"%{search}%"))102 103 if tag_ids:104 target_ids = TagService.get_target_ids_by_tag_ids("knowledge", tenant_id, tag_ids)105 if target_ids:106 query = query.filter(Dataset.id.in_(target_ids))107 else:108 return [], 0109 110 datasets = query.paginate(page=page, per_page=per_page, max_per_page=100, error_out=False)111 112 return datasets.items, datasets.total113 114 @staticmethod115 def get_process_rules(dataset_id):116 # get the latest process rule117 dataset_process_rule = (118 db.session.query(DatasetProcessRule)119 .filter(DatasetProcessRule.dataset_id == dataset_id)120 .order_by(DatasetProcessRule.created_at.desc())121 .limit(1)122 .one_or_none()123 )124 if dataset_process_rule:125 mode = dataset_process_rule.mode126 rules = dataset_process_rule.rules_dict127 else:128 mode = DocumentService.DEFAULT_RULES["mode"]129 rules = DocumentService.DEFAULT_RULES["rules"]130 return {"mode": mode, "rules": rules}131 132 @staticmethod133 def get_datasets_by_ids(ids, tenant_id):134 datasets = Dataset.query.filter(Dataset.id.in_(ids), Dataset.tenant_id == tenant_id).paginate(135 page=1, per_page=len(ids), max_per_page=len(ids), error_out=False136 )137 return datasets.items, datasets.total138 139 @staticmethod140 def create_empty_dataset(141 tenant_id: str,142 name: str,143 description: Optional[str],144 indexing_technique: Optional[str],145 account: Account,146 permission: Optional[str] = None,147 provider: str = "vendor",148 external_knowledge_api_id: Optional[str] = None,149 external_knowledge_id: Optional[str] = None,150 ):151 # check if dataset name already exists152 if Dataset.query.filter_by(name=name, tenant_id=tenant_id).first():153 raise DatasetNameDuplicateError(f"Dataset with name {name} already exists.")154 embedding_model = None155 if indexing_technique == "high_quality":156 model_manager = ModelManager()157 embedding_model = model_manager.get_default_model_instance(158 tenant_id=tenant_id, model_type=ModelType.TEXT_EMBEDDING159 )160 dataset = Dataset(name=name, indexing_technique=indexing_technique)161 # dataset = Dataset(name=name, provider=provider, config=config)162 dataset.description = description163 dataset.created_by = account.id164 dataset.updated_by = account.id165 dataset.tenant_id = tenant_id166 dataset.embedding_model_provider = embedding_model.provider if embedding_model else None167 dataset.embedding_model = embedding_model.model if embedding_model else None168 dataset.permission = permission or DatasetPermissionEnum.ONLY_ME169 dataset.provider = provider170 db.session.add(dataset)171 db.session.flush()172 173 if provider == "external" and external_knowledge_api_id:174 external_knowledge_api = ExternalDatasetService.get_external_knowledge_api(external_knowledge_api_id)175 if not external_knowledge_api:176 raise ValueError("External API template not found.")177 external_knowledge_binding = ExternalKnowledgeBindings(178 tenant_id=tenant_id,179 dataset_id=dataset.id,180 external_knowledge_api_id=external_knowledge_api_id,181 external_knowledge_id=external_knowledge_id,182 created_by=account.id,183 )184 db.session.add(external_knowledge_binding)185 186 db.session.commit()187 return dataset188 189 @staticmethod190 def get_dataset(dataset_id) -> Dataset:191 return Dataset.query.filter_by(id=dataset_id).first()192 193 @staticmethod194 def check_dataset_model_setting(dataset):195 if dataset.indexing_technique == "high_quality":196 try:197 model_manager = ModelManager()198 model_manager.get_model_instance(199 tenant_id=dataset.tenant_id,200 provider=dataset.embedding_model_provider,201 model_type=ModelType.TEXT_EMBEDDING,202 model=dataset.embedding_model,203 )204 except LLMBadRequestError:205 raise ValueError(206 "No Embedding Model available. Please configure a valid provider "207 "in the Settings -> Model Provider."208 )209 except ProviderTokenNotInitError as ex:210 raise ValueError(f"The dataset in unavailable, due to: {ex.description}")211 212 @staticmethod213 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 229 @staticmethod230 def update_dataset(dataset_id, data, user):231 dataset = DatasetService.get_dataset(dataset_id)232 233 DatasetService.check_dataset_permission(dataset, user)234 if dataset.provider == "external":235 dataset.retrieval_model = data.get("external_retrieval_model", None)236 dataset.name = data.get("name", dataset.name)237 dataset.description = data.get("description", "")238 external_knowledge_id = data.get("external_knowledge_id", None)239 dataset.permission = data.get("permission")240 db.session.add(dataset)241 if not external_knowledge_id:242 raise ValueError("External knowledge id is required.")243 external_knowledge_api_id = data.get("external_knowledge_api_id", None)244 if not external_knowledge_api_id:245 raise ValueError("External knowledge api id is required.")246 external_knowledge_binding = ExternalKnowledgeBindings.query.filter_by(dataset_id=dataset_id).first()247 if (248 external_knowledge_binding.external_knowledge_id != external_knowledge_id249 or external_knowledge_binding.external_knowledge_api_id != external_knowledge_api_id250 ):251 external_knowledge_binding.external_knowledge_id = external_knowledge_id252 external_knowledge_binding.external_knowledge_api_id = external_knowledge_api_id253 db.session.add(external_knowledge_binding)254 db.session.commit()255 else:256 data.pop("partial_member_list", None)257 data.pop("external_knowledge_api_id", None)258 data.pop("external_knowledge_id", None)259 data.pop("external_retrieval_model", None)260 filtered_data = {k: v for k, v in data.items() if v is not None or k == "description"}261 action = None262 if dataset.indexing_technique != data["indexing_technique"]:263 # if update indexing_technique264 if data["indexing_technique"] == "economy":265 action = "remove"266 filtered_data["embedding_model"] = None267 filtered_data["embedding_model_provider"] = None268 filtered_data["collection_binding_id"] = None269 elif data["indexing_technique"] == "high_quality":270 action = "add"271 # get embedding model setting272 try:273 model_manager = ModelManager()274 embedding_model = model_manager.get_model_instance(275 tenant_id=current_user.current_tenant_id,276 provider=data["embedding_model_provider"],277 model_type=ModelType.TEXT_EMBEDDING,278 model=data["embedding_model"],279 )280 filtered_data["embedding_model"] = embedding_model.model281 filtered_data["embedding_model_provider"] = embedding_model.provider282 dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(283 embedding_model.provider, embedding_model.model284 )285 filtered_data["collection_binding_id"] = dataset_collection_binding.id286 except LLMBadRequestError:287 raise ValueError(288 "No Embedding Model available. Please configure a valid provider "289 "in the Settings -> Model Provider."290 )291 except ProviderTokenNotInitError as ex:292 raise ValueError(ex.description)293 else:294 if (295 data["embedding_model_provider"] != dataset.embedding_model_provider296 or data["embedding_model"] != dataset.embedding_model297 ):298 action = "update"299 try:300 model_manager = ModelManager()301 embedding_model = model_manager.get_model_instance(302 tenant_id=current_user.current_tenant_id,303 provider=data["embedding_model_provider"],304 model_type=ModelType.TEXT_EMBEDDING,305 model=data["embedding_model"],306 )307 filtered_data["embedding_model"] = embedding_model.model308 filtered_data["embedding_model_provider"] = embedding_model.provider309 dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(310 embedding_model.provider, embedding_model.model311 )312 filtered_data["collection_binding_id"] = dataset_collection_binding.id313 except LLMBadRequestError:314 raise ValueError(315 "No Embedding Model available. Please configure a valid provider "316 "in the Settings -> Model Provider."317 )318 except ProviderTokenNotInitError as ex:319 raise ValueError(ex.description)320 321 filtered_data["updated_by"] = user.id322 filtered_data["updated_at"] = datetime.datetime.now()323 324 # update Retrieval model325 filtered_data["retrieval_model"] = data["retrieval_model"]326 327 dataset.query.filter_by(id=dataset_id).update(filtered_data)328 329 db.session.commit()330 if action:331 deal_dataset_vector_index_task.delay(dataset_id, action)332 return dataset333 334 @staticmethod335 def delete_dataset(dataset_id, user):336 dataset = DatasetService.get_dataset(dataset_id)337 338 if dataset is None:339 return False340 341 DatasetService.check_dataset_permission(dataset, user)342 343 dataset_was_deleted.send(dataset)344 345 db.session.delete(dataset)346 db.session.commit()347 return True348 349 @staticmethod350 def dataset_use_check(dataset_id) -> bool:351 count = AppDatasetJoin.query.filter_by(dataset_id=dataset_id).count()352 if count > 0:353 return True354 return False355 356 @staticmethod357 def check_dataset_permission(dataset, user):358 if dataset.tenant_id != user.current_tenant_id:359 logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")360 raise NoPermissionError("You do not have permission to access this dataset.")361 if dataset.permission == DatasetPermissionEnum.ONLY_ME and dataset.created_by != user.id:362 logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")363 raise NoPermissionError("You do not have permission to access this dataset.")364 if dataset.permission == "partial_members":365 user_permission = DatasetPermission.query.filter_by(dataset_id=dataset.id, account_id=user.id).first()366 if not user_permission and dataset.tenant_id != user.current_tenant_id and dataset.created_by != user.id:367 logging.debug(f"User {user.id} does not have permission to access dataset {dataset.id}")368 raise NoPermissionError("You do not have permission to access this dataset.")369 370 @staticmethod371 def check_dataset_operator_permission(user: Account = None, dataset: Dataset = None):372 if dataset.permission == DatasetPermissionEnum.ONLY_ME:373 if dataset.created_by != user.id:374 raise NoPermissionError("You do not have permission to access this dataset.")375 376 elif dataset.permission == DatasetPermissionEnum.PARTIAL_TEAM:377 if not any(378 dp.dataset_id == dataset.id for dp in DatasetPermission.query.filter_by(account_id=user.id).all()379 ):380 raise NoPermissionError("You do not have permission to access this dataset.")381 382 @staticmethod383 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.total390 391 @staticmethod392 def get_related_apps(dataset_id: str):393 return (394 AppDatasetJoin.query.filter(AppDatasetJoin.dataset_id == dataset_id)395 .order_by(db.desc(AppDatasetJoin.created_at))396 .all()397 )398 399 400class DocumentService:401 DEFAULT_RULES = {402 "mode": "custom",403 "rules": {404 "pre_processing_rules": [405 {"id": "remove_extra_spaces", "enabled": True},406 {"id": "remove_urls_emails", "enabled": False},407 ],408 "segmentation": {"delimiter": "\n", "max_tokens": 500, "chunk_overlap": 50},409 },410 }411 412 DOCUMENT_METADATA_SCHEMA = {413 "book": {414 "title": str,415 "language": str,416 "author": str,417 "publisher": str,418 "publication_date": str,419 "isbn": str,420 "category": str,421 },422 "web_page": {423 "title": str,424 "url": str,425 "language": str,426 "publish_date": str,427 "author/publisher": str,428 "topic/keywords": str,429 "description": str,430 },431 "paper": {432 "title": str,433 "language": str,434 "author": str,435 "publish_date": str,436 "journal/conference_name": str,437 "volume/issue/page_numbers": str,438 "doi": str,439 "topic/keywords": str,440 "abstract": str,441 },442 "social_media_post": {443 "platform": str,444 "author/username": str,445 "publish_date": str,446 "post_url": str,447 "topic/tags": str,448 },449 "wikipedia_entry": {450 "title": str,451 "language": str,452 "web_page_url": str,453 "last_edit_date": str,454 "editor/contributor": str,455 "summary/introduction": str,456 },457 "personal_document": {458 "title": str,459 "author": str,460 "creation_date": str,461 "last_modified_date": str,462 "document_type": str,463 "tags/category": str,464 },465 "business_document": {466 "title": str,467 "author": str,468 "creation_date": str,469 "last_modified_date": str,470 "document_type": str,471 "department/team": str,472 },473 "im_chat_log": {474 "chat_platform": str,475 "chat_participants/group_name": str,476 "start_date": str,477 "end_date": str,478 "summary": str,479 },480 "synced_from_notion": {481 "title": str,482 "language": str,483 "author/creator": str,484 "creation_date": str,485 "last_modified_date": str,486 "notion_page_link": str,487 "category/tags": str,488 "description": str,489 },490 "synced_from_github": {491 "repository_name": str,492 "repository_description": str,493 "repository_owner/organization": str,494 "code_filename": str,495 "code_file_path": str,496 "programming_language": str,497 "github_link": str,498 "open_source_license": str,499 "commit_date": str,500 "commit_author": str,501 },502 "others": dict,503 }504 505 @staticmethod506 def get_document(dataset_id: str, document_id: str) -> Optional[Document]:507 document = (508 db.session.query(Document).filter(Document.id == document_id, Document.dataset_id == dataset_id).first()509 )510 511 return document512 513 @staticmethod514 def get_document_by_id(document_id: str) -> Optional[Document]:515 document = db.session.query(Document).filter(Document.id == document_id).first()516 517 return document518 519 @staticmethod520 def get_document_by_dataset_id(dataset_id: str) -> list[Document]:521 documents = db.session.query(Document).filter(Document.dataset_id == dataset_id, Document.enabled == True).all()522 523 return documents524 525 @staticmethod526 def get_error_documents_by_dataset_id(dataset_id: str) -> list[Document]:527 documents = (528 db.session.query(Document)529 .filter(Document.dataset_id == dataset_id, Document.indexing_status.in_(["error", "paused"]))530 .all()531 )532 return documents533 534 @staticmethod535 def get_batch_documents(dataset_id: str, batch: str) -> list[Document]:536 documents = (537 db.session.query(Document)538 .filter(539 Document.batch == batch,540 Document.dataset_id == dataset_id,541 Document.tenant_id == current_user.current_tenant_id,542 )543 .all()544 )545 546 return documents547 548 @staticmethod549 def get_document_file_detail(file_id: str):550 file_detail = db.session.query(UploadFile).filter(UploadFile.id == file_id).one_or_none()551 return file_detail552 553 @staticmethod554 def check_archived(document):555 if document.archived:556 return True557 else:558 return False559 560 @staticmethod561 def delete_document(document):562 # trigger document_was_deleted signal563 file_id = None564 if document.data_source_type == "upload_file":565 if document.data_source_info:566 data_source_info = document.data_source_info_dict567 if data_source_info and "upload_file_id" in data_source_info:568 file_id = data_source_info["upload_file_id"]569 document_was_deleted.send(570 document.id, dataset_id=document.dataset_id, doc_form=document.doc_form, file_id=file_id571 )572 573 db.session.delete(document)574 db.session.commit()575 576 @staticmethod577 def rename_document(dataset_id: str, document_id: str, name: str) -> Document:578 dataset = DatasetService.get_dataset(dataset_id)579 if not dataset:580 raise ValueError("Dataset not found.")581 582 document = DocumentService.get_document(dataset_id, document_id)583 584 if not document:585 raise ValueError("Document not found.")586 587 if document.tenant_id != current_user.current_tenant_id:588 raise ValueError("No permission.")589 590 document.name = name591 592 db.session.add(document)593 db.session.commit()594 595 return document596 597 @staticmethod598 def pause_document(document):599 if document.indexing_status not in {"waiting", "parsing", "cleaning", "splitting", "indexing"}:600 raise DocumentIndexingError()601 # update document to be paused602 document.is_paused = True603 document.paused_by = current_user.id604 document.paused_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)605 606 db.session.add(document)607 db.session.commit()608 # set document paused flag609 indexing_cache_key = "document_{}_is_paused".format(document.id)610 redis_client.setnx(indexing_cache_key, "True")611 612 @staticmethod613 def recover_document(document):614 if not document.is_paused:615 raise DocumentIndexingError()616 # update document to be recover617 document.is_paused = False618 document.paused_by = None619 document.paused_at = None620 621 db.session.add(document)622 db.session.commit()623 # delete paused flag624 indexing_cache_key = "document_{}_is_paused".format(document.id)625 redis_client.delete(indexing_cache_key)626 # trigger async task627 recover_document_indexing_task.delay(document.dataset_id, document.id)628 629 @staticmethod630 def retry_document(dataset_id: str, documents: list[Document]):631 for document in documents:632 # add retry flag633 retry_indexing_cache_key = "document_{}_is_retried".format(document.id)634 cache_result = redis_client.get(retry_indexing_cache_key)635 if cache_result is not None:636 raise ValueError("Document is being retried, please try again later")637 # retry document indexing638 document.indexing_status = "waiting"639 db.session.add(document)640 db.session.commit()641 642 redis_client.setex(retry_indexing_cache_key, 600, 1)643 # trigger async task644 document_ids = [document.id for document in documents]645 retry_document_indexing_task.delay(dataset_id, document_ids)646 647 @staticmethod648 def sync_website_document(dataset_id: str, document: Document):649 # add sync flag650 sync_indexing_cache_key = "document_{}_is_sync".format(document.id)651 cache_result = redis_client.get(sync_indexing_cache_key)652 if cache_result is not None:653 raise ValueError("Document is being synced, please try again later")654 # sync document indexing655 document.indexing_status = "waiting"656 data_source_info = document.data_source_info_dict657 data_source_info["mode"] = "scrape"658 document.data_source_info = json.dumps(data_source_info, ensure_ascii=False)659 db.session.add(document)660 db.session.commit()661 662 redis_client.setex(sync_indexing_cache_key, 600, 1)663 664 sync_website_document_indexing_task.delay(dataset_id, document.id)665 666 @staticmethod667 def get_documents_position(dataset_id):668 document = Document.query.filter_by(dataset_id=dataset_id).order_by(Document.position.desc()).first()669 if document:670 return document.position + 1671 else:672 return 1673 674 @staticmethod675 def save_document_with_dataset_id(676 dataset: Dataset,677 document_data: dict,678 account: Account | Any,679 dataset_process_rule: Optional[DatasetProcessRule] = None,680 created_from: str = "web",681 ):682 # check document limit683 features = FeatureService.get_features(current_user.current_tenant_id)684 685 if features.billing.enabled:686 if "original_document_id" not in document_data or not document_data["original_document_id"]:687 count = 0688 if document_data["data_source"]["type"] == "upload_file":689 upload_file_list = document_data["data_source"]["info_list"]["file_info_list"]["file_ids"]690 count = len(upload_file_list)691 elif document_data["data_source"]["type"] == "notion_import":692 notion_info_list = document_data["data_source"]["info_list"]["notion_info_list"]693 for notion_info in notion_info_list:694 count = count + len(notion_info["pages"])695 elif document_data["data_source"]["type"] == "website_crawl":696 website_info = document_data["data_source"]["info_list"]["website_info_list"]697 count = len(website_info["urls"])698 batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT)699 if count > batch_upload_limit:700 raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")701 702 DocumentService.check_documents_upload_quota(count, features)703 704 # if dataset is empty, update dataset data_source_type705 if not dataset.data_source_type:706 dataset.data_source_type = document_data["data_source"]["type"]707 708 if not dataset.indexing_technique:709 if (710 "indexing_technique" not in document_data711 or document_data["indexing_technique"] not in Dataset.INDEXING_TECHNIQUE_LIST712 ):713 raise ValueError("Indexing technique is required")714 715 dataset.indexing_technique = document_data["indexing_technique"]716 if document_data["indexing_technique"] == "high_quality":717 model_manager = ModelManager()718 embedding_model = model_manager.get_default_model_instance(719 tenant_id=current_user.current_tenant_id, model_type=ModelType.TEXT_EMBEDDING720 )721 dataset.embedding_model = embedding_model.model722 dataset.embedding_model_provider = embedding_model.provider723 dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(724 embedding_model.provider, embedding_model.model725 )726 dataset.collection_binding_id = dataset_collection_binding.id727 if not dataset.retrieval_model:728 default_retrieval_model = {729 "search_method": RetrievalMethod.SEMANTIC_SEARCH.value,730 "reranking_enable": False,731 "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},732 "top_k": 2,733 "score_threshold_enabled": False,734 }735 736 dataset.retrieval_model = document_data.get("retrieval_model") or default_retrieval_model737 738 documents = []739 if document_data.get("original_document_id"):740 document = DocumentService.update_document_with_dataset_id(dataset, document_data, account)741 documents.append(document)742 batch = document.batch743 else:744 batch = time.strftime("%Y%m%d%H%M%S") + str(random.randint(100000, 999999))745 # save process rule746 if not dataset_process_rule:747 process_rule = document_data["process_rule"]748 if process_rule["mode"] == "custom":749 dataset_process_rule = DatasetProcessRule(750 dataset_id=dataset.id,751 mode=process_rule["mode"],752 rules=json.dumps(process_rule["rules"]),753 created_by=account.id,754 )755 elif process_rule["mode"] == "automatic":756 dataset_process_rule = DatasetProcessRule(757 dataset_id=dataset.id,758 mode=process_rule["mode"],759 rules=json.dumps(DatasetProcessRule.AUTOMATIC_RULES),760 created_by=account.id,761 )762 db.session.add(dataset_process_rule)763 db.session.commit()764 lock_name = "add_document_lock_dataset_id_{}".format(dataset.id)765 with redis_client.lock(lock_name, timeout=600):766 position = DocumentService.get_documents_position(dataset.id)767 document_ids = []768 duplicate_document_ids = []769 if document_data["data_source"]["type"] == "upload_file":770 upload_file_list = document_data["data_source"]["info_list"]["file_info_list"]["file_ids"]771 for file_id in upload_file_list:772 file = (773 db.session.query(UploadFile)774 .filter(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id)775 .first()776 )777 778 # raise error if file not found779 if not file:780 raise FileNotExistsError()781 782 file_name = file.name783 data_source_info = {784 "upload_file_id": file_id,785 }786 # check duplicate787 if document_data.get("duplicate", False):788 document = Document.query.filter_by(789 dataset_id=dataset.id,790 tenant_id=current_user.current_tenant_id,791 data_source_type="upload_file",792 enabled=True,793 name=file_name,794 ).first()795 if document:796 document.dataset_process_rule_id = dataset_process_rule.id797 document.updated_at = datetime.datetime.utcnow()798 document.created_from = created_from799 document.doc_form = document_data["doc_form"]800 document.doc_language = document_data["doc_language"]801 document.data_source_info = json.dumps(data_source_info)802 document.batch = batch803 document.indexing_status = "waiting"804 db.session.add(document)805 documents.append(document)806 duplicate_document_ids.append(document.id)807 continue808 document = DocumentService.build_document(809 dataset,810 dataset_process_rule.id,811 document_data["data_source"]["type"],812 document_data["doc_form"],813 document_data["doc_language"],814 data_source_info,815 created_from,816 position,817 account,818 file_name,819 batch,820 )821 db.session.add(document)822 db.session.flush()823 document_ids.append(document.id)824 documents.append(document)825 position += 1826 elif document_data["data_source"]["type"] == "notion_import":827 notion_info_list = document_data["data_source"]["info_list"]["notion_info_list"]828 exist_page_ids = []829 exist_document = {}830 documents = Document.query.filter_by(831 dataset_id=dataset.id,832 tenant_id=current_user.current_tenant_id,833 data_source_type="notion_import",834 enabled=True,835 ).all()836 if documents:837 for document in documents:838 data_source_info = json.loads(document.data_source_info)839 exist_page_ids.append(data_source_info["notion_page_id"])840 exist_document[data_source_info["notion_page_id"]] = document.id841 for notion_info in notion_info_list:842 workspace_id = notion_info["workspace_id"]843 data_source_binding = DataSourceOauthBinding.query.filter(844 db.and_(845 DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,846 DataSourceOauthBinding.provider == "notion",847 DataSourceOauthBinding.disabled == False,848 DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',849 )850 ).first()851 if not data_source_binding:852 raise ValueError("Data source binding not found.")853 for page in notion_info["pages"]:854 if page["page_id"] not in exist_page_ids:855 data_source_info = {856 "notion_workspace_id": workspace_id,857 "notion_page_id": page["page_id"],858 "notion_page_icon": page["page_icon"],859 "type": page["type"],860 }861 document = DocumentService.build_document(862 dataset,863 dataset_process_rule.id,864 document_data["data_source"]["type"],865 document_data["doc_form"],866 document_data["doc_language"],867 data_source_info,868 created_from,869 position,870 account,871 page["page_name"],872 batch,873 )874 db.session.add(document)875 db.session.flush()876 document_ids.append(document.id)877 documents.append(document)878 position += 1879 else:880 exist_document.pop(page["page_id"])881 # delete not selected documents882 if len(exist_document) > 0:883 clean_notion_document_task.delay(list(exist_document.values()), dataset.id)884 elif document_data["data_source"]["type"] == "website_crawl":885 website_info = document_data["data_source"]["info_list"]["website_info_list"]886 urls = website_info["urls"]887 for url in urls:888 data_source_info = {889 "url": url,890 "provider": website_info["provider"],891 "job_id": website_info["job_id"],892 "only_main_content": website_info.get("only_main_content", False),893 "mode": "crawl",894 }895 if len(url) > 255:896 document_name = url[:200] + "..."897 else:898 document_name = url899 document = DocumentService.build_document(900 dataset,901 dataset_process_rule.id,902 document_data["data_source"]["type"],903 document_data["doc_form"],904 document_data["doc_language"],905 data_source_info,906 created_from,907 position,908 account,909 document_name,910 batch,911 )912 db.session.add(document)913 db.session.flush()914 document_ids.append(document.id)915 documents.append(document)916 position += 1917 db.session.commit()918 919 # trigger async task920 if document_ids:921 document_indexing_task.delay(dataset.id, document_ids)922 if duplicate_document_ids:923 duplicate_document_indexing_task.delay(dataset.id, duplicate_document_ids)924 925 return documents, batch926 927 @staticmethod928 def check_documents_upload_quota(count: int, features: FeatureModel):929 can_upload_size = features.documents_upload_quota.limit - features.documents_upload_quota.size930 if count > can_upload_size:931 raise ValueError(932 f"You have reached the limit of your subscription. Only {can_upload_size} documents can be uploaded."933 )934 935 @staticmethod936 def build_document(937 dataset: Dataset,938 process_rule_id: str,939 data_source_type: str,940 document_form: str,941 document_language: str,942 data_source_info: dict,943 created_from: str,944 position: int,945 account: Account,946 name: str,947 batch: str,948 ):949 document = Document(950 tenant_id=dataset.tenant_id,951 dataset_id=dataset.id,952 position=position,953 data_source_type=data_source_type,954 data_source_info=json.dumps(data_source_info),955 dataset_process_rule_id=process_rule_id,956 batch=batch,957 name=name,958 created_from=created_from,959 created_by=account.id,960 doc_form=document_form,961 doc_language=document_language,962 )963 return document964 965 @staticmethod966 def get_tenant_documents_count():967 documents_count = Document.query.filter(968 Document.completed_at.isnot(None),969 Document.enabled == True,970 Document.archived == False,971 Document.tenant_id == current_user.current_tenant_id,972 ).count()973 return documents_count974 975 @staticmethod976 def update_document_with_dataset_id(977 dataset: Dataset,978 document_data: dict,979 account: Account,980 dataset_process_rule: Optional[DatasetProcessRule] = None,981 created_from: str = "web",982 ):983 DatasetService.check_dataset_model_setting(dataset)984 document = DocumentService.get_document(dataset.id, document_data["original_document_id"])985 if document is None:986 raise NotFound("Document not found")987 if document.display_status != "available":988 raise ValueError("Document is not available")989 # save process rule990 if document_data.get("process_rule"):991 process_rule = document_data["process_rule"]992 if process_rule["mode"] == "custom":993 dataset_process_rule = DatasetProcessRule(994 dataset_id=dataset.id,995 mode=process_rule["mode"],996 rules=json.dumps(process_rule["rules"]),997 created_by=account.id,998 )999 elif process_rule["mode"] == "automatic":1000 dataset_process_rule = DatasetProcessRule(1001 dataset_id=dataset.id,1002 mode=process_rule["mode"],1003 rules=json.dumps(DatasetProcessRule.AUTOMATIC_RULES),1004 created_by=account.id,1005 )1006 db.session.add(dataset_process_rule)1007 db.session.commit()1008 document.dataset_process_rule_id = dataset_process_rule.id1009 # update document data source1010 if document_data.get("data_source"):1011 file_name = ""1012 data_source_info = {}1013 if document_data["data_source"]["type"] == "upload_file":1014 upload_file_list = document_data["data_source"]["info_list"]["file_info_list"]["file_ids"]1015 for file_id in upload_file_list:1016 file = (1017 db.session.query(UploadFile)1018 .filter(UploadFile.tenant_id == dataset.tenant_id, UploadFile.id == file_id)1019 .first()1020 )1021 1022 # raise error if file not found1023 if not file:1024 raise FileNotExistsError()1025 1026 file_name = file.name1027 data_source_info = {1028 "upload_file_id": file_id,1029 }1030 elif document_data["data_source"]["type"] == "notion_import":1031 notion_info_list = document_data["data_source"]["info_list"]["notion_info_list"]1032 for notion_info in notion_info_list:1033 workspace_id = notion_info["workspace_id"]1034 data_source_binding = DataSourceOauthBinding.query.filter(1035 db.and_(1036 DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,1037 DataSourceOauthBinding.provider == "notion",1038 DataSourceOauthBinding.disabled == False,1039 DataSourceOauthBinding.source_info["workspace_id"] == f'"{workspace_id}"',1040 )1041 ).first()1042 if not data_source_binding:1043 raise ValueError("Data source binding not found.")1044 for page in notion_info["pages"]:1045 data_source_info = {1046 "notion_workspace_id": workspace_id,1047 "notion_page_id": page["page_id"],1048 "notion_page_icon": page["page_icon"],1049 "type": page["type"],1050 }1051 elif document_data["data_source"]["type"] == "website_crawl":1052 website_info = document_data["data_source"]["info_list"]["website_info_list"]1053 urls = website_info["urls"]1054 for url in urls:1055 data_source_info = {1056 "url": url,1057 "provider": website_info["provider"],1058 "job_id": website_info["job_id"],1059 "only_main_content": website_info.get("only_main_content", False),1060 "mode": "crawl",1061 }1062 document.data_source_type = document_data["data_source"]["type"]1063 document.data_source_info = json.dumps(data_source_info)1064 document.name = file_name1065 1066 # update document name1067 if document_data.get("name"):1068 document.name = document_data["name"]1069 # update document to be waiting1070 document.indexing_status = "waiting"1071 document.completed_at = None1072 document.processing_started_at = None1073 document.parsing_completed_at = None1074 document.cleaning_completed_at = None1075 document.splitting_completed_at = None1076 document.updated_at = datetime.datetime.now(datetime.timezone.utc).replace(tzinfo=None)1077 document.created_from = created_from1078 document.doc_form = document_data["doc_form"]1079 db.session.add(document)1080 db.session.commit()1081 # update document segment1082 update_params = {DocumentSegment.status: "re_segment"}1083 DocumentSegment.query.filter_by(document_id=document.id).update(update_params)1084 db.session.commit()1085 # trigger async task1086 document_indexing_update_task.delay(document.dataset_id, document.id)1087 return document1088 1089 @staticmethod1090 def save_document_without_dataset_id(tenant_id: str, document_data: dict, account: Account):1091 features = FeatureService.get_features(current_user.current_tenant_id)1092 1093 if features.billing.enabled:1094 count = 01095 if document_data["data_source"]["type"] == "upload_file":1096 upload_file_list = document_data["data_source"]["info_list"]["file_info_list"]["file_ids"]1097 count = len(upload_file_list)1098 elif document_data["data_source"]["type"] == "notion_import":1099 notion_info_list = document_data["data_source"]["info_list"]["notion_info_list"]1100 for notion_info in notion_info_list:1101 count = count + len(notion_info["pages"])1102 elif document_data["data_source"]["type"] == "website_crawl":1103 website_info = document_data["data_source"]["info_list"]["website_info_list"]1104 count = len(website_info["urls"])1105 batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT)1106 if count > batch_upload_limit:1107 raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")1108 1109 DocumentService.check_documents_upload_quota(count, features)1110 1111 dataset_collection_binding_id = None1112 retrieval_model = None1113 if document_data["indexing_technique"] == "high_quality":1114 dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(1115 document_data["embedding_model_provider"], document_data["embedding_model"]1116 )1117 dataset_collection_binding_id = dataset_collection_binding.id1118 if document_data.get("retrieval_model"):1119 retrieval_model = document_data["retrieval_model"]1120 else:1121 default_retrieval_model = {1122 "search_method": RetrievalMethod.SEMANTIC_SEARCH.value,1123 "reranking_enable": False,1124 "reranking_model": {"reranking_provider_name": "", "reranking_model_name": ""},1125 "top_k": 2,1126 "score_threshold_enabled": False,1127 }1128 retrieval_model = default_retrieval_model1129 # save dataset1130 dataset = Dataset(1131 tenant_id=tenant_id,1132 name="",1133 data_source_type=document_data["data_source"]["type"],1134 indexing_technique=document_data.get("indexing_technique", "high_quality"),1135 created_by=account.id,1136 embedding_model=document_data.get("embedding_model"),1137 embedding_model_provider=document_data.get("embedding_model_provider"),1138 collection_binding_id=dataset_collection_binding_id,1139 retrieval_model=retrieval_model,1140 )1141 1142 db.session.add(dataset)1143 db.session.flush()1144 1145 documents, batch = DocumentService.save_document_with_dataset_id(dataset, document_data, account)1146 1147 cut_length = 181148 cut_name = documents[0].name[:cut_length]1149 dataset.name = cut_name + "..."1150 dataset.description = "useful for when you want to answer queries about the " + documents[0].name1151 db.session.commit()1152 1153 return dataset, documents, batch1154 1155 @classmethod1156 def document_create_args_validate(cls, args: dict):1157 if "original_document_id" not in args or not args["original_document_id"]:1158 DocumentService.data_source_args_validate(args)1159 DocumentService.process_rule_args_validate(args)1160 else:1161 if ("data_source" not in args or not args["data_source"]) and (1162 "process_rule" not in args or not args["process_rule"]1163 ):1164 raise ValueError("Data source or Process rule is required")1165 else:1166 if args.get("data_source"):1167 DocumentService.data_source_args_validate(args)1168 if args.get("process_rule"):1169 DocumentService.process_rule_args_validate(args)1170 1171 @classmethod1172 def data_source_args_validate(cls, args: dict):1173 if "data_source" not in args or not args["data_source"]:1174 raise ValueError("Data source is required")1175 1176 if not isinstance(args["data_source"], dict):1177 raise ValueError("Data source is invalid")1178 1179 if "type" not in args["data_source"] or not args["data_source"]["type"]:1180 raise ValueError("Data source type is required")1181 1182 if args["data_source"]["type"] not in Document.DATA_SOURCES:1183 raise ValueError("Data source type is invalid")1184 1185 if "info_list" not in args["data_source"] or not args["data_source"]["info_list"]:1186 raise ValueError("Data source info is required")1187 1188 if args["data_source"]["type"] == "upload_file":1189 if (1190 "file_info_list" not in args["data_source"]["info_list"]1191 or not args["data_source"]["info_list"]["file_info_list"]1192 ):1193 raise ValueError("File source info is required")1194 if args["data_source"]["type"] == "notion_import":1195 if (1196 "notion_info_list" not in args["data_source"]["info_list"]1197 or not args["data_source"]["info_list"]["notion_info_list"]1198 ):1199 raise ValueError("Notion source info is required")1200 if args["data_source"]["type"] == "website_crawl":