faiss_rag_enterprise/app/api/upload.py

33 lines
1.1 KiB
Python

from fastapi import APIRouter, UploadFile, File, Form, HTTPException
import os
import shutil
from scripts.rag_build_query import build_user_index
router = APIRouter()
ALLOWED_SUFFIXES = {".txt", ".md", ".pdf", ".docx"}
@router.post("/upload")
def upload_user_file(user_id: str = Form(...), file: UploadFile = File(...)):
filename = os.path.basename(file.filename)
suffix = os.path.splitext(filename)[-1].lower()
if suffix not in ALLOWED_SUFFIXES:
raise HTTPException(status_code=400, detail="不支持的文件类型")
user_doc_dir = os.path.join("docs", user_id)
os.makedirs(user_doc_dir, exist_ok=True)
file_path = os.path.join(user_doc_dir, filename)
try:
with open(file_path, "wb") as f:
shutil.copyfileobj(file.file, f)
print(f"[UPLOAD] 文件已保存至 {file_path}")
build_user_index(user_id)
print(f"[UPLOAD] 用户 {user_id} 的索引已重建")
except Exception as e:
print(f"[UPLOAD ERROR] {e}")
raise HTTPException(status_code=500, detail="索引构建失败")
return {"status": "ok", "filename": filename}