thinkthinking commited on
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('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,11 +32,12 @@ async def llm_model_func(
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
  # 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='YOUR_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, 'r', encoding='utf-8') as f:
127
  content = f.read()
128
  except UnicodeDecodeError:
129
  # If UTF-8 decoding fails, try other encodings
130
- with open(request.file_path, 'r', encoding='gbk') as f:
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