Consolidateion of document upload API (#14106)

### What problem does this PR solve?

Consolidation WEB API & HTTP API for document upload

Before consolidation
Web API: POST /v1/document/upload
Http API - POST /api/v1/datasets/<dataset_id>/documents

After consolidation, Restful API -- POST
/api/v1/datasets/<dataset_id>/documents

### Type of change

- [x] Refactoring

api/apps/document_app.py

@@ -62,56 +62,6 @@ def _is_safe_download_filename(name: str) -> bool:
     return True
 
 
-@manager.route("/upload", methods=["POST"])  # noqa: F821
-@login_required
-@validate_request("kb_id")
-async def upload():
-    form = await request.form
-    kb_id = form.get("kb_id")
-    if not kb_id:
-        return get_json_result(data=False, message='Lack of "KB ID"', code=RetCode.ARGUMENT_ERROR)
-    files = await request.files
-    if "file" not in files:
-        return get_json_result(data=False, message="No file part!", code=RetCode.ARGUMENT_ERROR)
-
-    file_objs = files.getlist("file")
-
-    def _close_file_objs(objs):
-        for obj in objs:
-            try:
-                obj.close()
-            except Exception:
-                try:
-                    obj.stream.close()
-                except Exception:
-                    pass
-
-    for file_obj in file_objs:
-        if file_obj.filename == "":
-            _close_file_objs(file_objs)
-            return get_json_result(data=False, message="No file selected!", code=RetCode.ARGUMENT_ERROR)
-        if len(file_obj.filename.encode("utf-8")) > FILE_NAME_LEN_LIMIT:
-            _close_file_objs(file_objs)
-            return get_json_result(data=False, message=f"File name must be {FILE_NAME_LEN_LIMIT} bytes or less.", code=RetCode.ARGUMENT_ERROR)
-
-    e, kb = KnowledgebaseService.get_by_id(kb_id)
-    if not e:
-        raise LookupError("Can't find this dataset!")
-    if not check_kb_team_permission(kb, current_user.id):
-        return get_json_result(data=False, message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
-
-    err, files = await thread_pool_exec(FileService.upload_document, kb, file_objs, current_user.id)
-    if err:
-        files = [f[0] for f in files] if files else []
-        return get_json_result(data=files, message="\n".join(err), code=RetCode.SERVER_ERROR)
-
-    if not files:
-        return get_json_result(data=files, message="There seems to be an issue with your file format. Please verify it is correct and not corrupted.", code=RetCode.DATA_ERROR)
-    files = [f[0] for f in files]  # remove the blob
-
-    return get_json_result(data=files)
-
-
 @manager.route("/web_crawl", methods=["POST"])  # noqa: F821
 @login_required
 @validate_request("kb_id", "name", "url")

api/apps/restful_apis/document_api.py

@@ -19,8 +19,8 @@ from quart import request
 from peewee import OperationalError
 from pydantic import ValidationError
 
-from api.apps.services.document_api_service import rename_doc_key, validate_document_update_fields, \
-    update_document_name_only, update_chunk_method_only, update_document_status_only
+from api.apps.services.document_api_service import map_doc_keys_with_run_status, validate_document_update_fields, \
+    update_document_name_only, update_chunk_method_only, update_document_status_only, map_doc_keys
 from api.db.services.doc_metadata_service import DocMetadataService
 from api.db.services.document_service import DocumentService
 from api.db.services.knowledgebase_service import KnowledgebaseService
@@ -143,7 +143,7 @@ async def update_document(tenant_id, dataset_id, document_id):
     except OperationalError as e:
         logging.exception(e)
         return get_error_data_result(message="Database operation failed")
-    renamed_doc = rename_doc_key(doc)
+    renamed_doc = map_doc_keys(doc)
     return get_result(data=renamed_doc)
 
 
@@ -183,3 +183,136 @@ async def metadata_summary(dataset_id, tenant_id):
         return get_result(data={"summary": summary})
     except Exception as e:
         return server_error_response(e)
+
+
+@manager.route("/datasets/<dataset_id>/documents", methods=["POST"])  # noqa: F821
+@login_required
+@add_tenant_id_to_kwargs
+async def upload_document(dataset_id, tenant_id):
+    """
+    Upload documents to a dataset.
+    ---
+    tags:
+      - Documents
+    security:
+      - ApiKeyAuth: []
+    parameters:
+      - in: path
+        name: dataset_id
+        type: string
+        required: true
+        description: ID of the dataset.
+      - in: header
+        name: Authorization
+        type: string
+        required: true
+        description: Bearer token for authentication.
+      - in: formData
+        name: file
+        type: file
+        required: true
+        description: Document files to upload.
+      - in: formData
+        name: parent_path
+        type: string
+        description: Optional nested path under the parent folder. Uses '/' separators.
+      - in: query
+        name: return_raw_files
+        type: boolean
+        required: false
+        default: false
+        description: Whether to skip document key mapping and return raw document data
+    responses:
+      200:
+        description: Successfully uploaded documents.
+        schema:
+          type: object
+          properties:
+            data:
+              type: array
+              items:
+                type: object
+                properties:
+                  id:
+                    type: string
+                    description: Document ID.
+                  name:
+                    type: string
+                    description: Document name.
+                  chunk_count:
+                    type: integer
+                    description: Number of chunks.
+                  token_count:
+                    type: integer
+                    description: Number of tokens.
+                  dataset_id:
+                    type: string
+                    description: ID of the dataset.
+                  chunk_method:
+                    type: string
+                    description: Chunking method used.
+                  run:
+                    type: string
+                    description: Processing status.
+    """
+    from api.constants import FILE_NAME_LEN_LIMIT
+    from api.common.check_team_permission import check_kb_team_permission
+    from api.db.services.file_service import FileService
+    from common.misc_utils import thread_pool_exec
+    
+    form = await request.form
+    files = await request.files
+    
+    # Validation
+    if "file" not in files:
+        logging.error("No file part!")
+        return get_error_data_result(message="No file part!", code=RetCode.ARGUMENT_ERROR)
+    
+    file_objs = files.getlist("file")
+    for file_obj in file_objs:
+        if file_obj is None or file_obj.filename is None or file_obj.filename == "":
+            logging.error("No file selected!")
+            return get_error_data_result(message="No file selected!", code=RetCode.ARGUMENT_ERROR)
+        if len(file_obj.filename.encode("utf-8")) > FILE_NAME_LEN_LIMIT:
+            msg = f"File name must be {FILE_NAME_LEN_LIMIT} bytes or less."
+            logging.error(msg)
+            return get_error_data_result(message=msg, code=RetCode.ARGUMENT_ERROR)
+
+    # KB Lookup
+    e, kb = KnowledgebaseService.get_by_id(dataset_id)
+    if not e:
+        logging.error(f"Can't find the dataset with ID {dataset_id}!")
+        return get_error_data_result(message=f"Can't find the dataset with ID {dataset_id}!", code=RetCode.DATA_ERROR)
+    
+    # Permission Check
+    if not check_kb_team_permission(kb, tenant_id):
+        logging.error("No authorization.")
+        return get_error_data_result(message="No authorization.", code=RetCode.AUTHENTICATION_ERROR)
+
+    # File Upload (async)
+    err, files = await thread_pool_exec(
+        FileService.upload_document, kb, file_objs, tenant_id,
+        parent_path=form.get("parent_path")
+    )
+    if err:
+        msg = "\n".join(err)
+        logging.error(msg)
+        return get_error_data_result(message=msg, code=RetCode.SERVER_ERROR)
+
+    if not files:
+        msg = "There seems to be an issue with your file format. please verify it is correct and not corrupted."
+        logging.error(msg)
+        return get_error_data_result(message=msg, code=RetCode.DATA_ERROR)
+    
+    files = [f[0] for f in files]  # remove the blob
+
+    # Check if we should return raw files without document key mapping
+    return_raw_files = request.args.get("return_raw_files", "false").lower() == "true"
+
+    if return_raw_files:
+        return get_result(data=files)
+
+    renamed_doc_list = [map_doc_keys_with_run_status(doc, run_status="0") for doc in files]
+    return get_result(data=renamed_doc_list)
+
+

api/apps/sdk/doc.py

@@ -22,7 +22,6 @@ import xxhash
 from pydantic import BaseModel, Field, validator
 from quart import request, send_file
 
-from api.constants import FILE_NAME_LEN_LIMIT
 from api.db.db_models import APIToken, Document, File, Task
 from api.db.joint_services.tenant_model_service import get_model_config_by_id, get_model_config_by_type_and_name, get_tenant_default_model_by_type
 from api.db.services.doc_metadata_service import DocMetadataService
@@ -69,117 +68,6 @@ class Chunk(BaseModel):
         return value
 
 
-@manager.route("/datasets/<dataset_id>/documents", methods=["POST"])  # noqa: F821
-@token_required
-async def upload(dataset_id, tenant_id):
-    """
-    Upload documents to a dataset.
-    ---
-    tags:
-      - Documents
-    security:
-      - ApiKeyAuth: []
-    parameters:
-      - in: path
-        name: dataset_id
-        type: string
-        required: true
-        description: ID of the dataset.
-      - in: header
-        name: Authorization
-        type: string
-        required: true
-        description: Bearer token for authentication.
-      - in: formData
-        name: file
-        type: file
-        required: true
-        description: Document files to upload.
-      - in: formData
-        name: parent_path
-        type: string
-        description: Optional nested path under the parent folder. Uses '/' separators.
-    responses:
-      200:
-        description: Successfully uploaded documents.
-        schema:
-          type: object
-          properties:
-            data:
-              type: array
-              items:
-                type: object
-                properties:
-                  id:
-                    type: string
-                    description: Document ID.
-                  name:
-                    type: string
-                    description: Document name.
-                  chunk_count:
-                    type: integer
-                    description: Number of chunks.
-                  token_count:
-                    type: integer
-                    description: Number of tokens.
-                  dataset_id:
-                    type: string
-                    description: ID of the dataset.
-                  chunk_method:
-                    type: string
-                    description: Chunking method used.
-                  run:
-                    type: string
-                    description: Processing status.
-    """
-    form = await request.form
-    files = await request.files
-    if "file" not in files:
-        return get_error_data_result(message="No file part!", code=RetCode.ARGUMENT_ERROR)
-    file_objs = files.getlist("file")
-    for file_obj in file_objs:
-        if file_obj.filename == "":
-            return get_result(message="No file selected!", code=RetCode.ARGUMENT_ERROR)
-        if len(file_obj.filename.encode("utf-8")) > FILE_NAME_LEN_LIMIT:
-            return get_result(message=f"File name must be {FILE_NAME_LEN_LIMIT} bytes or less.", code=RetCode.ARGUMENT_ERROR)
-    """
-    # total size
-    total_size = 0
-    for file_obj in file_objs:
-        file_obj.seek(0, os.SEEK_END)
-        total_size += file_obj.tell()
-        file_obj.seek(0)
-    MAX_TOTAL_FILE_SIZE = 10 * 1024 * 1024
-    if total_size > MAX_TOTAL_FILE_SIZE:
-        return get_result(
-            message=f"Total file size exceeds 10MB limit! ({total_size / (1024 * 1024):.2f} MB)",
-            code=RetCode.ARGUMENT_ERROR,
-        )
-    """
-    e, kb = KnowledgebaseService.get_by_id(dataset_id)
-    if not e:
-        return server_error_response(LookupError(f"Can't find the dataset with ID {dataset_id}!"))
-    err, files = FileService.upload_document(kb, file_objs, tenant_id, parent_path=form.get("parent_path"))
-    if err:
-        return get_result(message="\n".join(err), code=RetCode.SERVER_ERROR)
-    # rename key's name
-    renamed_doc_list = []
-    for file in files:
-        doc = file[0]
-        key_mapping = {
-            "chunk_num": "chunk_count",
-            "kb_id": "dataset_id",
-            "token_num": "token_count",
-            "parser_id": "chunk_method",
-        }
-        renamed_doc = {}
-        for key, value in doc.items():
-            new_key = key_mapping.get(key, key)
-            renamed_doc[new_key] = value
-        renamed_doc["run"] = "UNSTART"
-        renamed_doc_list.append(renamed_doc)
-    return get_result(data=renamed_doc_list)
-
 @manager.route("/datasets/<dataset_id>/documents/<document_id>", methods=["GET"])  # noqa: F821
 @token_required
 async def download(tenant_id, dataset_id, document_id):

api/apps/services/document_api_service.py

@@ -165,7 +165,7 @@ def validate_document_update_fields(update_doc_req:UpdateDocumentReq, doc, req):
 
     return None, None
 
-def rename_doc_key(doc):
+def map_doc_keys(doc):
     """
     Rename document keys to match API response format.
 
@@ -175,6 +175,46 @@ def rename_doc_key(doc):
     Args:
         doc: The document model from the database.
 
+    Returns:
+        A dictionary with renamed keys for API response.
+    """
+    renamed_doc = _process_key_mappings(doc)
+    if "run" in renamed_doc.keys():
+        renamed_doc = _process_run_mapping(renamed_doc, renamed_doc["run"])
+    return renamed_doc
+
+
+def map_doc_keys_with_run_status(doc, run_status):
+    """
+    Map document keys to match API response format.
+
+    Converts internal document model field names to the external API
+    response field names (e.g., 'chunk_num' -> 'chunk_count').
+
+    Args:
+        doc: The document model from the database OR a dictionary.
+        run_status: Optional explicit run status value. If not provided:
+            - If doc has 'run' field, it will be mapped using run_mapping
+            - Otherwise, 'run' will be set to 'UNSTART' (for new uploads)
+
+    Returns:
+        A dictionary with renamed keys for API response.
+    """
+    renamed_doc = _process_key_mappings(doc)
+    renamed_doc = _process_run_mapping(renamed_doc, run_status)
+    return renamed_doc
+
+
+def _process_key_mappings(doc):
+    """
+    Map document keys to match API response format.
+
+    Converts internal document model field names to the external API
+    response field names (e.g., 'chunk_num' -> 'chunk_count').
+
+    Args:
+        doc: The document model from the database OR a dictionary.
+
     Returns:
         A dictionary with renamed keys for API response.
     """
@@ -184,6 +224,30 @@ def rename_doc_key(doc):
         "token_num": "token_count",
         "parser_id": "chunk_method",
     }
+
+    # Handle both dict and model input
+    items = doc.to_dict().items() if hasattr(doc, 'to_dict') else doc.items()
+
+    renamed_doc = {}
+    for key, value in items:
+        new_key = key_mapping.get(key, key)
+        renamed_doc[new_key] = value
+    return renamed_doc
+
+
+def _process_run_mapping(doc, run_status):
+    """
+    Map document keys to match API response format.
+
+    Args:
+        doc: The document model from the database OR a dictionary.
+        run_status: Optional explicit run status value. If not provided:
+            - If doc has 'run' field, it will be mapped using run_mapping
+            - Otherwise, 'run' will be set to 'UNSTART' (for new uploads)
+
+    Returns:
+        A dictionary with renamed keys for API response.
+    """
     run_mapping = {
         "0": "UNSTART",
         "1": "RUNNING",
@@ -191,11 +255,12 @@ def rename_doc_key(doc):
         "3": "DONE",
         "4": "FAIL",
     }
-    renamed_doc = {}
-    for key, value in doc.to_dict().items():
-        new_key = key_mapping.get(key, key)
-        renamed_doc[new_key] = value
-        if key == "run":
-            renamed_doc["run"] = run_mapping.get(str(value))
-    return renamed_doc
+
+    # Handle run field
+    if run_status is None or run_status not in run_mapping.keys():
+        run_status = "0"
+
+    doc["run"] = run_mapping[run_status]
+    return doc
+