Commit
·
78e9f46
1
Parent(s):
4bc1d2e
Refactor code formatting in lightrag_api_openai_compatible_demo.py
Browse files
examples/lightrag_api_openai_compatible_demo.py
CHANGED
|
@@ -16,7 +16,7 @@ DEFAULT_RAG_DIR = "index_default"
|
|
| 16 |
app = FastAPI(title="LightRAG API", description="API for RAG operations")
|
| 17 |
|
| 18 |
# Configure working directory
|
| 19 |
-
WORKING_DIR = os.environ.get(
|
| 20 |
print(f"WORKING_DIR: {WORKING_DIR}")
|
| 21 |
if not os.path.exists(WORKING_DIR):
|
| 22 |
os.mkdir(WORKING_DIR)
|
|
@@ -32,11 +32,12 @@ async def llm_model_func(
|
|
| 32 |
prompt,
|
| 33 |
system_prompt=system_prompt,
|
| 34 |
history_messages=history_messages,
|
| 35 |
-
api_key=
|
| 36 |
base_url="YourURL/v1",
|
| 37 |
**kwargs,
|
| 38 |
)
|
| 39 |
|
|
|
|
| 40 |
# Embedding function
|
| 41 |
|
| 42 |
|
|
@@ -44,10 +45,11 @@ async def embedding_func(texts: list[str]) -> np.ndarray:
|
|
| 44 |
return await openai_embedding(
|
| 45 |
texts,
|
| 46 |
model="text-embedding-3-large",
|
| 47 |
-
api_key=
|
| 48 |
base_url="YourURL/v1",
|
| 49 |
)
|
| 50 |
|
|
|
|
| 51 |
# Initialize RAG instance
|
| 52 |
rag = LightRAG(
|
| 53 |
working_dir=WORKING_DIR,
|
|
@@ -78,6 +80,7 @@ class Response(BaseModel):
|
|
| 78 |
data: Optional[str] = None
|
| 79 |
message: Optional[str] = None
|
| 80 |
|
|
|
|
| 81 |
# API routes
|
| 82 |
|
| 83 |
|
|
@@ -86,14 +89,9 @@ async def query_endpoint(request: QueryRequest):
|
|
| 86 |
try:
|
| 87 |
loop = asyncio.get_event_loop()
|
| 88 |
result = await loop.run_in_executor(
|
| 89 |
-
None,
|
| 90 |
-
lambda: rag.query(
|
| 91 |
-
request.query, param=QueryParam(mode=request.mode))
|
| 92 |
-
)
|
| 93 |
-
return Response(
|
| 94 |
-
status="success",
|
| 95 |
-
data=result
|
| 96 |
)
|
|
|
|
| 97 |
except Exception as e:
|
| 98 |
raise HTTPException(status_code=500, detail=str(e))
|
| 99 |
|
|
@@ -103,10 +101,7 @@ async def insert_endpoint(request: InsertRequest):
|
|
| 103 |
try:
|
| 104 |
loop = asyncio.get_event_loop()
|
| 105 |
await loop.run_in_executor(None, lambda: rag.insert(request.text))
|
| 106 |
-
return Response(
|
| 107 |
-
status="success",
|
| 108 |
-
message="Text inserted successfully"
|
| 109 |
-
)
|
| 110 |
except Exception as e:
|
| 111 |
raise HTTPException(status_code=500, detail=str(e))
|
| 112 |
|
|
@@ -117,17 +112,16 @@ async def insert_file(request: InsertFileRequest):
|
|
| 117 |
# Check if file exists
|
| 118 |
if not os.path.exists(request.file_path):
|
| 119 |
raise HTTPException(
|
| 120 |
-
status_code=404,
|
| 121 |
-
detail=f"File not found: {request.file_path}"
|
| 122 |
)
|
| 123 |
|
| 124 |
# Read file content
|
| 125 |
try:
|
| 126 |
-
with open(request.file_path,
|
| 127 |
content = f.read()
|
| 128 |
except UnicodeDecodeError:
|
| 129 |
# If UTF-8 decoding fails, try other encodings
|
| 130 |
-
with open(request.file_path,
|
| 131 |
content = f.read()
|
| 132 |
|
| 133 |
# Insert file content
|
|
@@ -136,7 +130,7 @@ async def insert_file(request: InsertFileRequest):
|
|
| 136 |
|
| 137 |
return Response(
|
| 138 |
status="success",
|
| 139 |
-
message=f"File content from {request.file_path} inserted successfully"
|
| 140 |
)
|
| 141 |
except Exception as e:
|
| 142 |
raise HTTPException(status_code=500, detail=str(e))
|
|
@@ -146,8 +140,10 @@ async def insert_file(request: InsertFileRequest):
|
|
| 146 |
async def health_check():
|
| 147 |
return {"status": "healthy"}
|
| 148 |
|
|
|
|
| 149 |
if __name__ == "__main__":
|
| 150 |
import uvicorn
|
|
|
|
| 151 |
uvicorn.run(app, host="0.0.0.0", port=8020)
|
| 152 |
|
| 153 |
# Usage example
|
|
|
|
| 16 |
app = FastAPI(title="LightRAG API", description="API for RAG operations")
|
| 17 |
|
| 18 |
# Configure working directory
|
| 19 |
+
WORKING_DIR = os.environ.get("RAG_DIR", f"{DEFAULT_RAG_DIR}")
|
| 20 |
print(f"WORKING_DIR: {WORKING_DIR}")
|
| 21 |
if not os.path.exists(WORKING_DIR):
|
| 22 |
os.mkdir(WORKING_DIR)
|
|
|
|
| 32 |
prompt,
|
| 33 |
system_prompt=system_prompt,
|
| 34 |
history_messages=history_messages,
|
| 35 |
+
api_key="YOUR_API_KEY",
|
| 36 |
base_url="YourURL/v1",
|
| 37 |
**kwargs,
|
| 38 |
)
|
| 39 |
|
| 40 |
+
|
| 41 |
# Embedding function
|
| 42 |
|
| 43 |
|
|
|
|
| 45 |
return await openai_embedding(
|
| 46 |
texts,
|
| 47 |
model="text-embedding-3-large",
|
| 48 |
+
api_key="YOUR_API_KEY",
|
| 49 |
base_url="YourURL/v1",
|
| 50 |
)
|
| 51 |
|
| 52 |
+
|
| 53 |
# Initialize RAG instance
|
| 54 |
rag = LightRAG(
|
| 55 |
working_dir=WORKING_DIR,
|
|
|
|
| 80 |
data: Optional[str] = None
|
| 81 |
message: Optional[str] = None
|
| 82 |
|
| 83 |
+
|
| 84 |
# API routes
|
| 85 |
|
| 86 |
|
|
|
|
| 89 |
try:
|
| 90 |
loop = asyncio.get_event_loop()
|
| 91 |
result = await loop.run_in_executor(
|
| 92 |
+
None, lambda: rag.query(request.query, param=QueryParam(mode=request.mode))
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 93 |
)
|
| 94 |
+
return Response(status="success", data=result)
|
| 95 |
except Exception as e:
|
| 96 |
raise HTTPException(status_code=500, detail=str(e))
|
| 97 |
|
|
|
|
| 101 |
try:
|
| 102 |
loop = asyncio.get_event_loop()
|
| 103 |
await loop.run_in_executor(None, lambda: rag.insert(request.text))
|
| 104 |
+
return Response(status="success", message="Text inserted successfully")
|
|
|
|
|
|
|
|
|
|
| 105 |
except Exception as e:
|
| 106 |
raise HTTPException(status_code=500, detail=str(e))
|
| 107 |
|
|
|
|
| 112 |
# Check if file exists
|
| 113 |
if not os.path.exists(request.file_path):
|
| 114 |
raise HTTPException(
|
| 115 |
+
status_code=404, detail=f"File not found: {request.file_path}"
|
|
|
|
| 116 |
)
|
| 117 |
|
| 118 |
# Read file content
|
| 119 |
try:
|
| 120 |
+
with open(request.file_path, "r", encoding="utf-8") as f:
|
| 121 |
content = f.read()
|
| 122 |
except UnicodeDecodeError:
|
| 123 |
# If UTF-8 decoding fails, try other encodings
|
| 124 |
+
with open(request.file_path, "r", encoding="gbk") as f:
|
| 125 |
content = f.read()
|
| 126 |
|
| 127 |
# Insert file content
|
|
|
|
| 130 |
|
| 131 |
return Response(
|
| 132 |
status="success",
|
| 133 |
+
message=f"File content from {request.file_path} inserted successfully",
|
| 134 |
)
|
| 135 |
except Exception as e:
|
| 136 |
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
| 140 |
async def health_check():
|
| 141 |
return {"status": "healthy"}
|
| 142 |
|
| 143 |
+
|
| 144 |
if __name__ == "__main__":
|
| 145 |
import uvicorn
|
| 146 |
+
|
| 147 |
uvicorn.run(app, host="0.0.0.0", port=8020)
|
| 148 |
|
| 149 |
# Usage example
|