mirror of
https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git
synced 2025-04-23 14:19:42 +08:00
881 lines
36 KiB
Python
881 lines
36 KiB
Python
# -*- coding:utf-8 -*-
|
|
from datetime import datetime
|
|
from typing import List
|
|
|
|
import services
|
|
from controllers.console import api
|
|
from controllers.console.app.error import (ProviderModelCurrentlyNotSupportError, ProviderNotInitializeError,
|
|
ProviderQuotaExceededError)
|
|
from controllers.console.datasets.error import (ArchivedDocumentImmutableError, DocumentAlreadyFinishedError,
|
|
DocumentIndexingError, InvalidActionError, InvalidMetadataError)
|
|
from controllers.console.setup import setup_required
|
|
from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check
|
|
from core.errors.error import (LLMBadRequestError, ModelCurrentlyNotSupportError, ProviderTokenNotInitError,
|
|
QuotaExceededError)
|
|
from core.indexing_runner import IndexingRunner
|
|
from core.model_manager import ModelManager
|
|
from core.model_runtime.entities.model_entities import ModelType
|
|
from core.model_runtime.errors.invoke import InvokeAuthorizationError
|
|
from extensions.ext_database import db
|
|
from extensions.ext_redis import redis_client
|
|
from fields.document_fields import (dataset_and_document_fields, document_fields, document_status_fields,
|
|
document_with_segments_fields)
|
|
from flask import request
|
|
from flask_login import current_user
|
|
from flask_restful import Resource, fields, marshal, marshal_with, reqparse
|
|
from libs.login import login_required
|
|
from models.dataset import Dataset, DatasetProcessRule, Document, DocumentSegment
|
|
from models.model import UploadFile
|
|
from services.dataset_service import DatasetService, DocumentService
|
|
from sqlalchemy import asc, desc
|
|
from tasks.add_document_to_index_task import add_document_to_index_task
|
|
from tasks.remove_document_from_index_task import remove_document_from_index_task
|
|
from werkzeug.exceptions import Forbidden, NotFound
|
|
|
|
|
|
class DocumentResource(Resource):
|
|
def get_document(self, dataset_id: str, document_id: str) -> Document:
|
|
dataset = DatasetService.get_dataset(dataset_id)
|
|
if not dataset:
|
|
raise NotFound('Dataset not found.')
|
|
|
|
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:
|
|
raise NotFound('Document not found.')
|
|
|
|
if document.tenant_id != current_user.current_tenant_id:
|
|
raise Forbidden('No permission.')
|
|
|
|
return document
|
|
|
|
def get_batch_documents(self, dataset_id: str, batch: str) -> List[Document]:
|
|
dataset = DatasetService.get_dataset(dataset_id)
|
|
if not dataset:
|
|
raise NotFound('Dataset not found.')
|
|
|
|
try:
|
|
DatasetService.check_dataset_permission(dataset, current_user)
|
|
except services.errors.account.NoPermissionError as e:
|
|
raise Forbidden(str(e))
|
|
|
|
documents = DocumentService.get_batch_documents(dataset_id, batch)
|
|
|
|
if not documents:
|
|
raise NotFound('Documents not found.')
|
|
|
|
return documents
|
|
|
|
|
|
class GetProcessRuleApi(Resource):
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def get(self):
|
|
req_data = request.args
|
|
|
|
document_id = req_data.get('document_id')
|
|
|
|
# get default rules
|
|
mode = DocumentService.DEFAULT_RULES['mode']
|
|
rules = DocumentService.DEFAULT_RULES['rules']
|
|
if document_id:
|
|
# get the latest process rule
|
|
document = Document.query.get_or_404(document_id)
|
|
|
|
dataset = DatasetService.get_dataset(document.dataset_id)
|
|
|
|
if not dataset:
|
|
raise NotFound('Dataset not found.')
|
|
|
|
try:
|
|
DatasetService.check_dataset_permission(dataset, current_user)
|
|
except services.errors.account.NoPermissionError as e:
|
|
raise Forbidden(str(e))
|
|
|
|
# get the latest process rule
|
|
dataset_process_rule = db.session.query(DatasetProcessRule). \
|
|
filter(DatasetProcessRule.dataset_id == document.dataset_id). \
|
|
order_by(DatasetProcessRule.created_at.desc()). \
|
|
limit(1). \
|
|
one_or_none()
|
|
if dataset_process_rule:
|
|
mode = dataset_process_rule.mode
|
|
rules = dataset_process_rule.rules_dict
|
|
|
|
return {
|
|
'mode': mode,
|
|
'rules': rules
|
|
}
|
|
|
|
|
|
class DatasetDocumentListApi(Resource):
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def get(self, dataset_id):
|
|
dataset_id = str(dataset_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)
|
|
sort = request.args.get('sort', default='-created_at', type=str)
|
|
fetch = request.args.get('fetch', default=False, type=bool)
|
|
dataset = DatasetService.get_dataset(dataset_id)
|
|
if not dataset:
|
|
raise NotFound('Dataset not found.')
|
|
|
|
try:
|
|
DatasetService.check_dataset_permission(dataset, current_user)
|
|
except services.errors.account.NoPermissionError as e:
|
|
raise Forbidden(str(e))
|
|
|
|
query = Document.query.filter_by(
|
|
dataset_id=str(dataset_id), tenant_id=current_user.current_tenant_id)
|
|
|
|
if search:
|
|
search = f'%{search}%'
|
|
query = query.filter(Document.name.like(search))
|
|
|
|
if sort.startswith('-'):
|
|
sort_logic = desc
|
|
sort = sort[1:]
|
|
else:
|
|
sort_logic = asc
|
|
|
|
if sort == 'hit_count':
|
|
sub_query = db.select(DocumentSegment.document_id,
|
|
db.func.sum(DocumentSegment.hit_count).label("total_hit_count")) \
|
|
.group_by(DocumentSegment.document_id) \
|
|
.subquery()
|
|
|
|
query = query.outerjoin(sub_query, sub_query.c.document_id == Document.id) \
|
|
.order_by(sort_logic(db.func.coalesce(sub_query.c.total_hit_count, 0)))
|
|
elif sort == 'created_at':
|
|
query = query.order_by(sort_logic(Document.created_at))
|
|
else:
|
|
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
|
|
if fetch:
|
|
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
|
|
data = marshal(documents, document_with_segments_fields)
|
|
else:
|
|
data = marshal(documents, document_fields)
|
|
response = {
|
|
'data': data,
|
|
'has_more': len(documents) == limit,
|
|
'limit': limit,
|
|
'total': paginated_documents.total,
|
|
'page': page
|
|
}
|
|
|
|
return response
|
|
|
|
documents_and_batch_fields = {
|
|
'documents': fields.List(fields.Nested(document_fields)),
|
|
'batch': fields.String
|
|
}
|
|
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
@marshal_with(documents_and_batch_fields)
|
|
@cloud_edition_billing_resource_check('vector_space')
|
|
def post(self, dataset_id):
|
|
dataset_id = str(dataset_id)
|
|
|
|
dataset = DatasetService.get_dataset(dataset_id)
|
|
|
|
if not dataset:
|
|
raise NotFound('Dataset not found.')
|
|
|
|
# The role of the current user in the ta table must be admin or owner
|
|
if not current_user.is_admin_or_owner:
|
|
raise Forbidden()
|
|
|
|
try:
|
|
DatasetService.check_dataset_permission(dataset, current_user)
|
|
except services.errors.account.NoPermissionError as e:
|
|
raise Forbidden(str(e))
|
|
|
|
parser = reqparse.RequestParser()
|
|
parser.add_argument('indexing_technique', type=str, choices=Dataset.INDEXING_TECHNIQUE_LIST, nullable=False,
|
|
location='json')
|
|
parser.add_argument('data_source', type=dict, required=False, location='json')
|
|
parser.add_argument('process_rule', type=dict, required=False, location='json')
|
|
parser.add_argument('duplicate', type=bool, nullable=False, 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('retrieval_model', type=dict, required=False, nullable=False,
|
|
location='json')
|
|
args = parser.parse_args()
|
|
|
|
if not dataset.indexing_technique and not args['indexing_technique']:
|
|
raise ValueError('indexing_technique is required.')
|
|
|
|
# validate args
|
|
DocumentService.document_create_args_validate(args)
|
|
|
|
try:
|
|
documents, batch = DocumentService.save_document_with_dataset_id(dataset, args, current_user)
|
|
except ProviderTokenNotInitError as ex:
|
|
raise ProviderNotInitializeError(ex.description)
|
|
except QuotaExceededError:
|
|
raise ProviderQuotaExceededError()
|
|
except ModelCurrentlyNotSupportError:
|
|
raise ProviderModelCurrentlyNotSupportError()
|
|
|
|
return {
|
|
'documents': documents,
|
|
'batch': batch
|
|
}
|
|
|
|
|
|
class DatasetInitApi(Resource):
|
|
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
@marshal_with(dataset_and_document_fields)
|
|
@cloud_edition_billing_resource_check('vector_space')
|
|
def post(self):
|
|
# The role of the current user in the ta table must be admin or owner
|
|
if not current_user.is_admin_or_owner:
|
|
raise Forbidden()
|
|
|
|
parser = reqparse.RequestParser()
|
|
parser.add_argument('indexing_technique', type=str, choices=Dataset.INDEXING_TECHNIQUE_LIST, required=True,
|
|
nullable=False, location='json')
|
|
parser.add_argument('data_source', type=dict, required=True, nullable=True, location='json')
|
|
parser.add_argument('process_rule', type=dict, required=True, 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')
|
|
parser.add_argument('retrieval_model', type=dict, required=False, nullable=False,
|
|
location='json')
|
|
args = parser.parse_args()
|
|
if args['indexing_technique'] == 'high_quality':
|
|
try:
|
|
model_manager = ModelManager()
|
|
model_manager.get_default_model_instance(
|
|
tenant_id=current_user.current_tenant_id,
|
|
model_type=ModelType.TEXT_EMBEDDING
|
|
)
|
|
except InvokeAuthorizationError:
|
|
raise ProviderNotInitializeError(
|
|
f"No Embedding Model available. Please configure a valid provider "
|
|
f"in the Settings -> Model Provider.")
|
|
except ProviderTokenNotInitError as ex:
|
|
raise ProviderNotInitializeError(ex.description)
|
|
|
|
# validate args
|
|
DocumentService.document_create_args_validate(args)
|
|
|
|
try:
|
|
dataset, documents, batch = DocumentService.save_document_without_dataset_id(
|
|
tenant_id=current_user.current_tenant_id,
|
|
document_data=args,
|
|
account=current_user
|
|
)
|
|
except ProviderTokenNotInitError as ex:
|
|
raise ProviderNotInitializeError(ex.description)
|
|
except QuotaExceededError:
|
|
raise ProviderQuotaExceededError()
|
|
except ModelCurrentlyNotSupportError:
|
|
raise ProviderModelCurrentlyNotSupportError()
|
|
|
|
response = {
|
|
'dataset': dataset,
|
|
'documents': documents,
|
|
'batch': batch
|
|
}
|
|
|
|
return response
|
|
|
|
|
|
class DocumentIndexingEstimateApi(DocumentResource):
|
|
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def get(self, dataset_id, document_id):
|
|
dataset_id = str(dataset_id)
|
|
document_id = str(document_id)
|
|
document = self.get_document(dataset_id, document_id)
|
|
|
|
if document.indexing_status in ['completed', 'error']:
|
|
raise DocumentAlreadyFinishedError()
|
|
|
|
data_process_rule = document.dataset_process_rule
|
|
data_process_rule_dict = data_process_rule.to_dict()
|
|
|
|
response = {
|
|
"tokens": 0,
|
|
"total_price": 0,
|
|
"currency": "USD",
|
|
"total_segments": 0,
|
|
"preview": []
|
|
}
|
|
|
|
if document.data_source_type == 'upload_file':
|
|
data_source_info = document.data_source_info_dict
|
|
if data_source_info and 'upload_file_id' in data_source_info:
|
|
file_id = data_source_info['upload_file_id']
|
|
|
|
file = db.session.query(UploadFile).filter(
|
|
UploadFile.tenant_id == document.tenant_id,
|
|
UploadFile.id == file_id
|
|
).first()
|
|
|
|
# raise error if file not found
|
|
if not file:
|
|
raise NotFound('File not found.')
|
|
|
|
indexing_runner = IndexingRunner()
|
|
|
|
try:
|
|
response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, [file],
|
|
data_process_rule_dict, None,
|
|
'English', dataset_id)
|
|
except LLMBadRequestError:
|
|
raise ProviderNotInitializeError(
|
|
f"No Embedding Model available. Please configure a valid provider "
|
|
f"in the Settings -> Model Provider.")
|
|
except ProviderTokenNotInitError as ex:
|
|
raise ProviderNotInitializeError(ex.description)
|
|
|
|
return response
|
|
|
|
|
|
class DocumentBatchIndexingEstimateApi(DocumentResource):
|
|
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def get(self, dataset_id, batch):
|
|
dataset_id = str(dataset_id)
|
|
batch = str(batch)
|
|
dataset = DatasetService.get_dataset(dataset_id)
|
|
if dataset is None:
|
|
raise NotFound("Dataset not found.")
|
|
documents = self.get_batch_documents(dataset_id, batch)
|
|
response = {
|
|
"tokens": 0,
|
|
"total_price": 0,
|
|
"currency": "USD",
|
|
"total_segments": 0,
|
|
"preview": []
|
|
}
|
|
if not documents:
|
|
return response
|
|
data_process_rule = documents[0].dataset_process_rule
|
|
data_process_rule_dict = data_process_rule.to_dict()
|
|
info_list = []
|
|
for document in documents:
|
|
if document.indexing_status in ['completed', 'error']:
|
|
raise DocumentAlreadyFinishedError()
|
|
data_source_info = document.data_source_info_dict
|
|
# format document files info
|
|
if data_source_info and 'upload_file_id' in data_source_info:
|
|
file_id = data_source_info['upload_file_id']
|
|
info_list.append(file_id)
|
|
# format document notion info
|
|
elif data_source_info and 'notion_workspace_id' in data_source_info and 'notion_page_id' in data_source_info:
|
|
pages = []
|
|
page = {
|
|
'page_id': data_source_info['notion_page_id'],
|
|
'type': data_source_info['type']
|
|
}
|
|
pages.append(page)
|
|
notion_info = {
|
|
'workspace_id': data_source_info['notion_workspace_id'],
|
|
'pages': pages
|
|
}
|
|
info_list.append(notion_info)
|
|
|
|
if dataset.data_source_type == 'upload_file':
|
|
file_details = db.session.query(UploadFile).filter(
|
|
UploadFile.tenant_id == current_user.current_tenant_id,
|
|
UploadFile.id.in_(info_list)
|
|
).all()
|
|
|
|
if file_details is None:
|
|
raise NotFound("File not found.")
|
|
|
|
indexing_runner = IndexingRunner()
|
|
try:
|
|
response = indexing_runner.file_indexing_estimate(current_user.current_tenant_id, file_details,
|
|
data_process_rule_dict, None,
|
|
'English', dataset_id)
|
|
except LLMBadRequestError:
|
|
raise ProviderNotInitializeError(
|
|
f"No Embedding Model available. Please configure a valid provider "
|
|
f"in the Settings -> Model Provider.")
|
|
except ProviderTokenNotInitError as ex:
|
|
raise ProviderNotInitializeError(ex.description)
|
|
elif dataset.data_source_type == 'notion_import':
|
|
|
|
indexing_runner = IndexingRunner()
|
|
try:
|
|
response = indexing_runner.notion_indexing_estimate(current_user.current_tenant_id,
|
|
info_list,
|
|
data_process_rule_dict,
|
|
None, 'English', dataset_id)
|
|
except LLMBadRequestError:
|
|
raise ProviderNotInitializeError(
|
|
f"No Embedding Model available. Please configure a valid provider "
|
|
f"in the Settings -> Model Provider.")
|
|
except ProviderTokenNotInitError as ex:
|
|
raise ProviderNotInitializeError(ex.description)
|
|
else:
|
|
raise ValueError('Data source type not support')
|
|
return response
|
|
|
|
|
|
class DocumentBatchIndexingStatusApi(DocumentResource):
|
|
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def get(self, dataset_id, batch):
|
|
dataset_id = str(dataset_id)
|
|
batch = str(batch)
|
|
documents = self.get_batch_documents(dataset_id, batch)
|
|
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
|
|
|
|
|
|
class DocumentIndexingStatusApi(DocumentResource):
|
|
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def get(self, dataset_id, document_id):
|
|
dataset_id = str(dataset_id)
|
|
document_id = str(document_id)
|
|
document = self.get_document(dataset_id, document_id)
|
|
|
|
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'
|
|
return marshal(document, document_status_fields)
|
|
|
|
|
|
class DocumentDetailApi(DocumentResource):
|
|
METADATA_CHOICES = {'all', 'only', 'without'}
|
|
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def get(self, dataset_id, document_id):
|
|
dataset_id = str(dataset_id)
|
|
document_id = str(document_id)
|
|
document = self.get_document(dataset_id, document_id)
|
|
|
|
metadata = request.args.get('metadata', 'all')
|
|
if metadata not in self.METADATA_CHOICES:
|
|
raise InvalidMetadataError(f'Invalid metadata value: {metadata}')
|
|
|
|
if metadata == 'only':
|
|
response = {
|
|
'id': document.id,
|
|
'doc_type': document.doc_type,
|
|
'doc_metadata': document.doc_metadata
|
|
}
|
|
elif metadata == 'without':
|
|
process_rules = DatasetService.get_process_rules(dataset_id)
|
|
data_source_info = document.data_source_detail_dict
|
|
response = {
|
|
'id': document.id,
|
|
'position': document.position,
|
|
'data_source_type': document.data_source_type,
|
|
'data_source_info': data_source_info,
|
|
'dataset_process_rule_id': document.dataset_process_rule_id,
|
|
'dataset_process_rule': process_rules,
|
|
'name': document.name,
|
|
'created_from': document.created_from,
|
|
'created_by': document.created_by,
|
|
'created_at': document.created_at.timestamp(),
|
|
'tokens': document.tokens,
|
|
'indexing_status': document.indexing_status,
|
|
'completed_at': int(document.completed_at.timestamp()) if document.completed_at else None,
|
|
'updated_at': int(document.updated_at.timestamp()) if document.updated_at else None,
|
|
'indexing_latency': document.indexing_latency,
|
|
'error': document.error,
|
|
'enabled': document.enabled,
|
|
'disabled_at': int(document.disabled_at.timestamp()) if document.disabled_at else None,
|
|
'disabled_by': document.disabled_by,
|
|
'archived': document.archived,
|
|
'segment_count': document.segment_count,
|
|
'average_segment_length': document.average_segment_length,
|
|
'hit_count': document.hit_count,
|
|
'display_status': document.display_status,
|
|
'doc_form': document.doc_form
|
|
}
|
|
else:
|
|
process_rules = DatasetService.get_process_rules(dataset_id)
|
|
data_source_info = document.data_source_detail_dict
|
|
response = {
|
|
'id': document.id,
|
|
'position': document.position,
|
|
'data_source_type': document.data_source_type,
|
|
'data_source_info': data_source_info,
|
|
'dataset_process_rule_id': document.dataset_process_rule_id,
|
|
'dataset_process_rule': process_rules,
|
|
'name': document.name,
|
|
'created_from': document.created_from,
|
|
'created_by': document.created_by,
|
|
'created_at': document.created_at.timestamp(),
|
|
'tokens': document.tokens,
|
|
'indexing_status': document.indexing_status,
|
|
'completed_at': int(document.completed_at.timestamp()) if document.completed_at else None,
|
|
'updated_at': int(document.updated_at.timestamp()) if document.updated_at else None,
|
|
'indexing_latency': document.indexing_latency,
|
|
'error': document.error,
|
|
'enabled': document.enabled,
|
|
'disabled_at': int(document.disabled_at.timestamp()) if document.disabled_at else None,
|
|
'disabled_by': document.disabled_by,
|
|
'archived': document.archived,
|
|
'doc_type': document.doc_type,
|
|
'doc_metadata': document.doc_metadata,
|
|
'segment_count': document.segment_count,
|
|
'average_segment_length': document.average_segment_length,
|
|
'hit_count': document.hit_count,
|
|
'display_status': document.display_status,
|
|
'doc_form': document.doc_form
|
|
}
|
|
|
|
return response, 200
|
|
|
|
|
|
class DocumentProcessingApi(DocumentResource):
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def patch(self, dataset_id, document_id, action):
|
|
dataset_id = str(dataset_id)
|
|
document_id = str(document_id)
|
|
document = self.get_document(dataset_id, document_id)
|
|
|
|
# The role of the current user in the ta table must be admin or owner
|
|
if not current_user.is_admin_or_owner:
|
|
raise Forbidden()
|
|
|
|
if action == "pause":
|
|
if document.indexing_status != "indexing":
|
|
raise InvalidActionError('Document not in indexing state.')
|
|
|
|
document.paused_by = current_user.id
|
|
document.paused_at = datetime.utcnow()
|
|
document.is_paused = True
|
|
db.session.commit()
|
|
|
|
elif action == "resume":
|
|
if document.indexing_status not in ["paused", "error"]:
|
|
raise InvalidActionError('Document not in paused or error state.')
|
|
|
|
document.paused_by = None
|
|
document.paused_at = None
|
|
document.is_paused = False
|
|
db.session.commit()
|
|
else:
|
|
raise InvalidActionError()
|
|
|
|
return {'result': 'success'}, 200
|
|
|
|
|
|
class DocumentDeleteApi(DocumentResource):
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def delete(self, dataset_id, document_id):
|
|
dataset_id = str(dataset_id)
|
|
document_id = str(document_id)
|
|
dataset = DatasetService.get_dataset(dataset_id)
|
|
if dataset is None:
|
|
raise NotFound("Dataset not found.")
|
|
# check user's model setting
|
|
DatasetService.check_dataset_model_setting(dataset)
|
|
|
|
document = self.get_document(dataset_id, document_id)
|
|
|
|
try:
|
|
DocumentService.delete_document(document)
|
|
except services.errors.document.DocumentIndexingError:
|
|
raise DocumentIndexingError('Cannot delete document during indexing.')
|
|
|
|
return {'result': 'success'}, 204
|
|
|
|
|
|
class DocumentMetadataApi(DocumentResource):
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def put(self, dataset_id, document_id):
|
|
dataset_id = str(dataset_id)
|
|
document_id = str(document_id)
|
|
document = self.get_document(dataset_id, document_id)
|
|
|
|
req_data = request.get_json()
|
|
|
|
doc_type = req_data.get('doc_type')
|
|
doc_metadata = req_data.get('doc_metadata')
|
|
|
|
# The role of the current user in the ta table must be admin or owner
|
|
if not current_user.is_admin_or_owner:
|
|
raise Forbidden()
|
|
|
|
if doc_type is None or doc_metadata is None:
|
|
raise ValueError('Both doc_type and doc_metadata must be provided.')
|
|
|
|
if doc_type not in DocumentService.DOCUMENT_METADATA_SCHEMA:
|
|
raise ValueError('Invalid doc_type.')
|
|
|
|
if not isinstance(doc_metadata, dict):
|
|
raise ValueError('doc_metadata must be a dictionary.')
|
|
|
|
metadata_schema = DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type]
|
|
|
|
document.doc_metadata = {}
|
|
if doc_type == 'others':
|
|
document.doc_metadata = doc_metadata
|
|
else:
|
|
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
|
|
|
|
document.doc_type = doc_type
|
|
document.updated_at = datetime.utcnow()
|
|
db.session.commit()
|
|
|
|
return {'result': 'success', 'message': 'Document metadata updated.'}, 200
|
|
|
|
|
|
class DocumentStatusApi(DocumentResource):
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
@cloud_edition_billing_resource_check('vector_space')
|
|
def patch(self, dataset_id, document_id, action):
|
|
dataset_id = str(dataset_id)
|
|
document_id = str(document_id)
|
|
dataset = DatasetService.get_dataset(dataset_id)
|
|
if dataset is None:
|
|
raise NotFound("Dataset not found.")
|
|
# check user's model setting
|
|
DatasetService.check_dataset_model_setting(dataset)
|
|
|
|
document = self.get_document(dataset_id, document_id)
|
|
|
|
# The role of the current user in the ta table must be admin or owner
|
|
if not current_user.is_admin_or_owner:
|
|
raise Forbidden()
|
|
|
|
indexing_cache_key = 'document_{}_indexing'.format(document.id)
|
|
cache_result = redis_client.get(indexing_cache_key)
|
|
if cache_result is not None:
|
|
raise InvalidActionError("Document is being indexed, please try again later")
|
|
|
|
if action == "enable":
|
|
if document.enabled:
|
|
raise InvalidActionError('Document already enabled.')
|
|
|
|
document.enabled = True
|
|
document.disabled_at = None
|
|
document.disabled_by = None
|
|
document.updated_at = datetime.utcnow()
|
|
db.session.commit()
|
|
|
|
# Set cache to prevent indexing the same document multiple times
|
|
redis_client.setex(indexing_cache_key, 600, 1)
|
|
|
|
add_document_to_index_task.delay(document_id)
|
|
|
|
return {'result': 'success'}, 200
|
|
|
|
elif action == "disable":
|
|
if not document.completed_at or document.indexing_status != 'completed':
|
|
raise InvalidActionError('Document is not completed.')
|
|
if not document.enabled:
|
|
raise InvalidActionError('Document already disabled.')
|
|
|
|
document.enabled = False
|
|
document.disabled_at = datetime.utcnow()
|
|
document.disabled_by = current_user.id
|
|
document.updated_at = datetime.utcnow()
|
|
db.session.commit()
|
|
|
|
# Set cache to prevent indexing the same document multiple times
|
|
redis_client.setex(indexing_cache_key, 600, 1)
|
|
|
|
remove_document_from_index_task.delay(document_id)
|
|
|
|
return {'result': 'success'}, 200
|
|
|
|
elif action == "archive":
|
|
if document.archived:
|
|
raise InvalidActionError('Document already archived.')
|
|
|
|
document.archived = True
|
|
document.archived_at = datetime.utcnow()
|
|
document.archived_by = current_user.id
|
|
document.updated_at = datetime.utcnow()
|
|
db.session.commit()
|
|
|
|
if document.enabled:
|
|
# Set cache to prevent indexing the same document multiple times
|
|
redis_client.setex(indexing_cache_key, 600, 1)
|
|
|
|
remove_document_from_index_task.delay(document_id)
|
|
|
|
return {'result': 'success'}, 200
|
|
elif action == "un_archive":
|
|
if not document.archived:
|
|
raise InvalidActionError('Document is not archived.')
|
|
|
|
document.archived = False
|
|
document.archived_at = None
|
|
document.archived_by = None
|
|
document.updated_at = datetime.utcnow()
|
|
db.session.commit()
|
|
|
|
# Set cache to prevent indexing the same document multiple times
|
|
redis_client.setex(indexing_cache_key, 600, 1)
|
|
|
|
add_document_to_index_task.delay(document_id)
|
|
|
|
return {'result': 'success'}, 200
|
|
else:
|
|
raise InvalidActionError()
|
|
|
|
|
|
class DocumentPauseApi(DocumentResource):
|
|
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def patch(self, dataset_id, document_id):
|
|
"""pause document."""
|
|
dataset_id = str(dataset_id)
|
|
document_id = str(document_id)
|
|
|
|
dataset = DatasetService.get_dataset(dataset_id)
|
|
if not dataset:
|
|
raise NotFound('Dataset not found.')
|
|
|
|
document = DocumentService.get_document(dataset.id, document_id)
|
|
|
|
# 404 if document not found
|
|
if document is None:
|
|
raise NotFound("Document Not Exists.")
|
|
|
|
# 403 if document is archived
|
|
if DocumentService.check_archived(document):
|
|
raise ArchivedDocumentImmutableError()
|
|
|
|
try:
|
|
# pause document
|
|
DocumentService.pause_document(document)
|
|
except services.errors.document.DocumentIndexingError:
|
|
raise DocumentIndexingError('Cannot pause completed document.')
|
|
|
|
return {'result': 'success'}, 204
|
|
|
|
|
|
class DocumentRecoverApi(DocumentResource):
|
|
@setup_required
|
|
@login_required
|
|
@account_initialization_required
|
|
def patch(self, dataset_id, document_id):
|
|
"""recover document."""
|
|
dataset_id = str(dataset_id)
|
|
document_id = str(document_id)
|
|
dataset = DatasetService.get_dataset(dataset_id)
|
|
if not dataset:
|
|
raise NotFound('Dataset not found.')
|
|
document = DocumentService.get_document(dataset.id, document_id)
|
|
|
|
# 404 if document not found
|
|
if document is None:
|
|
raise NotFound("Document Not Exists.")
|
|
|
|
# 403 if document is archived
|
|
if DocumentService.check_archived(document):
|
|
raise ArchivedDocumentImmutableError()
|
|
try:
|
|
# pause document
|
|
DocumentService.recover_document(document)
|
|
except services.errors.document.DocumentIndexingError:
|
|
raise DocumentIndexingError('Document is not in paused status.')
|
|
|
|
return {'result': 'success'}, 204
|
|
|
|
|
|
api.add_resource(GetProcessRuleApi, '/datasets/process-rule')
|
|
api.add_resource(DatasetDocumentListApi,
|
|
'/datasets/<uuid:dataset_id>/documents')
|
|
api.add_resource(DatasetInitApi,
|
|
'/datasets/init')
|
|
api.add_resource(DocumentIndexingEstimateApi,
|
|
'/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-estimate')
|
|
api.add_resource(DocumentBatchIndexingEstimateApi,
|
|
'/datasets/<uuid:dataset_id>/batch/<string:batch>/indexing-estimate')
|
|
api.add_resource(DocumentBatchIndexingStatusApi,
|
|
'/datasets/<uuid:dataset_id>/batch/<string:batch>/indexing-status')
|
|
api.add_resource(DocumentIndexingStatusApi,
|
|
'/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/indexing-status')
|
|
api.add_resource(DocumentDetailApi,
|
|
'/datasets/<uuid:dataset_id>/documents/<uuid:document_id>')
|
|
api.add_resource(DocumentProcessingApi,
|
|
'/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/<string:action>')
|
|
api.add_resource(DocumentDeleteApi,
|
|
'/datasets/<uuid:dataset_id>/documents/<uuid:document_id>')
|
|
api.add_resource(DocumentMetadataApi,
|
|
'/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/metadata')
|
|
api.add_resource(DocumentStatusApi,
|
|
'/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/status/<string:action>')
|
|
api.add_resource(DocumentPauseApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/pause')
|
|
api.add_resource(DocumentRecoverApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/resume')
|