64 lines
2.5 KiB
Python
64 lines
2.5 KiB
Python
from fastapi import APIRouter, UploadFile, File, Form, HTTPException
|
||
import os
|
||
import shutil
|
||
import logging
|
||
from scripts.rag_build_query import build_user_index
|
||
|
||
# 设置日志
|
||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||
logger = logging.getLogger(__name__)
|
||
|
||
router = APIRouter()
|
||
ALLOWED_SUFFIXES = {".txt", ".md", ".pdf", ".docx"}
|
||
MAX_FILE_SIZE = 50 * 1024 * 1024 # 设置最大文件大小限制，50MB
|
||
|
||
@router.post("/upload")
|
||
def upload_user_file(user_id: str = Form(...), file: UploadFile = File(...)):
|
||
# 参数检查
|
||
if not user_id or len(user_id.strip()) == 0:
|
||
raise HTTPException(status_code=400, detail="用户ID不能为空")
|
||
|
||
# 检查文件类型
|
||
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="不支持的文件类型")
|
||
|
||
# 检查文件大小
|
||
file_size = len(file.file.read()) # 获取文件大小（字节）
|
||
if file_size > MAX_FILE_SIZE:
|
||
raise HTTPException(status_code=400, detail="文件太大，请上传不超过50MB的文件")
|
||
|
||
# 创建用户文档目录
|
||
user_doc_dir = os.path.join("docs", user_id)
|
||
if not os.path.exists(user_doc_dir): # 检查目录是否存在
|
||
os.makedirs(user_doc_dir, exist_ok=True) # 如果不存在，创建目录
|
||
|
||
# 创建索引数据目录，确保路径存在
|
||
index_data_dir = os.path.join("index_data", user_id)
|
||
if not os.path.exists(index_data_dir): # 检查目录是否存在
|
||
os.makedirs(index_data_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)
|
||
logger.info(f"文件已保存至 {file_path}")
|
||
|
||
# 创建索引并保存 docstore.json
|
||
docstore_path = os.path.join(index_data_dir, "docstore.json")
|
||
if not os.path.exists(docstore_path): # 检查文件是否存在
|
||
with open(docstore_path, "w") as docstore_file:
|
||
docstore_file.write("{}") # 可根据需要保存实际内容
|
||
|
||
# 重新构建用户索引
|
||
build_user_index(user_id)
|
||
logger.info(f"用户 {user_id} 的索引已重建")
|
||
|
||
except Exception as e:
|
||
logger.error(f"[UPLOAD ERROR] {e}")
|
||
raise HTTPException(status_code=500, detail="索引构建失败")
|
||
|
||
return {"status": "ok", "filename": filename}
|