refactor(lightrag): 优化文件上传接口

- 移除 InsertFileRequest 模型,改用 FastAPI 的 File 和 UploadFile
- 修改 insert_file 函数,以适应新的文件上传方式
- 更新函数参数和逻辑,支持直接上传文件
- 优化错误处理和响应消息
This commit is contained in:
90houlaoheshang
2024-11-05 17:22:04 +08:00
parent 11b791c72b
commit a140f744a3

View File

@@ -1,4 +1,4 @@
from fastapi import FastAPI, HTTPException
from fastapi import FastAPI, HTTPException, File, UploadFile
from pydantic import BaseModel
import os
from lightrag import LightRAG, QueryParam
@@ -78,10 +78,6 @@ class InsertRequest(BaseModel):
text: str
class InsertFileRequest(BaseModel):
file_path: str
class Response(BaseModel):
status: str
data: Optional[str] = None
@@ -115,30 +111,22 @@ async def insert_endpoint(request: InsertRequest):
@app.post("/insert_file", response_model=Response)
async def insert_file(request: InsertFileRequest):
async def insert_file(file: UploadFile = File(...)):
try:
# Check if file exists
if not os.path.exists(request.file_path):
raise HTTPException(
status_code=404, detail=f"File not found: {request.file_path}"
)
file_content = await file.read()
# Read file content
try:
with open(request.file_path, "r", encoding="utf-8") as f:
content = f.read()
content = file_content.decode("utf-8")
except UnicodeDecodeError:
# If UTF-8 decoding fails, try other encodings
with open(request.file_path, "r", encoding="gbk") as f:
content = f.read()
content = file_content.decode("gbk")
# Insert file content
loop = asyncio.get_event_loop()
await loop.run_in_executor(None, lambda: rag.insert(content))
return Response(
status="success",
message=f"File content from {request.file_path} inserted successfully",
message=f"File content from {file.filename} inserted successfully",
)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))