mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-05-06 10:17:49 +08:00
Refactor: migrate document run api (#14351)
### What problem does this PR solve? Before migration: POST /v1/document/run After migration: POST /api/v1/documents/ingest/ ### Type of change - [x] Refactoring
This commit is contained in:
Jack
GitHub
@ -20,11 +20,8 @@ from quart import make_response, request
|
||||
from api.apps import current_user, login_required
|
||||
from api.constants import IMG_BASE64_PREFIX
|
||||
from api.db import FileType
|
||||
from api.db.db_models import Task
|
||||
from api.db.services.document_service import DocumentService
|
||||
from api.db.services.file2document_service import File2DocumentService
|
||||
from api.db.services.knowledgebase_service import KnowledgebaseService
|
||||
from api.db.services.task_service import TaskService, cancel_all_task_of
|
||||
from api.utils.api_utils import (
|
||||
get_data_error_result,
|
||||
get_json_result,
|
||||
@ -58,69 +55,6 @@ def thumbnails():
|
||||
return server_error_response(e)
|
||||
|
||||
|
||||
@manager.route("/run", methods=["POST"]) # noqa: F821
|
||||
@login_required
|
||||
@validate_request("doc_ids", "run")
|
||||
async def run():
|
||||
req = await get_request_json()
|
||||
uid = current_user.id
|
||||
try:
|
||||
|
||||
def _run_sync():
|
||||
for doc_id in req["doc_ids"]:
|
||||
if not DocumentService.accessible(doc_id, uid):
|
||||
return get_json_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
|
||||
|
||||
kb_table_num_map = {}
|
||||
for id in req["doc_ids"]:
|
||||
info = {"run": str(req["run"]), "progress": 0}
|
||||
if str(req["run"]) == TaskStatus.RUNNING.value and req.get("delete", False):
|
||||
info["progress_msg"] = ""
|
||||
info["chunk_num"] = 0
|
||||
info["token_num"] = 0
|
||||
|
||||
tenant_id = DocumentService.get_tenant_id(id)
|
||||
if not tenant_id:
|
||||
return get_data_error_result(message="Tenant not found!")
|
||||
e, doc = DocumentService.get_by_id(id)
|
||||
if not e:
|
||||
return get_data_error_result(message="Document not found!")
|
||||
|
||||
if str(req["run"]) == TaskStatus.CANCEL.value:
|
||||
tasks = list(TaskService.query(doc_id=id))
|
||||
has_unfinished_task = any((task.progress or 0) < 1 for task in tasks)
|
||||
if str(doc.run) in [TaskStatus.RUNNING.value, TaskStatus.CANCEL.value] or has_unfinished_task:
|
||||
cancel_all_task_of(id)
|
||||
else:
|
||||
return get_data_error_result(message="Cannot cancel a task that is not in RUNNING status")
|
||||
if all([("delete" not in req or req["delete"]), str(req["run"]) == TaskStatus.RUNNING.value, str(doc.run) == TaskStatus.DONE.value]):
|
||||
DocumentService.clear_chunk_num_when_rerun(doc.id)
|
||||
|
||||
DocumentService.update_by_id(id, info)
|
||||
if req.get("delete", False):
|
||||
TaskService.filter_delete([Task.doc_id == id])
|
||||
if settings.docStoreConn.index_exist(search.index_name(tenant_id), doc.kb_id):
|
||||
settings.docStoreConn.delete({"doc_id": id}, search.index_name(tenant_id), doc.kb_id)
|
||||
|
||||
if str(req["run"]) == TaskStatus.RUNNING.value:
|
||||
if req.get("apply_kb"):
|
||||
e, kb = KnowledgebaseService.get_by_id(doc.kb_id)
|
||||
if not e:
|
||||
raise LookupError("Can't find this dataset!")
|
||||
doc.parser_config["llm_id"] = kb.parser_config.get("llm_id")
|
||||
doc.parser_config["enable_metadata"] = kb.parser_config.get("enable_metadata", False)
|
||||
doc.parser_config["metadata"] = kb.parser_config.get("metadata", {})
|
||||
DocumentService.update_parser_config(doc.id, doc.parser_config)
|
||||
doc_dict = doc.to_dict()
|
||||
DocumentService.run(tenant_id, doc_dict, kb_table_num_map)
|
||||
|
||||
return get_json_result(data=True)
|
||||
|
||||
return await thread_pool_exec(_run_sync)
|
||||
except Exception as e:
|
||||
return server_error_response(e)
|
||||
|
||||
|
||||
@manager.route("/get/<doc_id>", methods=["GET"]) # noqa: F821
|
||||
@login_required
|
||||
async def get(doc_id):
|
||||
|
||||
@ -34,13 +34,14 @@ from api.db.db_models import Task
|
||||
from api.db.services.document_service import DocumentService
|
||||
from api.db.services.file_service import FileService
|
||||
from api.db.services.knowledgebase_service import KnowledgebaseService
|
||||
from api.db.services.task_service import TaskService, cancel_all_task_of
|
||||
from api.common.check_team_permission import check_kb_team_permission
|
||||
from api.db.services.task_service import TaskService, cancel_all_task_of
|
||||
from api.utils.api_utils import get_data_error_result, get_error_data_result, get_result, get_json_result, \
|
||||
server_error_response, add_tenant_id_to_kwargs, get_request_json, get_error_argument_result, check_duplicate_ids
|
||||
from api.utils.validation_utils import (
|
||||
UpdateDocumentReq, format_validation_error_message, validate_and_parse_json_request, DeleteDocumentReq,
|
||||
)
|
||||
|
||||
from common import settings
|
||||
from common.constants import ParserType, RetCode, TaskStatus, SANDBOX_ARTIFACT_BUCKET
|
||||
from common.metadata_utils import convert_conditions, meta_filter, turn2jsonschema
|
||||
@ -1295,6 +1296,77 @@ async def update_metadata(tenant_id, dataset_id):
|
||||
return get_result(data={"updated": updated, "matched_docs": len(target_doc_ids)})
|
||||
|
||||
|
||||
@manager.route("/documents/ingest", methods=["POST"]) # noqa: F821
|
||||
@login_required
|
||||
@add_tenant_id_to_kwargs
|
||||
async def ingest(tenant_id):
|
||||
req = await get_request_json()
|
||||
try:
|
||||
user_id = tenant_id
|
||||
|
||||
error_code, error_message = await thread_pool_exec(_run_sync, user_id, req)
|
||||
|
||||
if error_code:
|
||||
logging.error(f"error when ingest documents:{req}, error message:{error_message}")
|
||||
return get_json_result(error_code, error_message)
|
||||
|
||||
return get_json_result(data=True)
|
||||
except Exception as e:
|
||||
logging.exception("document ingest/run failed")
|
||||
return server_error_response(e)
|
||||
|
||||
def _run_sync(user_id:str, req):
|
||||
for doc_id in req["doc_ids"]:
|
||||
if not DocumentService.accessible(doc_id, user_id):
|
||||
return RetCode.AUTHENTICATION_ERROR, "No authorization."
|
||||
|
||||
kb_table_num_map = {}
|
||||
for doc_id in req["doc_ids"]:
|
||||
info = {"run": str(req["run"]), "progress": 0}
|
||||
rerun_with_delete = str(req["run"]) == TaskStatus.RUNNING.value and req.get("delete", False)
|
||||
if rerun_with_delete:
|
||||
info["progress_msg"] = ""
|
||||
info["chunk_num"] = 0
|
||||
info["token_num"] = 0
|
||||
|
||||
doc_tenant_id = DocumentService.get_tenant_id(doc_id)
|
||||
if not doc_tenant_id:
|
||||
return RetCode.DATA_ERROR, "Tenant not found!"
|
||||
e, doc = DocumentService.get_by_id(doc_id)
|
||||
if not e:
|
||||
return RetCode.DATA_ERROR, "Document not found!"
|
||||
|
||||
if str(req["run"]) == TaskStatus.CANCEL.value:
|
||||
tasks = list(TaskService.query(doc_id=doc_id))
|
||||
has_unfinished_task = any((task.progress or 0) < 1 for task in tasks)
|
||||
if str(doc.run) in [TaskStatus.RUNNING.value, TaskStatus.CANCEL.value] or has_unfinished_task:
|
||||
cancel_all_task_of(doc_id)
|
||||
else:
|
||||
return RetCode.DATA_ERROR, "Cannot cancel a task that is not in RUNNING status"
|
||||
if all([rerun_with_delete, str(doc.run) == TaskStatus.DONE.value]):
|
||||
DocumentService.clear_chunk_num_when_rerun(doc_id)
|
||||
|
||||
DocumentService.update_by_id(doc_id, info)
|
||||
if req.get("delete", False):
|
||||
TaskService.filter_delete([Task.doc_id == doc_id])
|
||||
if settings.docStoreConn.index_exist(search.index_name(doc_tenant_id), doc.kb_id):
|
||||
settings.docStoreConn.delete({"doc_id": doc_id}, search.index_name(doc_tenant_id), doc.kb_id)
|
||||
|
||||
if str(req["run"]) == TaskStatus.RUNNING.value:
|
||||
if req.get("apply_kb"):
|
||||
e, kb = KnowledgebaseService.get_by_id(doc.kb_id)
|
||||
if not e:
|
||||
raise LookupError("Can't find this dataset!")
|
||||
doc.parser_config["llm_id"] = kb.parser_config.get("llm_id")
|
||||
doc.parser_config["enable_metadata"] = kb.parser_config.get("enable_metadata", False)
|
||||
doc.parser_config["metadata"] = kb.parser_config.get("metadata", {})
|
||||
DocumentService.update_parser_config(doc.id, doc.parser_config)
|
||||
doc_dict = doc.to_dict()
|
||||
DocumentService.run(doc_tenant_id, doc_dict, kb_table_num_map)
|
||||
|
||||
return None, None
|
||||
|
||||
|
||||
@manager.route("/datasets/<dataset_id>/documents/parse", methods=["POST"]) # noqa: F821
|
||||
@login_required
|
||||
@add_tenant_id_to_kwargs
|
||||
|
||||
Reference in New Issue
Block a user