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