| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  | import uuid | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-01-12 12:34:01 +08:00
										 |  |  | import pandas as pd | 
					
						
							| 
									
										
										
										
											2024-02-06 13:21:13 +08:00
										 |  |  | from flask import request | 
					
						
							| 
									
										
										
										
											2025-05-06 11:58:49 +08:00
										 |  |  | from flask_login import current_user | 
					
						
							|  |  |  | from flask_restful import Resource, marshal, reqparse | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  | from sqlalchemy import select | 
					
						
							| 
									
										
										
										
											2024-02-06 13:21:13 +08:00
										 |  |  | from werkzeug.exceptions import Forbidden, NotFound | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | import services | 
					
						
							|  |  |  | from controllers.console import api | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  | from controllers.console.app.error import ProviderNotInitializeError | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | from controllers.console.datasets.error import ( | 
					
						
							|  |  |  |     ChildChunkDeleteIndexError, | 
					
						
							|  |  |  |     ChildChunkIndexingError, | 
					
						
							|  |  |  |     InvalidActionError, | 
					
						
							|  |  |  |     NoFileUploadedError, | 
					
						
							|  |  |  |     TooManyFilesError, | 
					
						
							|  |  |  | ) | 
					
						
							| 
									
										
										
										
											2024-04-02 17:55:49 +08:00
										 |  |  | from controllers.console.wraps import ( | 
					
						
							|  |  |  |     account_initialization_required, | 
					
						
							|  |  |  |     cloud_edition_billing_knowledge_limit_check, | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     cloud_edition_billing_rate_limit_check, | 
					
						
							| 
									
										
										
										
											2024-04-02 17:55:49 +08:00
										 |  |  |     cloud_edition_billing_resource_check, | 
					
						
							| 
									
										
										
										
											2024-11-01 15:51:22 +08:00
										 |  |  |     setup_required, | 
					
						
							| 
									
										
										
										
											2024-04-02 17:55:49 +08:00
										 |  |  | ) | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  | from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError | 
					
						
							|  |  |  | from core.model_manager import ModelManager | 
					
						
							|  |  |  | from core.model_runtime.entities.model_entities import ModelType | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  | from extensions.ext_database import db | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | from extensions.ext_redis import redis_client | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | from fields.segment_fields import child_chunk_fields, segment_fields | 
					
						
							| 
									
										
										
										
											2024-01-12 12:34:01 +08:00
										 |  |  | from libs.login import login_required | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | from models.dataset import ChildChunk, DocumentSegment | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  | from services.dataset_service import DatasetService, DocumentService, SegmentService | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | from services.entities.knowledge_entities.knowledge_entities import ChildChunkUpdateArgs, SegmentUpdateArgs | 
					
						
							|  |  |  | from services.errors.chunk import ChildChunkDeleteIndexError as ChildChunkDeleteIndexServiceError | 
					
						
							|  |  |  | from services.errors.chunk import ChildChunkIndexingError as ChildChunkIndexingServiceError | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  | from tasks.batch_create_segment_to_index_task import batch_create_segment_to_index_task | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | class DatasetDocumentSegmentListApi(Resource): | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							|  |  |  |     def get(self, dataset_id, document_id): | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Dataset not found.") | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             DatasetService.check_dataset_permission(dataset, current_user) | 
					
						
							|  |  |  |         except services.errors.account.NoPermissionError as e: | 
					
						
							|  |  |  |             raise Forbidden(str(e)) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         if not document: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Document not found.") | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |         parser = reqparse.RequestParser() | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         parser.add_argument("limit", type=int, default=20, location="args") | 
					
						
							|  |  |  |         parser.add_argument("status", type=str, action="append", default=[], location="args") | 
					
						
							|  |  |  |         parser.add_argument("hit_count_gte", type=int, default=None, location="args") | 
					
						
							|  |  |  |         parser.add_argument("enabled", type=str, default="all", location="args") | 
					
						
							|  |  |  |         parser.add_argument("keyword", type=str, default=None, location="args") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         parser.add_argument("page", type=int, default=1, location="args") | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         args = parser.parse_args() | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         page = args["page"] | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         limit = min(args["limit"], 100) | 
					
						
							|  |  |  |         status_list = args["status"] | 
					
						
							|  |  |  |         hit_count_gte = args["hit_count_gte"] | 
					
						
							|  |  |  |         keyword = args["keyword"] | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         query = ( | 
					
						
							|  |  |  |             select(DocumentSegment) | 
					
						
							|  |  |  |             .filter( | 
					
						
							|  |  |  |                 DocumentSegment.document_id == str(document_id), | 
					
						
							|  |  |  |                 DocumentSegment.tenant_id == current_user.current_tenant_id, | 
					
						
							|  |  |  |             ) | 
					
						
							|  |  |  |             .order_by(DocumentSegment.position.asc()) | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |         if status_list: | 
					
						
							|  |  |  |             query = query.filter(DocumentSegment.status.in_(status_list)) | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         if hit_count_gte is not None: | 
					
						
							|  |  |  |             query = query.filter(DocumentSegment.hit_count >= hit_count_gte) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-06-07 00:45:25 +08:00
										 |  |  |         if keyword: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             query = query.where(DocumentSegment.content.ilike(f"%{keyword}%")) | 
					
						
							| 
									
										
										
										
											2023-06-07 00:45:25 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         if args["enabled"].lower() != "all": | 
					
						
							|  |  |  |             if args["enabled"].lower() == "true": | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |                 query = query.filter(DocumentSegment.enabled == True) | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             elif args["enabled"].lower() == "false": | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |                 query = query.filter(DocumentSegment.enabled == False) | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         segments = db.paginate(select=query, page=page, per_page=limit, max_per_page=100, error_out=False) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         response = { | 
					
						
							|  |  |  |             "data": marshal(segments.items, segment_fields), | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             "limit": limit, | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |             "total": segments.total, | 
					
						
							|  |  |  |             "total_pages": segments.pages, | 
					
						
							|  |  |  |             "page": page, | 
					
						
							|  |  |  |         } | 
					
						
							|  |  |  |         return response, 200 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     @cloud_edition_billing_rate_limit_check("knowledge") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |     def delete(self, dataset_id, document_id): | 
					
						
							|  |  |  |         # check dataset | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							|  |  |  |             raise NotFound("Dataset not found.") | 
					
						
							|  |  |  |         # check user's model setting | 
					
						
							|  |  |  |         DatasetService.check_dataset_model_setting(dataset) | 
					
						
							|  |  |  |         # check document | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							|  |  |  |             raise NotFound("Document not found.") | 
					
						
							|  |  |  |         segment_ids = request.args.getlist("segment_id") | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2025-03-12 12:51:00 +08:00
										 |  |  |         # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor | 
					
						
							|  |  |  |         if not current_user.is_dataset_editor: | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |             raise Forbidden() | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             DatasetService.check_dataset_permission(dataset, current_user) | 
					
						
							|  |  |  |         except services.errors.account.NoPermissionError as e: | 
					
						
							|  |  |  |             raise Forbidden(str(e)) | 
					
						
							|  |  |  |         SegmentService.delete_segments(segment_ids, document, dataset) | 
					
						
							| 
									
										
										
										
											2025-04-27 12:12:46 +08:00
										 |  |  |         return {"result": "success"}, 204 | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | class DatasetDocumentSegmentApi(Resource): | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |     @cloud_edition_billing_resource_check("vector_space") | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     @cloud_edition_billing_rate_limit_check("knowledge") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |     def patch(self, dataset_id, document_id, action): | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Dataset not found.") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							|  |  |  |             raise NotFound("Document not found.") | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |         # check user's model setting | 
					
						
							|  |  |  |         DatasetService.check_dataset_model_setting(dataset) | 
					
						
							| 
									
										
										
										
											2025-03-12 12:51:00 +08:00
										 |  |  |         # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor | 
					
						
							|  |  |  |         if not current_user.is_dataset_editor: | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |             raise Forbidden() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             DatasetService.check_dataset_permission(dataset, current_user) | 
					
						
							|  |  |  |         except services.errors.account.NoPermissionError as e: | 
					
						
							|  |  |  |             raise Forbidden(str(e)) | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         if dataset.indexing_technique == "high_quality": | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |             # check embedding model setting | 
					
						
							|  |  |  |             try: | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |                 model_manager = ModelManager() | 
					
						
							|  |  |  |                 model_manager.get_model_instance( | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |                     tenant_id=current_user.current_tenant_id, | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |                     provider=dataset.embedding_model_provider, | 
					
						
							|  |  |  |                     model_type=ModelType.TEXT_EMBEDDING, | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |                     model=dataset.embedding_model, | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |                 ) | 
					
						
							|  |  |  |             except LLMBadRequestError: | 
					
						
							|  |  |  |                 raise ProviderNotInitializeError( | 
					
						
							| 
									
										
										
										
											2025-01-21 10:12:29 +08:00
										 |  |  |                     "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |                 ) | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |             except ProviderTokenNotInitError as ex: | 
					
						
							|  |  |  |                 raise ProviderNotInitializeError(ex.description) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         segment_ids = request.args.getlist("segment_id") | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         document_indexing_cache_key = "document_{}_indexing".format(document.id) | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  |         cache_result = redis_client.get(document_indexing_cache_key) | 
					
						
							|  |  |  |         if cache_result is not None: | 
					
						
							|  |  |  |             raise InvalidActionError("Document is being indexed, please try again later") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         try: | 
					
						
							|  |  |  |             SegmentService.update_segments_status(segment_ids, action, dataset, document) | 
					
						
							|  |  |  |         except Exception as e: | 
					
						
							|  |  |  |             raise InvalidActionError(str(e)) | 
					
						
							|  |  |  |         return {"result": "success"}, 200 | 
					
						
							| 
									
										
										
										
											2023-05-15 08:51:32 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  | class DatasetDocumentSegmentAddApi(Resource): | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |     @cloud_edition_billing_resource_check("vector_space") | 
					
						
							|  |  |  |     @cloud_edition_billing_knowledge_limit_check("add_segment") | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     @cloud_edition_billing_rate_limit_check("knowledge") | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |     def post(self, dataset_id, document_id): | 
					
						
							|  |  |  |         # check dataset | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Dataset not found.") | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |         # check document | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Document not found.") | 
					
						
							| 
									
										
										
										
											2025-03-12 12:51:00 +08:00
										 |  |  |         if not current_user.is_dataset_editor: | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |             raise Forbidden() | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         # check embedding model setting | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         if dataset.indexing_technique == "high_quality": | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |             try: | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |                 model_manager = ModelManager() | 
					
						
							|  |  |  |                 model_manager.get_model_instance( | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |                     tenant_id=current_user.current_tenant_id, | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |                     provider=dataset.embedding_model_provider, | 
					
						
							|  |  |  |                     model_type=ModelType.TEXT_EMBEDDING, | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |                     model=dataset.embedding_model, | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |                 ) | 
					
						
							|  |  |  |             except LLMBadRequestError: | 
					
						
							|  |  |  |                 raise ProviderNotInitializeError( | 
					
						
							| 
									
										
										
										
											2025-01-21 10:12:29 +08:00
										 |  |  |                     "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |                 ) | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |             except ProviderTokenNotInitError as ex: | 
					
						
							|  |  |  |                 raise ProviderNotInitializeError(ex.description) | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |         try: | 
					
						
							|  |  |  |             DatasetService.check_dataset_permission(dataset, current_user) | 
					
						
							|  |  |  |         except services.errors.account.NoPermissionError as e: | 
					
						
							|  |  |  |             raise Forbidden(str(e)) | 
					
						
							|  |  |  |         # validate args | 
					
						
							|  |  |  |         parser = reqparse.RequestParser() | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         parser.add_argument("content", type=str, required=True, nullable=False, location="json") | 
					
						
							|  |  |  |         parser.add_argument("answer", type=str, required=False, nullable=True, location="json") | 
					
						
							|  |  |  |         parser.add_argument("keywords", type=list, required=False, nullable=True, location="json") | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |         args = parser.parse_args() | 
					
						
							|  |  |  |         SegmentService.segment_create_args_validate(args, document) | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         segment = SegmentService.create_segment(args, document, dataset) | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         return {"data": marshal(segment, segment_fields), "doc_form": document.doc_form}, 200 | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | class DatasetDocumentSegmentUpdateApi(Resource): | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |     @cloud_edition_billing_resource_check("vector_space") | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     @cloud_edition_billing_rate_limit_check("knowledge") | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |     def patch(self, dataset_id, document_id, segment_id): | 
					
						
							|  |  |  |         # check dataset | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Dataset not found.") | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |         # check user's model setting | 
					
						
							|  |  |  |         DatasetService.check_dataset_model_setting(dataset) | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |         # check document | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Document not found.") | 
					
						
							|  |  |  |         if dataset.indexing_technique == "high_quality": | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |             # check embedding model setting | 
					
						
							|  |  |  |             try: | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |                 model_manager = ModelManager() | 
					
						
							|  |  |  |                 model_manager.get_model_instance( | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |                     tenant_id=current_user.current_tenant_id, | 
					
						
							| 
									
										
										
										
											2024-01-02 23:42:00 +08:00
										 |  |  |                     provider=dataset.embedding_model_provider, | 
					
						
							|  |  |  |                     model_type=ModelType.TEXT_EMBEDDING, | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |                     model=dataset.embedding_model, | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |                 ) | 
					
						
							|  |  |  |             except LLMBadRequestError: | 
					
						
							|  |  |  |                 raise ProviderNotInitializeError( | 
					
						
							| 
									
										
										
										
											2025-01-21 10:12:29 +08:00
										 |  |  |                     "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |                 ) | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |             except ProviderTokenNotInitError as ex: | 
					
						
							|  |  |  |                 raise ProviderNotInitializeError(ex.description) | 
					
						
							|  |  |  |             # check segment | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |         segment_id = str(segment_id) | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         segment = ( | 
					
						
							|  |  |  |             db.session.query(DocumentSegment) | 
					
						
							|  |  |  |             .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | 
					
						
							|  |  |  |             .first() | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |         if not segment: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Segment not found.") | 
					
						
							| 
									
										
										
										
											2025-03-12 12:51:00 +08:00
										 |  |  |         # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor | 
					
						
							|  |  |  |         if not current_user.is_dataset_editor: | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |             raise Forbidden() | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             DatasetService.check_dataset_permission(dataset, current_user) | 
					
						
							|  |  |  |         except services.errors.account.NoPermissionError as e: | 
					
						
							|  |  |  |             raise Forbidden(str(e)) | 
					
						
							|  |  |  |         # validate args | 
					
						
							|  |  |  |         parser = reqparse.RequestParser() | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         parser.add_argument("content", type=str, required=True, nullable=False, location="json") | 
					
						
							|  |  |  |         parser.add_argument("answer", type=str, required=False, nullable=True, location="json") | 
					
						
							|  |  |  |         parser.add_argument("keywords", type=list, required=False, nullable=True, location="json") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         parser.add_argument( | 
					
						
							|  |  |  |             "regenerate_child_chunks", type=bool, required=False, nullable=True, default=False, location="json" | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  |         args = parser.parse_args() | 
					
						
							|  |  |  |         SegmentService.segment_create_args_validate(args, document) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         segment = SegmentService.update_segment(SegmentUpdateArgs(**args), segment, document, dataset) | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         return {"data": marshal(segment, segment_fields), "doc_form": document.doc_form}, 200 | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     @cloud_edition_billing_rate_limit_check("knowledge") | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |     def delete(self, dataset_id, document_id, segment_id): | 
					
						
							|  |  |  |         # check dataset | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Dataset not found.") | 
					
						
							| 
									
										
										
										
											2023-08-29 03:37:45 +08:00
										 |  |  |         # check user's model setting | 
					
						
							|  |  |  |         DatasetService.check_dataset_model_setting(dataset) | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         # check document | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Document not found.") | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         # check segment | 
					
						
							|  |  |  |         segment_id = str(segment_id) | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         segment = ( | 
					
						
							|  |  |  |             db.session.query(DocumentSegment) | 
					
						
							|  |  |  |             .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | 
					
						
							|  |  |  |             .first() | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         if not segment: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Segment not found.") | 
					
						
							| 
									
										
										
										
											2025-03-12 12:51:00 +08:00
										 |  |  |         # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor | 
					
						
							|  |  |  |         if not current_user.is_dataset_editor: | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |             raise Forbidden() | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             DatasetService.check_dataset_permission(dataset, current_user) | 
					
						
							|  |  |  |         except services.errors.account.NoPermissionError as e: | 
					
						
							|  |  |  |             raise Forbidden(str(e)) | 
					
						
							|  |  |  |         SegmentService.delete_segment(segment, document, dataset) | 
					
						
							| 
									
										
										
										
											2025-04-27 12:12:46 +08:00
										 |  |  |         return {"result": "success"}, 204 | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | class DatasetDocumentSegmentBatchImportApi(Resource): | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |     @cloud_edition_billing_resource_check("vector_space") | 
					
						
							|  |  |  |     @cloud_edition_billing_knowledge_limit_check("add_segment") | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     @cloud_edition_billing_rate_limit_check("knowledge") | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |     def post(self, dataset_id, document_id): | 
					
						
							|  |  |  |         # check dataset | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Dataset not found.") | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         # check document | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             raise NotFound("Document not found.") | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         # get file from request | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         file = request.files["file"] | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         # check file | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         if "file" not in request.files: | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |             raise NoFileUploadedError() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         if len(request.files) > 1: | 
					
						
							|  |  |  |             raise TooManyFilesError() | 
					
						
							|  |  |  |         # check file type | 
					
						
							| 
									
										
										
										
											2025-06-11 01:01:33 +08:00
										 |  |  |         if not file.filename or not file.filename.endswith(".csv"): | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |             raise ValueError("Invalid file type. Only CSV files are allowed") | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             # Skip the first row | 
					
						
							|  |  |  |             df = pd.read_csv(file) | 
					
						
							|  |  |  |             result = [] | 
					
						
							|  |  |  |             for index, row in df.iterrows(): | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |                 if document.doc_form == "qa_model": | 
					
						
							| 
									
										
										
										
											2025-01-13 09:06:59 +08:00
										 |  |  |                     data = {"content": row.iloc[0], "answer": row.iloc[1]} | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |                 else: | 
					
						
							| 
									
										
										
										
											2025-01-13 09:06:59 +08:00
										 |  |  |                     data = {"content": row.iloc[0]} | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |                 result.append(data) | 
					
						
							|  |  |  |             if len(result) == 0: | 
					
						
							|  |  |  |                 raise ValueError("The CSV file is empty.") | 
					
						
							|  |  |  |             # async job | 
					
						
							|  |  |  |             job_id = str(uuid.uuid4()) | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             indexing_cache_key = "segment_batch_import_{}".format(str(job_id)) | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |             # send batch add segments task | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             redis_client.setnx(indexing_cache_key, "waiting") | 
					
						
							|  |  |  |             batch_create_segment_to_index_task.delay( | 
					
						
							|  |  |  |                 str(job_id), result, dataset_id, document_id, current_user.current_tenant_id, current_user.id | 
					
						
							|  |  |  |             ) | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         except Exception as e: | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |             return {"error": str(e)}, 500 | 
					
						
							|  |  |  |         return {"job_id": job_id, "job_status": "waiting"}, 200 | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							|  |  |  |     def get(self, job_id): | 
					
						
							|  |  |  |         job_id = str(job_id) | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         indexing_cache_key = "segment_batch_import_{}".format(job_id) | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  |         cache_result = redis_client.get(indexing_cache_key) | 
					
						
							|  |  |  |         if cache_result is None: | 
					
						
							| 
									
										
										
										
											2025-04-14 10:24:01 +08:00
										 |  |  |             raise ValueError("The job does not exist.") | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  |         return {"job_id": job_id, "job_status": cache_result.decode()}, 200 | 
					
						
							| 
									
										
										
										
											2023-08-18 17:37:31 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2023-07-28 20:47:15 +08:00
										 |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | class ChildChunkAddApi(Resource): | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							|  |  |  |     @cloud_edition_billing_resource_check("vector_space") | 
					
						
							|  |  |  |     @cloud_edition_billing_knowledge_limit_check("add_segment") | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     @cloud_edition_billing_rate_limit_check("knowledge") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |     def post(self, dataset_id, document_id, segment_id): | 
					
						
							|  |  |  |         # check dataset | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							|  |  |  |             raise NotFound("Dataset not found.") | 
					
						
							|  |  |  |         # check document | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							|  |  |  |             raise NotFound("Document not found.") | 
					
						
							|  |  |  |         # check segment | 
					
						
							|  |  |  |         segment_id = str(segment_id) | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         segment = ( | 
					
						
							|  |  |  |             db.session.query(DocumentSegment) | 
					
						
							|  |  |  |             .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | 
					
						
							|  |  |  |             .first() | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         if not segment: | 
					
						
							|  |  |  |             raise NotFound("Segment not found.") | 
					
						
							| 
									
										
										
										
											2025-03-12 12:51:00 +08:00
										 |  |  |         if not current_user.is_dataset_editor: | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |             raise Forbidden() | 
					
						
							|  |  |  |         # check embedding model setting | 
					
						
							|  |  |  |         if dataset.indexing_technique == "high_quality": | 
					
						
							|  |  |  |             try: | 
					
						
							|  |  |  |                 model_manager = ModelManager() | 
					
						
							|  |  |  |                 model_manager.get_model_instance( | 
					
						
							|  |  |  |                     tenant_id=current_user.current_tenant_id, | 
					
						
							|  |  |  |                     provider=dataset.embedding_model_provider, | 
					
						
							|  |  |  |                     model_type=ModelType.TEXT_EMBEDDING, | 
					
						
							|  |  |  |                     model=dataset.embedding_model, | 
					
						
							|  |  |  |                 ) | 
					
						
							|  |  |  |             except LLMBadRequestError: | 
					
						
							|  |  |  |                 raise ProviderNotInitializeError( | 
					
						
							| 
									
										
										
										
											2025-01-21 10:12:29 +08:00
										 |  |  |                     "No Embedding Model available. Please configure a valid provider in the Settings -> Model Provider." | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |                 ) | 
					
						
							|  |  |  |             except ProviderTokenNotInitError as ex: | 
					
						
							|  |  |  |                 raise ProviderNotInitializeError(ex.description) | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             DatasetService.check_dataset_permission(dataset, current_user) | 
					
						
							|  |  |  |         except services.errors.account.NoPermissionError as e: | 
					
						
							|  |  |  |             raise Forbidden(str(e)) | 
					
						
							|  |  |  |         # validate args | 
					
						
							|  |  |  |         parser = reqparse.RequestParser() | 
					
						
							|  |  |  |         parser.add_argument("content", type=str, required=True, nullable=False, location="json") | 
					
						
							|  |  |  |         args = parser.parse_args() | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             child_chunk = SegmentService.create_child_chunk(args.get("content"), segment, document, dataset) | 
					
						
							|  |  |  |         except ChildChunkIndexingServiceError as e: | 
					
						
							|  |  |  |             raise ChildChunkIndexingError(str(e)) | 
					
						
							|  |  |  |         return {"data": marshal(child_chunk, child_chunk_fields)}, 200 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							|  |  |  |     def get(self, dataset_id, document_id, segment_id): | 
					
						
							|  |  |  |         # check dataset | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							|  |  |  |             raise NotFound("Dataset not found.") | 
					
						
							|  |  |  |         # check user's model setting | 
					
						
							|  |  |  |         DatasetService.check_dataset_model_setting(dataset) | 
					
						
							|  |  |  |         # check document | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							|  |  |  |             raise NotFound("Document not found.") | 
					
						
							|  |  |  |         # check segment | 
					
						
							|  |  |  |         segment_id = str(segment_id) | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         segment = ( | 
					
						
							|  |  |  |             db.session.query(DocumentSegment) | 
					
						
							|  |  |  |             .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | 
					
						
							|  |  |  |             .first() | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         if not segment: | 
					
						
							|  |  |  |             raise NotFound("Segment not found.") | 
					
						
							|  |  |  |         parser = reqparse.RequestParser() | 
					
						
							|  |  |  |         parser.add_argument("limit", type=int, default=20, location="args") | 
					
						
							|  |  |  |         parser.add_argument("keyword", type=str, default=None, location="args") | 
					
						
							|  |  |  |         parser.add_argument("page", type=int, default=1, location="args") | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         args = parser.parse_args() | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         page = args["page"] | 
					
						
							|  |  |  |         limit = min(args["limit"], 100) | 
					
						
							|  |  |  |         keyword = args["keyword"] | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |         child_chunks = SegmentService.get_child_chunks(segment_id, document_id, dataset_id, page, limit, keyword) | 
					
						
							|  |  |  |         return { | 
					
						
							|  |  |  |             "data": marshal(child_chunks.items, child_chunk_fields), | 
					
						
							|  |  |  |             "total": child_chunks.total, | 
					
						
							|  |  |  |             "total_pages": child_chunks.pages, | 
					
						
							|  |  |  |             "page": page, | 
					
						
							|  |  |  |             "limit": limit, | 
					
						
							|  |  |  |         }, 200 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							|  |  |  |     @cloud_edition_billing_resource_check("vector_space") | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     @cloud_edition_billing_rate_limit_check("knowledge") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |     def patch(self, dataset_id, document_id, segment_id): | 
					
						
							|  |  |  |         # check dataset | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							|  |  |  |             raise NotFound("Dataset not found.") | 
					
						
							|  |  |  |         # check user's model setting | 
					
						
							|  |  |  |         DatasetService.check_dataset_model_setting(dataset) | 
					
						
							|  |  |  |         # check document | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							|  |  |  |             raise NotFound("Document not found.") | 
					
						
							|  |  |  |             # check segment | 
					
						
							|  |  |  |         segment_id = str(segment_id) | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         segment = ( | 
					
						
							|  |  |  |             db.session.query(DocumentSegment) | 
					
						
							|  |  |  |             .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | 
					
						
							|  |  |  |             .first() | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         if not segment: | 
					
						
							|  |  |  |             raise NotFound("Segment not found.") | 
					
						
							| 
									
										
										
										
											2025-03-12 12:51:00 +08:00
										 |  |  |         # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor | 
					
						
							|  |  |  |         if not current_user.is_dataset_editor: | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |             raise Forbidden() | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             DatasetService.check_dataset_permission(dataset, current_user) | 
					
						
							|  |  |  |         except services.errors.account.NoPermissionError as e: | 
					
						
							|  |  |  |             raise Forbidden(str(e)) | 
					
						
							|  |  |  |         # validate args | 
					
						
							|  |  |  |         parser = reqparse.RequestParser() | 
					
						
							|  |  |  |         parser.add_argument("chunks", type=list, required=True, nullable=False, location="json") | 
					
						
							|  |  |  |         args = parser.parse_args() | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             chunks = [ChildChunkUpdateArgs(**chunk) for chunk in args.get("chunks")] | 
					
						
							|  |  |  |             child_chunks = SegmentService.update_child_chunks(chunks, segment, document, dataset) | 
					
						
							|  |  |  |         except ChildChunkIndexingServiceError as e: | 
					
						
							|  |  |  |             raise ChildChunkIndexingError(str(e)) | 
					
						
							|  |  |  |         return {"data": marshal(child_chunks, child_chunk_fields)}, 200 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | class ChildChunkUpdateApi(Resource): | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     @cloud_edition_billing_rate_limit_check("knowledge") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |     def delete(self, dataset_id, document_id, segment_id, child_chunk_id): | 
					
						
							|  |  |  |         # check dataset | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							|  |  |  |             raise NotFound("Dataset not found.") | 
					
						
							|  |  |  |         # check user's model setting | 
					
						
							|  |  |  |         DatasetService.check_dataset_model_setting(dataset) | 
					
						
							|  |  |  |         # check document | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							|  |  |  |             raise NotFound("Document not found.") | 
					
						
							|  |  |  |         # check segment | 
					
						
							|  |  |  |         segment_id = str(segment_id) | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         segment = ( | 
					
						
							|  |  |  |             db.session.query(DocumentSegment) | 
					
						
							|  |  |  |             .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | 
					
						
							|  |  |  |             .first() | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         if not segment: | 
					
						
							|  |  |  |             raise NotFound("Segment not found.") | 
					
						
							|  |  |  |         # check child chunk | 
					
						
							|  |  |  |         child_chunk_id = str(child_chunk_id) | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         child_chunk = ( | 
					
						
							|  |  |  |             db.session.query(ChildChunk) | 
					
						
							|  |  |  |             .filter(ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id) | 
					
						
							|  |  |  |             .first() | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         if not child_chunk: | 
					
						
							|  |  |  |             raise NotFound("Child chunk not found.") | 
					
						
							| 
									
										
										
										
											2025-03-12 12:51:00 +08:00
										 |  |  |         # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor | 
					
						
							|  |  |  |         if not current_user.is_dataset_editor: | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |             raise Forbidden() | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             DatasetService.check_dataset_permission(dataset, current_user) | 
					
						
							|  |  |  |         except services.errors.account.NoPermissionError as e: | 
					
						
							|  |  |  |             raise Forbidden(str(e)) | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             SegmentService.delete_child_chunk(child_chunk, dataset) | 
					
						
							|  |  |  |         except ChildChunkDeleteIndexServiceError as e: | 
					
						
							|  |  |  |             raise ChildChunkDeleteIndexError(str(e)) | 
					
						
							| 
									
										
										
										
											2025-04-27 12:12:46 +08:00
										 |  |  |         return {"result": "success"}, 204 | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | 
 | 
					
						
							|  |  |  |     @setup_required | 
					
						
							|  |  |  |     @login_required | 
					
						
							|  |  |  |     @account_initialization_required | 
					
						
							|  |  |  |     @cloud_edition_billing_resource_check("vector_space") | 
					
						
							| 
									
										
										
										
											2025-03-10 19:50:11 +08:00
										 |  |  |     @cloud_edition_billing_rate_limit_check("knowledge") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |     def patch(self, dataset_id, document_id, segment_id, child_chunk_id): | 
					
						
							|  |  |  |         # check dataset | 
					
						
							|  |  |  |         dataset_id = str(dataset_id) | 
					
						
							|  |  |  |         dataset = DatasetService.get_dataset(dataset_id) | 
					
						
							|  |  |  |         if not dataset: | 
					
						
							|  |  |  |             raise NotFound("Dataset not found.") | 
					
						
							|  |  |  |         # check user's model setting | 
					
						
							|  |  |  |         DatasetService.check_dataset_model_setting(dataset) | 
					
						
							|  |  |  |         # check document | 
					
						
							|  |  |  |         document_id = str(document_id) | 
					
						
							|  |  |  |         document = DocumentService.get_document(dataset_id, document_id) | 
					
						
							|  |  |  |         if not document: | 
					
						
							|  |  |  |             raise NotFound("Document not found.") | 
					
						
							|  |  |  |             # check segment | 
					
						
							|  |  |  |         segment_id = str(segment_id) | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         segment = ( | 
					
						
							|  |  |  |             db.session.query(DocumentSegment) | 
					
						
							|  |  |  |             .filter(DocumentSegment.id == str(segment_id), DocumentSegment.tenant_id == current_user.current_tenant_id) | 
					
						
							|  |  |  |             .first() | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         if not segment: | 
					
						
							|  |  |  |             raise NotFound("Segment not found.") | 
					
						
							|  |  |  |         # check child chunk | 
					
						
							|  |  |  |         child_chunk_id = str(child_chunk_id) | 
					
						
							| 
									
										
										
										
											2025-05-12 13:52:33 +08:00
										 |  |  |         child_chunk = ( | 
					
						
							|  |  |  |             db.session.query(ChildChunk) | 
					
						
							|  |  |  |             .filter(ChildChunk.id == str(child_chunk_id), ChildChunk.tenant_id == current_user.current_tenant_id) | 
					
						
							|  |  |  |             .first() | 
					
						
							|  |  |  |         ) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |         if not child_chunk: | 
					
						
							|  |  |  |             raise NotFound("Child chunk not found.") | 
					
						
							| 
									
										
										
										
											2025-03-12 12:51:00 +08:00
										 |  |  |         # The role of the current user in the ta table must be admin, owner, dataset_operator, or editor | 
					
						
							|  |  |  |         if not current_user.is_dataset_editor: | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  |             raise Forbidden() | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             DatasetService.check_dataset_permission(dataset, current_user) | 
					
						
							|  |  |  |         except services.errors.account.NoPermissionError as e: | 
					
						
							|  |  |  |             raise Forbidden(str(e)) | 
					
						
							|  |  |  |         # validate args | 
					
						
							|  |  |  |         parser = reqparse.RequestParser() | 
					
						
							|  |  |  |         parser.add_argument("content", type=str, required=True, nullable=False, location="json") | 
					
						
							|  |  |  |         args = parser.parse_args() | 
					
						
							|  |  |  |         try: | 
					
						
							|  |  |  |             child_chunk = SegmentService.update_child_chunk( | 
					
						
							|  |  |  |                 args.get("content"), child_chunk, segment, document, dataset | 
					
						
							|  |  |  |             ) | 
					
						
							|  |  |  |         except ChildChunkIndexingServiceError as e: | 
					
						
							|  |  |  |             raise ChildChunkIndexingError(str(e)) | 
					
						
							|  |  |  |         return {"data": marshal(child_chunk, child_chunk_fields)}, 200 | 
					
						
							|  |  |  | 
 | 
					
						
							|  |  |  | 
 | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  | api.add_resource(DatasetDocumentSegmentListApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments") | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | api.add_resource( | 
					
						
							|  |  |  |     DatasetDocumentSegmentApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segment/<string:action>" | 
					
						
							|  |  |  | ) | 
					
						
							| 
									
										
										
										
											2024-08-26 15:29:10 +08:00
										 |  |  | api.add_resource(DatasetDocumentSegmentAddApi, "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segment") | 
					
						
							|  |  |  | api.add_resource( | 
					
						
							|  |  |  |     DatasetDocumentSegmentUpdateApi, | 
					
						
							|  |  |  |     "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/<uuid:segment_id>", | 
					
						
							|  |  |  | ) | 
					
						
							|  |  |  | api.add_resource( | 
					
						
							|  |  |  |     DatasetDocumentSegmentBatchImportApi, | 
					
						
							|  |  |  |     "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/batch_import", | 
					
						
							|  |  |  |     "/datasets/batch_import_status/<uuid:job_id>", | 
					
						
							|  |  |  | ) | 
					
						
							| 
									
										
										
										
											2024-12-25 19:49:07 +08:00
										 |  |  | api.add_resource( | 
					
						
							|  |  |  |     ChildChunkAddApi, | 
					
						
							|  |  |  |     "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/<uuid:segment_id>/child_chunks", | 
					
						
							|  |  |  | ) | 
					
						
							|  |  |  | api.add_resource( | 
					
						
							|  |  |  |     ChildChunkUpdateApi, | 
					
						
							|  |  |  |     "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/<uuid:segment_id>/child_chunks/<uuid:child_chunk_id>", | 
					
						
							|  |  |  | ) |