|
@@ -1,114 +1,291 @@
|
|
|
import datetime
|
|
|
+import json
|
|
|
import uuid
|
|
|
|
|
|
-from flask import current_app
|
|
|
-from flask_restful import reqparse
|
|
|
+from flask import current_app, request
|
|
|
+from flask_restful import reqparse, marshal
|
|
|
+from sqlalchemy import desc
|
|
|
from werkzeug.exceptions import NotFound
|
|
|
|
|
|
import services.dataset_service
|
|
|
from controllers.service_api import api
|
|
|
from controllers.service_api.app.error import ProviderNotInitializeError
|
|
|
from controllers.service_api.dataset.error import ArchivedDocumentImmutableError, DocumentIndexingError, \
|
|
|
- DatasetNotInitedError
|
|
|
+ NoFileUploadedError, TooManyFilesError
|
|
|
from controllers.service_api.wraps import DatasetApiResource
|
|
|
+from core.login.login import current_user
|
|
|
from core.model_providers.error import ProviderTokenNotInitError
|
|
|
from extensions.ext_database import db
|
|
|
from extensions.ext_storage import storage
|
|
|
+from fields.document_fields import document_fields, document_status_fields
|
|
|
+from models.dataset import Dataset, Document, DocumentSegment
|
|
|
from models.model import UploadFile
|
|
|
from services.dataset_service import DocumentService
|
|
|
+from services.file_service import FileService
|
|
|
|
|
|
|
|
|
-class DocumentListApi(DatasetApiResource):
|
|
|
+class DocumentAddByTextApi(DatasetApiResource):
|
|
|
"""Resource for documents."""
|
|
|
|
|
|
- def post(self, dataset):
|
|
|
- """Create document."""
|
|
|
+ def post(self, tenant_id, dataset_id):
|
|
|
+ """Create document by text."""
|
|
|
parser = reqparse.RequestParser()
|
|
|
parser.add_argument('name', type=str, required=True, nullable=False, location='json')
|
|
|
parser.add_argument('text', type=str, required=True, nullable=False, location='json')
|
|
|
- parser.add_argument('doc_type', type=str, location='json')
|
|
|
- parser.add_argument('doc_metadata', type=dict, location='json')
|
|
|
+ parser.add_argument('process_rule', type=dict, required=False, nullable=True, location='json')
|
|
|
+ parser.add_argument('original_document_id', type=str, required=False, location='json')
|
|
|
+ parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
|
|
|
+ parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False,
|
|
|
+ location='json')
|
|
|
+ parser.add_argument('indexing_technique', type=str, choices=Dataset.INDEXING_TECHNIQUE_LIST, nullable=False,
|
|
|
+ location='json')
|
|
|
+ args = parser.parse_args()
|
|
|
+ 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 ValueError('Dataset is not exist.')
|
|
|
+
|
|
|
+ if not dataset.indexing_technique and not args['indexing_technique']:
|
|
|
+ raise ValueError('indexing_technique is required.')
|
|
|
+
|
|
|
+ upload_file = FileService.upload_text(args.get('text'), args.get('name'))
|
|
|
+ data_source = {
|
|
|
+ 'type': 'upload_file',
|
|
|
+ 'info_list': {
|
|
|
+ 'data_source_type': 'upload_file',
|
|
|
+ 'file_info_list': {
|
|
|
+ 'file_ids': [upload_file.id]
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ args['data_source'] = data_source
|
|
|
+ # validate args
|
|
|
+ DocumentService.document_create_args_validate(args)
|
|
|
+
|
|
|
+ try:
|
|
|
+ documents, batch = DocumentService.save_document_with_dataset_id(
|
|
|
+ dataset=dataset,
|
|
|
+ document_data=args,
|
|
|
+ account=current_user,
|
|
|
+ dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
|
|
|
+ created_from='api'
|
|
|
+ )
|
|
|
+ except ProviderTokenNotInitError as ex:
|
|
|
+ raise ProviderNotInitializeError(ex.description)
|
|
|
+ document = documents[0]
|
|
|
+
|
|
|
+ documents_and_batch_fields = {
|
|
|
+ 'document': marshal(document, document_fields),
|
|
|
+ 'batch': batch
|
|
|
+ }
|
|
|
+ return documents_and_batch_fields, 200
|
|
|
+
|
|
|
+
|
|
|
+class DocumentUpdateByTextApi(DatasetApiResource):
|
|
|
+ """Resource for update documents."""
|
|
|
+
|
|
|
+ def post(self, tenant_id, dataset_id, document_id):
|
|
|
+ """Update document by text."""
|
|
|
+ parser = reqparse.RequestParser()
|
|
|
+ parser.add_argument('name', type=str, required=False, nullable=True, location='json')
|
|
|
+ parser.add_argument('text', type=str, required=False, nullable=True, location='json')
|
|
|
+ parser.add_argument('process_rule', type=dict, required=False, nullable=True, location='json')
|
|
|
+ parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
|
|
|
+ parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False,
|
|
|
+ location='json')
|
|
|
args = parser.parse_args()
|
|
|
+ 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.indexing_technique:
|
|
|
- raise DatasetNotInitedError("Dataset indexing technique must be set.")
|
|
|
-
|
|
|
- doc_type = args.get('doc_type')
|
|
|
- doc_metadata = args.get('doc_metadata')
|
|
|
-
|
|
|
- if doc_type and doc_type not in DocumentService.DOCUMENT_METADATA_SCHEMA:
|
|
|
- raise ValueError('Invalid doc_type.')
|
|
|
-
|
|
|
- # user uuid as file name
|
|
|
- file_uuid = str(uuid.uuid4())
|
|
|
- file_key = 'upload_files/' + dataset.tenant_id + '/' + file_uuid + '.txt'
|
|
|
-
|
|
|
- # save file to storage
|
|
|
- storage.save(file_key, args.get('text'))
|
|
|
-
|
|
|
- # save file to db
|
|
|
- config = current_app.config
|
|
|
- upload_file = UploadFile(
|
|
|
- tenant_id=dataset.tenant_id,
|
|
|
- storage_type=config['STORAGE_TYPE'],
|
|
|
- key=file_key,
|
|
|
- name=args.get('name') + '.txt',
|
|
|
- size=len(args.get('text')),
|
|
|
- extension='txt',
|
|
|
- mime_type='text/plain',
|
|
|
- created_by=dataset.created_by,
|
|
|
- created_at=datetime.datetime.utcnow(),
|
|
|
- used=True,
|
|
|
- used_by=dataset.created_by,
|
|
|
- used_at=datetime.datetime.utcnow()
|
|
|
- )
|
|
|
-
|
|
|
- db.session.add(upload_file)
|
|
|
- db.session.commit()
|
|
|
-
|
|
|
- document_data = {
|
|
|
- 'data_source': {
|
|
|
+ if not dataset:
|
|
|
+ raise ValueError('Dataset is not exist.')
|
|
|
+
|
|
|
+ if args['text']:
|
|
|
+ upload_file = FileService.upload_text(args.get('text'), args.get('name'))
|
|
|
+ data_source = {
|
|
|
'type': 'upload_file',
|
|
|
- 'info': [
|
|
|
- {
|
|
|
- 'upload_file_id': upload_file.id
|
|
|
+ 'info_list': {
|
|
|
+ 'data_source_type': 'upload_file',
|
|
|
+ 'file_info_list': {
|
|
|
+ 'file_ids': [upload_file.id]
|
|
|
}
|
|
|
- ]
|
|
|
+ }
|
|
|
+ }
|
|
|
+ args['data_source'] = data_source
|
|
|
+ # validate args
|
|
|
+ args['original_document_id'] = str(document_id)
|
|
|
+ DocumentService.document_create_args_validate(args)
|
|
|
+
|
|
|
+ try:
|
|
|
+ documents, batch = DocumentService.save_document_with_dataset_id(
|
|
|
+ dataset=dataset,
|
|
|
+ document_data=args,
|
|
|
+ account=current_user,
|
|
|
+ dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
|
|
|
+ created_from='api'
|
|
|
+ )
|
|
|
+ except ProviderTokenNotInitError as ex:
|
|
|
+ raise ProviderNotInitializeError(ex.description)
|
|
|
+ document = documents[0]
|
|
|
+
|
|
|
+ documents_and_batch_fields = {
|
|
|
+ 'document': marshal(document, document_fields),
|
|
|
+ 'batch': batch
|
|
|
+ }
|
|
|
+ return documents_and_batch_fields, 200
|
|
|
+
|
|
|
+
|
|
|
+class DocumentAddByFileApi(DatasetApiResource):
|
|
|
+ """Resource for documents."""
|
|
|
+ def post(self, tenant_id, dataset_id):
|
|
|
+ """Create document by upload file."""
|
|
|
+ args = {}
|
|
|
+ if 'data' in request.form:
|
|
|
+ args = json.loads(request.form['data'])
|
|
|
+ if 'doc_form' not in args:
|
|
|
+ args['doc_form'] = 'text_model'
|
|
|
+ if 'doc_language' not in args:
|
|
|
+ args['doc_language'] = 'English'
|
|
|
+ # get dataset info
|
|
|
+ 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 ValueError('Dataset is not exist.')
|
|
|
+ if not dataset.indexing_technique and not args['indexing_technique']:
|
|
|
+ raise ValueError('indexing_technique is required.')
|
|
|
+
|
|
|
+ # save file info
|
|
|
+ file = request.files['file']
|
|
|
+ # check file
|
|
|
+ if 'file' not in request.files:
|
|
|
+ raise NoFileUploadedError()
|
|
|
+
|
|
|
+ if len(request.files) > 1:
|
|
|
+ raise TooManyFilesError()
|
|
|
+
|
|
|
+ upload_file = FileService.upload_file(file)
|
|
|
+ data_source = {
|
|
|
+ 'type': 'upload_file',
|
|
|
+ 'info_list': {
|
|
|
+ 'file_info_list': {
|
|
|
+ 'file_ids': [upload_file.id]
|
|
|
+ }
|
|
|
}
|
|
|
}
|
|
|
+ args['data_source'] = data_source
|
|
|
+ # validate args
|
|
|
+ DocumentService.document_create_args_validate(args)
|
|
|
|
|
|
try:
|
|
|
documents, batch = DocumentService.save_document_with_dataset_id(
|
|
|
dataset=dataset,
|
|
|
- document_data=document_data,
|
|
|
+ document_data=args,
|
|
|
account=dataset.created_by_account,
|
|
|
- dataset_process_rule=dataset.latest_process_rule,
|
|
|
+ dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
|
|
|
created_from='api'
|
|
|
)
|
|
|
except ProviderTokenNotInitError as ex:
|
|
|
raise ProviderNotInitializeError(ex.description)
|
|
|
document = documents[0]
|
|
|
- if doc_type and doc_metadata:
|
|
|
- metadata_schema = DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type]
|
|
|
+ documents_and_batch_fields = {
|
|
|
+ 'document': marshal(document, document_fields),
|
|
|
+ 'batch': batch
|
|
|
+ }
|
|
|
+ return documents_and_batch_fields, 200
|
|
|
|
|
|
- document.doc_metadata = {}
|
|
|
|
|
|
- for key, value_type in metadata_schema.items():
|
|
|
- value = doc_metadata.get(key)
|
|
|
- if value is not None and isinstance(value, value_type):
|
|
|
- document.doc_metadata[key] = value
|
|
|
+class DocumentUpdateByFileApi(DatasetApiResource):
|
|
|
+ """Resource for update documents."""
|
|
|
|
|
|
- document.doc_type = doc_type
|
|
|
- document.updated_at = datetime.datetime.utcnow()
|
|
|
- db.session.commit()
|
|
|
+ def post(self, tenant_id, dataset_id, document_id):
|
|
|
+ """Update document by upload file."""
|
|
|
+ args = {}
|
|
|
+ if 'data' in request.form:
|
|
|
+ args = json.loads(request.form['data'])
|
|
|
+ if 'doc_form' not in args:
|
|
|
+ args['doc_form'] = 'text_model'
|
|
|
+ if 'doc_language' not in args:
|
|
|
+ args['doc_language'] = 'English'
|
|
|
|
|
|
- return {'id': document.id}
|
|
|
+ # get dataset info
|
|
|
+ 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 ValueError('Dataset is not exist.')
|
|
|
+ if 'file' in request.files:
|
|
|
+ # save file info
|
|
|
+ file = request.files['file']
|
|
|
+
|
|
|
+
|
|
|
+ if len(request.files) > 1:
|
|
|
+ raise TooManyFilesError()
|
|
|
+
|
|
|
+ upload_file = FileService.upload_file(file)
|
|
|
+ data_source = {
|
|
|
+ 'type': 'upload_file',
|
|
|
+ 'info_list': {
|
|
|
+ 'file_info_list': {
|
|
|
+ 'file_ids': [upload_file.id]
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ args['data_source'] = data_source
|
|
|
+ # validate args
|
|
|
+ args['original_document_id'] = str(document_id)
|
|
|
+ DocumentService.document_create_args_validate(args)
|
|
|
+
|
|
|
+ try:
|
|
|
+ documents, batch = DocumentService.save_document_with_dataset_id(
|
|
|
+ dataset=dataset,
|
|
|
+ document_data=args,
|
|
|
+ account=dataset.created_by_account,
|
|
|
+ dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
|
|
|
+ created_from='api'
|
|
|
+ )
|
|
|
+ except ProviderTokenNotInitError as ex:
|
|
|
+ raise ProviderNotInitializeError(ex.description)
|
|
|
+ document = documents[0]
|
|
|
+ documents_and_batch_fields = {
|
|
|
+ 'document': marshal(document, document_fields),
|
|
|
+ 'batch': batch
|
|
|
+ }
|
|
|
+ return documents_and_batch_fields, 200
|
|
|
|
|
|
|
|
|
-class DocumentApi(DatasetApiResource):
|
|
|
- def delete(self, dataset, document_id):
|
|
|
+class DocumentDeleteApi(DatasetApiResource):
|
|
|
+ def delete(self, tenant_id, dataset_id, document_id):
|
|
|
"""Delete document."""
|
|
|
document_id = str(document_id)
|
|
|
+ dataset_id = str(dataset_id)
|
|
|
+ tenant_id = str(tenant_id)
|
|
|
+
|
|
|
+ # get dataset info
|
|
|
+ dataset = db.session.query(Dataset).filter(
|
|
|
+ Dataset.tenant_id == tenant_id,
|
|
|
+ Dataset.id == dataset_id
|
|
|
+ ).first()
|
|
|
+
|
|
|
+ if not dataset:
|
|
|
+ raise ValueError('Dataset is not exist.')
|
|
|
|
|
|
document = DocumentService.get_document(dataset.id, document_id)
|
|
|
|
|
@@ -126,8 +303,85 @@ class DocumentApi(DatasetApiResource):
|
|
|
except services.errors.document.DocumentIndexingError:
|
|
|
raise DocumentIndexingError('Cannot delete document during indexing.')
|
|
|
|
|
|
- return {'result': 'success'}, 204
|
|
|
+ return {'result': 'success'}, 200
|
|
|
+
|
|
|
+
|
|
|
+class DocumentListApi(DatasetApiResource):
|
|
|
+ def get(self, tenant_id, dataset_id):
|
|
|
+ dataset_id = str(dataset_id)
|
|
|
+ tenant_id = str(tenant_id)
|
|
|
+ page = request.args.get('page', default=1, type=int)
|
|
|
+ limit = request.args.get('limit', default=20, type=int)
|
|
|
+ search = request.args.get('keyword', default=None, type=str)
|
|
|
+ dataset = db.session.query(Dataset).filter(
|
|
|
+ Dataset.tenant_id == tenant_id,
|
|
|
+ Dataset.id == dataset_id
|
|
|
+ ).first()
|
|
|
+ if not dataset:
|
|
|
+ raise NotFound('Dataset not found.')
|
|
|
+
|
|
|
+ query = Document.query.filter_by(
|
|
|
+ dataset_id=str(dataset_id), tenant_id=tenant_id)
|
|
|
+
|
|
|
+ if search:
|
|
|
+ search = f'%{search}%'
|
|
|
+ query = query.filter(Document.name.like(search))
|
|
|
+
|
|
|
+ query = query.order_by(desc(Document.created_at))
|
|
|
+
|
|
|
+ paginated_documents = query.paginate(
|
|
|
+ page=page, per_page=limit, max_per_page=100, error_out=False)
|
|
|
+ documents = paginated_documents.items
|
|
|
+
|
|
|
+ response = {
|
|
|
+ 'data': marshal(documents, document_fields),
|
|
|
+ 'has_more': len(documents) == limit,
|
|
|
+ 'limit': limit,
|
|
|
+ 'total': paginated_documents.total,
|
|
|
+ 'page': page
|
|
|
+ }
|
|
|
+
|
|
|
+ return response
|
|
|
+
|
|
|
+
|
|
|
+class DocumentIndexingStatusApi(DatasetApiResource):
|
|
|
+ def get(self, tenant_id, dataset_id, batch):
|
|
|
+ dataset_id = str(dataset_id)
|
|
|
+ batch = str(batch)
|
|
|
+ tenant_id = str(tenant_id)
|
|
|
+ # get dataset
|
|
|
+ dataset = db.session.query(Dataset).filter(
|
|
|
+ Dataset.tenant_id == tenant_id,
|
|
|
+ Dataset.id == dataset_id
|
|
|
+ ).first()
|
|
|
+ if not dataset:
|
|
|
+ raise NotFound('Dataset not found.')
|
|
|
+ # get documents
|
|
|
+ documents = DocumentService.get_batch_documents(dataset_id, batch)
|
|
|
+ if not documents:
|
|
|
+ raise NotFound('Documents not found.')
|
|
|
+ documents_status = []
|
|
|
+ for document in documents:
|
|
|
+ completed_segments = DocumentSegment.query.filter(DocumentSegment.completed_at.isnot(None),
|
|
|
+ DocumentSegment.document_id == str(document.id),
|
|
|
+ DocumentSegment.status != 're_segment').count()
|
|
|
+ total_segments = DocumentSegment.query.filter(DocumentSegment.document_id == str(document.id),
|
|
|
+ DocumentSegment.status != 're_segment').count()
|
|
|
+ document.completed_segments = completed_segments
|
|
|
+ document.total_segments = total_segments
|
|
|
+ if document.is_paused:
|
|
|
+ document.indexing_status = 'paused'
|
|
|
+ documents_status.append(marshal(document, document_status_fields))
|
|
|
+ data = {
|
|
|
+ 'data': documents_status
|
|
|
+ }
|
|
|
+ return data
|
|
|
|
|
|
|
|
|
-api.add_resource(DocumentListApi, '/documents')
|
|
|
-api.add_resource(DocumentApi, '/documents/<uuid:document_id>')
|
|
|
+api.add_resource(DocumentAddByTextApi, '/datasets/<uuid:dataset_id>/document/create_by_text')
|
|
|
+api.add_resource(DocumentAddByFileApi, '/datasets/<uuid:dataset_id>/document/create_by_file')
|
|
|
+api.add_resource(DocumentUpdateByTextApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_text')
|
|
|
+api.add_resource(DocumentUpdateByFileApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_file')
|
|
|
+api.add_resource(DocumentDeleteApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>')
|
|
|
+api.add_resource(DocumentListApi, '/datasets/<uuid:dataset_id>/documents')
|
|
|
+api.add_resource(DocumentIndexingStatusApi, '/datasets/<uuid:dataset_id>/documents/<string:batch>/indexing-status')
|