Spaces:
Running
Running
| from flask_login import current_user | |
| from flask_restful import marshal, reqparse | |
| from werkzeug.exceptions import NotFound | |
| from controllers.service_api import api | |
| from controllers.service_api.app.error import ProviderNotInitializeError | |
| from controllers.service_api.wraps import ( | |
| DatasetApiResource, | |
| cloud_edition_billing_knowledge_limit_check, | |
| cloud_edition_billing_resource_check, | |
| ) | |
| from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError | |
| from core.model_manager import ModelManager | |
| from core.model_runtime.entities.model_entities import ModelType | |
| from extensions.ext_database import db | |
| from fields.segment_fields import segment_fields | |
| from models.dataset import Dataset, DocumentSegment | |
| from services.dataset_service import DatasetService, DocumentService, SegmentService | |
| class SegmentApi(DatasetApiResource): | |
| """Resource for segments.""" | |
| def post(self, tenant_id, dataset_id, document_id): | |
| """Create single segment.""" | |
| # check dataset | |
| dataset_id = str(dataset_id) | |
| tenant_id = str(tenant_id) | |
| dataset = db.session.query(Dataset).filter( | |
| Dataset.tenant_id == tenant_id, | |
| Dataset.id == dataset_id | |
| ).first() | |
| 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 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( | |
| "No Embedding Model available. Please configure a valid provider " | |
| "in the Settings -> Model Provider.") | |
| except ProviderTokenNotInitError as ex: | |
| raise ProviderNotInitializeError(ex.description) | |
| # validate args | |
| parser = reqparse.RequestParser() | |
| parser.add_argument('segments', type=list, required=False, nullable=True, location='json') | |
| args = parser.parse_args() | |
| for args_item in args['segments']: | |
| SegmentService.segment_create_args_validate(args_item, document) | |
| segments = SegmentService.multi_create_segment(args['segments'], document, dataset) | |
| return { | |
| 'data': marshal(segments, segment_fields), | |
| 'doc_form': document.doc_form | |
| }, 200 | |
| def get(self, tenant_id, dataset_id, document_id): | |
| """Create single segment.""" | |
| # check dataset | |
| dataset_id = str(dataset_id) | |
| tenant_id = str(tenant_id) | |
| dataset = db.session.query(Dataset).filter( | |
| Dataset.tenant_id == tenant_id, | |
| Dataset.id == dataset_id | |
| ).first() | |
| 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 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( | |
| "No Embedding Model available. Please configure a valid provider " | |
| "in the Settings -> Model Provider.") | |
| except ProviderTokenNotInitError as ex: | |
| raise ProviderNotInitializeError(ex.description) | |
| parser = reqparse.RequestParser() | |
| parser.add_argument('status', type=str, | |
| action='append', default=[], location='args') | |
| parser.add_argument('keyword', type=str, default=None, location='args') | |
| args = parser.parse_args() | |
| status_list = args['status'] | |
| keyword = args['keyword'] | |
| query = DocumentSegment.query.filter( | |
| DocumentSegment.document_id == str(document_id), | |
| DocumentSegment.tenant_id == current_user.current_tenant_id | |
| ) | |
| if status_list: | |
| query = query.filter(DocumentSegment.status.in_(status_list)) | |
| if keyword: | |
| query = query.where(DocumentSegment.content.ilike(f'%{keyword}%')) | |
| total = query.count() | |
| segments = query.order_by(DocumentSegment.position).all() | |
| return { | |
| 'data': marshal(segments, segment_fields), | |
| 'doc_form': document.doc_form, | |
| 'total': total | |
| }, 200 | |
| class DatasetSegmentApi(DatasetApiResource): | |
| def delete(self, tenant_id, dataset_id, document_id, segment_id): | |
| # check dataset | |
| dataset_id = str(dataset_id) | |
| tenant_id = str(tenant_id) | |
| dataset = db.session.query(Dataset).filter( | |
| Dataset.tenant_id == tenant_id, | |
| Dataset.id == dataset_id | |
| ).first() | |
| 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 = DocumentSegment.query.filter( | |
| DocumentSegment.id == str(segment_id), | |
| DocumentSegment.tenant_id == current_user.current_tenant_id | |
| ).first() | |
| if not segment: | |
| raise NotFound('Segment not found.') | |
| SegmentService.delete_segment(segment, document, dataset) | |
| return {'result': 'success'}, 200 | |
| def post(self, tenant_id, dataset_id, document_id, segment_id): | |
| # check dataset | |
| dataset_id = str(dataset_id) | |
| tenant_id = str(tenant_id) | |
| dataset = db.session.query(Dataset).filter( | |
| Dataset.tenant_id == tenant_id, | |
| Dataset.id == dataset_id | |
| ).first() | |
| 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.') | |
| if dataset.indexing_technique == 'high_quality': | |
| # check embedding model setting | |
| 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( | |
| "No Embedding Model available. Please configure a valid provider " | |
| "in the Settings -> Model Provider.") | |
| except ProviderTokenNotInitError as ex: | |
| raise ProviderNotInitializeError(ex.description) | |
| # check segment | |
| segment_id = str(segment_id) | |
| segment = DocumentSegment.query.filter( | |
| DocumentSegment.id == str(segment_id), | |
| DocumentSegment.tenant_id == current_user.current_tenant_id | |
| ).first() | |
| if not segment: | |
| raise NotFound('Segment not found.') | |
| # validate args | |
| parser = reqparse.RequestParser() | |
| parser.add_argument('segment', type=dict, required=False, nullable=True, location='json') | |
| args = parser.parse_args() | |
| SegmentService.segment_create_args_validate(args['segment'], document) | |
| segment = SegmentService.update_segment(args['segment'], segment, document, dataset) | |
| return { | |
| 'data': marshal(segment, segment_fields), | |
| 'doc_form': document.doc_form | |
| }, 200 | |
| api.add_resource(SegmentApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments') | |
| api.add_resource(DatasetSegmentApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/<uuid:segment_id>') | |