fix(avatar): upload knowledge files in chunks

2026-09-07 17:45:22 +08:00
parent 359e558dbe
commit e71267cf86
3 changed files with 336 additions and 21 deletions
@@ -1,4 +1,7 @@
 import os
+import json
+import shutil
+import time
 import uuid

 from fastapi import APIRouter, UploadFile, File, Depends, Header, HTTPException
@@ -19,6 +22,9 @@ os.makedirs(UPLOAD_DIR, exist_ok=True)
 ALLOWED_EXT = {".md", ".txt", ".pdf", ".doc", ".docx", ".xlsx"}
 MAX_UPLOAD_BYTES = 50 * 1024 * 1024
 UPLOAD_CHUNK_BYTES = 1024 * 1024
+MULTIPART_CHUNK_BYTES = 5 * 1024 * 1024
+MULTIPART_ROOT = ".multipart"
+MULTIPART_TTL_SECONDS = 24 * 60 * 60


 class QAIn(BaseModel):
@@ -31,6 +37,74 @@ class EnabledIn(BaseModel):
    enabled: bool = True


+class MultipartUploadIn(BaseModel):
+    filename: str
+    fileSize: int
+    totalChunks: int
+
+
+def _validate_document(filename: str, file_size: int):
+    ext = os.path.splitext(filename or "")[1].lower()
+    if ext not in ALLOWED_EXT:
+        return None, f"不支持的文件类型：{ext or '空'}，仅支持 md/txt/pdf/doc/docx/xlsx"
+    if file_size <= 0:
+        return None, "文件内容不能为空"
+    if file_size > MAX_UPLOAD_BYTES:
+        return None, "文件不能超过 50MB"
+    return ext, ""
+
+
+def _multipart_dir(avatar_id: str, upload_id: str) -> str:
+    safe_avatar_id = os.path.basename(avatar_id)
+    safe_upload_id = os.path.basename(upload_id)
+    if (
+        safe_avatar_id != avatar_id
+        or safe_upload_id != upload_id
+        or len(upload_id) != 32
+        or any(character not in "0123456789abcdef" for character in upload_id)
+    ):
+        raise HTTPException(status_code=400, detail="上传标识无效")
+    return os.path.join(UPLOAD_DIR, MULTIPART_ROOT, safe_avatar_id, safe_upload_id)
+
+
+def _purge_stale_multipart_uploads(avatar_id: str):
+    avatar_upload_root = os.path.join(UPLOAD_DIR, MULTIPART_ROOT, os.path.basename(avatar_id))
+    if not os.path.isdir(avatar_upload_root):
+        return
+    cutoff = time.time() - MULTIPART_TTL_SECONDS
+    for entry in os.scandir(avatar_upload_root):
+        if entry.is_dir(follow_symlinks=False) and entry.stat(follow_symlinks=False).st_mtime < cutoff:
+            shutil.rmtree(entry.path, ignore_errors=True)
+
+
+def _read_multipart_metadata(avatar_id: str, upload_id: str) -> tuple[str, dict]:
+    upload_dir = _multipart_dir(avatar_id, upload_id)
+    metadata_path = os.path.join(upload_dir, "metadata.json")
+    if not os.path.isfile(metadata_path):
+        raise HTTPException(status_code=404, detail="上传任务不存在或已过期")
+    with open(metadata_path, "r", encoding="utf-8") as stream:
+        return upload_dir, json.load(stream)
+
+
+def _create_knowledge_doc(db: Session, avatar_id: str, filename: str, ext: str, file_size: int, stored: str):
+    doc = KnowledgeDoc(
+        id=uuid.uuid4().hex,
+        avatar_id=avatar_id,
+        filename=filename,
+        file_type=ext.lstrip("."),
+        file_size=file_size,
+        file_url=f"/api/files/{avatar_id}/{stored}",
+        status="parsing",
+        index_stage="queued",
+        index_progress=0,
+    )
+    db.add(doc)
+    db.commit()
+    db.refresh(doc)
+    knowledge_vectorizer.enqueue(doc.id)
+    return doc
+
+
 def _doc_payload(doc: KnowledgeDoc) -> dict:
    payload = doc.to_dict()
    stored_name = os.path.basename(doc.file_url or "")
@@ -74,9 +148,9 @@ def list_docs(avatar_id: str, authorization: str = Header(None), db: Session = D
@router.post("/avatar/{avatar_id}/knowledge/docs")
 async def upload_doc(avatar_id: str, file: UploadFile = File(...), authorization: str = Header(None), db: Session = Depends(get_db)):
    _require_owned_avatar(db, avatar_id, authorization)
-    ext = os.path.splitext(file.filename or "")[1].lower()
-    if ext not in ALLOWED_EXT:
-        return fail(f"不支持的文件类型：{ext or '空'}，仅支持 md/txt/pdf/doc/docx/xlsx", code=400)
+    ext, validation_error = _validate_document(file.filename or "", 1)
+    if validation_error:
+        return fail(validation_error, code=400)
    avatar_dir = os.path.join(UPLOAD_DIR, avatar_id)
    os.makedirs(avatar_dir, exist_ok=True)
    stored = f"{uuid.uuid4().hex}{ext}"
@@ -94,25 +168,124 @@ async def upload_doc(avatar_id: str, file: UploadFile = File(...), authorization
        if os.path.exists(path):
            os.remove(path)
        return fail(str(exc), code=400)
-    doc = KnowledgeDoc(
-        id=uuid.uuid4().hex,
-        avatar_id=avatar_id,
-        filename=file.filename,
-        file_type=ext.lstrip("."),
-        file_size=file_size,
-        file_url=f"/api/files/{avatar_id}/{stored}",
-        status="parsing",
-        index_stage="queued",
-        index_progress=0,
+    if file_size == 0:
+        if os.path.exists(path):
+            os.remove(path)
+        return fail("文件内容不能为空", code=400)
+
+    doc = _create_knowledge_doc(db, avatar_id, file.filename or stored, ext, file_size, stored)
+    return ok(_doc_payload(doc))
+
+
+@router.post("/avatar/{avatar_id}/knowledge/uploads")
+def create_multipart_upload(
+    avatar_id: str,
+    body: MultipartUploadIn,
+    authorization: str = Header(None),
+    db: Session = Depends(get_db),
+):
+    _require_owned_avatar(db, avatar_id, authorization)
+    ext, validation_error = _validate_document(body.filename, body.fileSize)
+    if validation_error:
+        return fail(validation_error, code=400)
+    expected_chunks = (body.fileSize + MULTIPART_CHUNK_BYTES - 1) // MULTIPART_CHUNK_BYTES
+    if body.totalChunks != expected_chunks:
+        return fail("文件分片数量不正确", code=400)
+
+    _purge_stale_multipart_uploads(avatar_id)
+    upload_id = uuid.uuid4().hex
+    upload_dir = _multipart_dir(avatar_id, upload_id)
+    os.makedirs(upload_dir, exist_ok=False)
+    metadata = {
+        "filename": body.filename,
+        "fileSize": body.fileSize,
+        "totalChunks": body.totalChunks,
+        "extension": ext,
+    }
+    with open(os.path.join(upload_dir, "metadata.json"), "w", encoding="utf-8") as stream:
+        json.dump(metadata, stream, ensure_ascii=False)
+    return ok({"uploadId": upload_id, "chunkSize": MULTIPART_CHUNK_BYTES})
+
+
+@router.post("/avatar/{avatar_id}/knowledge/uploads/{upload_id}/chunks/{chunk_index}")
+async def upload_multipart_chunk(
+    avatar_id: str,
+    upload_id: str,
+    chunk_index: int,
+    file: UploadFile = File(...),
+    authorization: str = Header(None),
+    db: Session = Depends(get_db),
+):
+    _require_owned_avatar(db, avatar_id, authorization)
+    upload_dir, metadata = _read_multipart_metadata(avatar_id, upload_id)
+    total_chunks = int(metadata["totalChunks"])
+    if chunk_index < 0 or chunk_index >= total_chunks:
+        return fail("文件分片序号不正确", code=400)
+
+    expected_size = min(
+        MULTIPART_CHUNK_BYTES,
+        int(metadata["fileSize"]) - chunk_index * MULTIPART_CHUNK_BYTES,
    )
+    part_path = os.path.join(upload_dir, f"{chunk_index}.part")
+    temporary_path = f"{part_path}.uploading"
+    received = 0
+    try:
+        with open(temporary_path, "wb") as stream:
+            while chunk := await file.read(UPLOAD_CHUNK_BYTES):
+                received += len(chunk)
+                if received > expected_size:
+                    raise ValueError("文件分片大小不正确")
+                stream.write(chunk)
+        if received != expected_size:
+            raise ValueError("文件分片大小不正确")
+        os.replace(temporary_path, part_path)
+    except ValueError as exc:
+        if os.path.exists(temporary_path):
+            os.remove(temporary_path)
+        return fail(str(exc), code=400)
+    return ok({"chunkIndex": chunk_index, "uploadedBytes": received})

-    # Persist and acknowledge the upload first. Extraction and embeddings may take
-    # minutes for a PDF and must never consume the browser request timeout.
-    db.add(doc)
-    db.commit()
-    db.refresh(doc)
-    knowledge_vectorizer.enqueue(doc.id)

+@router.post("/avatar/{avatar_id}/knowledge/uploads/{upload_id}/complete")
+def complete_multipart_upload(
+    avatar_id: str,
+    upload_id: str,
+    authorization: str = Header(None),
+    db: Session = Depends(get_db),
+):
+    _require_owned_avatar(db, avatar_id, authorization)
+    upload_dir, metadata = _read_multipart_metadata(avatar_id, upload_id)
+    total_chunks = int(metadata["totalChunks"])
+    part_paths = [os.path.join(upload_dir, f"{index}.part") for index in range(total_chunks)]
+    if not all(os.path.isfile(path) for path in part_paths):
+        return fail("文件分片尚未上传完整", code=400)
+    if sum(os.path.getsize(path) for path in part_paths) != int(metadata["fileSize"]):
+        return fail("文件分片总大小不正确", code=400)
+
+    avatar_dir = os.path.join(UPLOAD_DIR, avatar_id)
+    os.makedirs(avatar_dir, exist_ok=True)
+    stored = f"{uuid.uuid4().hex}{metadata['extension']}"
+    final_path = os.path.join(avatar_dir, stored)
+    temporary_path = f"{final_path}.assembling"
+    try:
+        with open(temporary_path, "wb") as output:
+            for part_path in part_paths:
+                with open(part_path, "rb") as source:
+                    shutil.copyfileobj(source, output, UPLOAD_CHUNK_BYTES)
+        os.replace(temporary_path, final_path)
+        doc = _create_knowledge_doc(
+            db,
+            avatar_id,
+            metadata["filename"],
+            metadata["extension"],
+            int(metadata["fileSize"]),
+            stored,
+        )
+    except Exception:
+        if os.path.exists(temporary_path):
+            os.remove(temporary_path)
+        raise
+    shutil.rmtree(upload_dir, ignore_errors=True)
    return ok(_doc_payload(doc))