Refactor code formatting in lightrag_api_openai_compatible_demo.py

This commit is contained in:
zhenjie.ye
2024-10-26 16:00:30 +08:00
parent 02f94ab228
commit 08feac942a

View File

@@ -12,7 +12,7 @@ import nest_asyncio
# Apply nest_asyncio to solve event loop issues
nest_asyncio.apply()
DEFAULT_RAG_DIR="index_default"
DEFAULT_RAG_DIR = "index_default"
app = FastAPI(title="LightRAG API", description="API for RAG operations")
# Configure working directory
@@ -22,6 +22,8 @@ if not os.path.exists(WORKING_DIR):
os.mkdir(WORKING_DIR)
# LLM model function
async def llm_model_func(
prompt, system_prompt=None, history_messages=[], **kwargs
) -> str:
@@ -36,6 +38,8 @@ async def llm_model_func(
)
# Embedding function
async def embedding_func(texts: list[str]) -> np.ndarray:
return await openai_embedding(
texts,
@@ -54,29 +58,37 @@ rag = LightRAG(
)
# Data models
class QueryRequest(BaseModel):
query: str
mode: str = "hybrid"
class InsertRequest(BaseModel):
text: str
class InsertFileRequest(BaseModel):
file_path: str
class Response(BaseModel):
status: str
data: Optional[str] = None
message: Optional[str] = None
# API routes
@app.post("/query", response_model=Response)
async def query_endpoint(request: QueryRequest):
try:
loop = asyncio.get_event_loop()
result = await loop.run_in_executor(
None,
lambda: rag.query(request.query, param=QueryParam(mode=request.mode))
lambda: rag.query(
request.query, param=QueryParam(mode=request.mode))
)
return Response(
status="success",
@@ -85,6 +97,7 @@ async def query_endpoint(request: QueryRequest):
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/insert", response_model=Response)
async def insert_endpoint(request: InsertRequest):
try:
@@ -97,6 +110,7 @@ async def insert_endpoint(request: InsertRequest):
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/insert_file", response_model=Response)
async def insert_file(request: InsertFileRequest):
try:
@@ -127,6 +141,7 @@ async def insert_file(request: InsertFileRequest):
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/health")
async def health_check():
return {"status": "healthy"}