samuel-z-chen commited on
Commit
4d8e9a6
·
2 Parent(s): d6836a4 ab2732b

Merge remote-tracking branch 'origin/main'

Browse files
README.md CHANGED
@@ -1025,6 +1025,7 @@ Each server has its own specific configuration options:
1025
  | --max-embed-tokens | 8192 | Maximum embedding token size |
1026
  | --input-file | ./book.txt | Initial input file |
1027
  | --log-level | INFO | Logging level |
 
1028
 
1029
  #### Ollama Server Options
1030
 
@@ -1042,6 +1043,7 @@ Each server has its own specific configuration options:
1042
  | --max-embed-tokens | 8192 | Maximum embedding token size |
1043
  | --input-file | ./book.txt | Initial input file |
1044
  | --log-level | INFO | Logging level |
 
1045
 
1046
  #### OpenAI Server Options
1047
 
@@ -1056,6 +1058,7 @@ Each server has its own specific configuration options:
1056
  | --max-embed-tokens | 8192 | Maximum embedding token size |
1057
  | --input-dir | ./inputs | Input directory for documents |
1058
  | --log-level | INFO | Logging level |
 
1059
 
1060
  #### OpenAI AZURE Server Options
1061
 
@@ -1071,8 +1074,10 @@ Each server has its own specific configuration options:
1071
  | --input-dir | ./inputs | Input directory for documents |
1072
  | --enable-cache | True | Enable response cache |
1073
  | --log-level | INFO | Logging level |
 
1074
 
1075
 
 
1076
  ### Example Usage
1077
 
1078
  #### LoLLMs RAG Server
@@ -1083,6 +1088,10 @@ lollms-lightrag-server --model mistral-nemo --port 8080 --working-dir ./custom_r
1083
 
1084
  # Using specific models (ensure they are installed in your LoLLMs instance)
1085
  lollms-lightrag-server --model mistral-nemo:latest --embedding-model bge-m3 --embedding-dim 1024
 
 
 
 
1086
  ```
1087
 
1088
  #### Ollama RAG Server
 
1025
  | --max-embed-tokens | 8192 | Maximum embedding token size |
1026
  | --input-file | ./book.txt | Initial input file |
1027
  | --log-level | INFO | Logging level |
1028
+ | --key | none | Access Key to protect the lightrag service |
1029
 
1030
  #### Ollama Server Options
1031
 
 
1043
  | --max-embed-tokens | 8192 | Maximum embedding token size |
1044
  | --input-file | ./book.txt | Initial input file |
1045
  | --log-level | INFO | Logging level |
1046
+ | --key | none | Access Key to protect the lightrag service |
1047
 
1048
  #### OpenAI Server Options
1049
 
 
1058
  | --max-embed-tokens | 8192 | Maximum embedding token size |
1059
  | --input-dir | ./inputs | Input directory for documents |
1060
  | --log-level | INFO | Logging level |
1061
+ | --key | none | Access Key to protect the lightrag service |
1062
 
1063
  #### OpenAI AZURE Server Options
1064
 
 
1074
  | --input-dir | ./inputs | Input directory for documents |
1075
  | --enable-cache | True | Enable response cache |
1076
  | --log-level | INFO | Logging level |
1077
+ | --key | none | Access Key to protect the lightrag service |
1078
 
1079
 
1080
+ For protecting the server using an authentication key, you can also use an environment variable named `LIGHTRAG_API_KEY`.
1081
  ### Example Usage
1082
 
1083
  #### LoLLMs RAG Server
 
1088
 
1089
  # Using specific models (ensure they are installed in your LoLLMs instance)
1090
  lollms-lightrag-server --model mistral-nemo:latest --embedding-model bge-m3 --embedding-dim 1024
1091
+
1092
+ # Using specific models and an authentication key
1093
+ lollms-lightrag-server --model mistral-nemo:latest --embedding-model bge-m3 --embedding-dim 1024 --key ky-mykey
1094
+
1095
  ```
1096
 
1097
  #### Ollama RAG Server
lightrag/api/azure_openai_lightrag_server.py CHANGED
@@ -21,6 +21,12 @@ import inspect
21
  import json
22
  from fastapi.responses import StreamingResponse
23
 
 
 
 
 
 
 
24
  load_dotenv()
25
 
26
  AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION")
@@ -93,6 +99,13 @@ def parse_args():
93
  help="Logging level (default: INFO)",
94
  )
95
 
 
 
 
 
 
 
 
96
  return parser.parse_args()
97
 
98
 
@@ -155,6 +168,31 @@ class InsertResponse(BaseModel):
155
  document_count: int
156
 
157
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
158
  async def get_embedding_dim(embedding_model: str) -> int:
159
  """Get embedding dimensions for the specified model"""
160
  test_text = ["This is a test sentence."]
@@ -168,12 +206,32 @@ def create_app(args):
168
  format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
169
  )
170
 
171
- # Initialize FastAPI app
 
 
 
172
  app = FastAPI(
173
  title="LightRAG API",
174
- description="API for querying text using LightRAG with OpenAI integration",
 
 
 
 
 
175
  )
176
 
 
 
 
 
 
 
 
 
 
 
 
 
177
  # Create working directory if it doesn't exist
178
  Path(args.working_dir).mkdir(parents=True, exist_ok=True)
179
 
@@ -239,7 +297,7 @@ def create_app(args):
239
  except Exception as e:
240
  logging.error(f"Error during startup indexing: {str(e)}")
241
 
242
- @app.post("/documents/scan")
243
  async def scan_for_new_documents():
244
  """Manually trigger scanning for new documents"""
245
  try:
@@ -264,7 +322,7 @@ def create_app(args):
264
  except Exception as e:
265
  raise HTTPException(status_code=500, detail=str(e))
266
 
267
- @app.post("/resetcache")
268
  async def reset_cache():
269
  """Manually reset cache"""
270
  try:
@@ -276,7 +334,7 @@ def create_app(args):
276
  except Exception as e:
277
  raise HTTPException(status_code=500, detail=str(e))
278
 
279
- @app.post("/documents/upload")
280
  async def upload_to_input_dir(file: UploadFile = File(...)):
281
  """Upload a file to the input directory"""
282
  try:
@@ -304,7 +362,9 @@ def create_app(args):
304
  except Exception as e:
305
  raise HTTPException(status_code=500, detail=str(e))
306
 
307
- @app.post("/query", response_model=QueryResponse)
 
 
308
  async def query_text(request: QueryRequest):
309
  try:
310
  response = await rag.aquery(
@@ -319,7 +379,7 @@ def create_app(args):
319
  except Exception as e:
320
  raise HTTPException(status_code=500, detail=str(e))
321
 
322
- @app.post("/query/stream")
323
  async def query_text_stream(request: QueryRequest):
324
  try:
325
  response = await rag.aquery(
@@ -345,7 +405,11 @@ def create_app(args):
345
  except Exception as e:
346
  raise HTTPException(status_code=500, detail=str(e))
347
 
348
- @app.post("/documents/text", response_model=InsertResponse)
 
 
 
 
349
  async def insert_text(request: InsertTextRequest):
350
  try:
351
  await rag.ainsert(request.text)
@@ -357,7 +421,11 @@ def create_app(args):
357
  except Exception as e:
358
  raise HTTPException(status_code=500, detail=str(e))
359
 
360
- @app.post("/documents/file", response_model=InsertResponse)
 
 
 
 
361
  async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
362
  try:
363
  content = await file.read()
@@ -381,7 +449,11 @@ def create_app(args):
381
  except Exception as e:
382
  raise HTTPException(status_code=500, detail=str(e))
383
 
384
- @app.post("/documents/batch", response_model=InsertResponse)
 
 
 
 
385
  async def insert_batch(files: List[UploadFile] = File(...)):
386
  try:
387
  inserted_count = 0
@@ -411,7 +483,11 @@ def create_app(args):
411
  except Exception as e:
412
  raise HTTPException(status_code=500, detail=str(e))
413
 
414
- @app.delete("/documents", response_model=InsertResponse)
 
 
 
 
415
  async def clear_documents():
416
  try:
417
  rag.text_chunks = []
@@ -425,7 +501,7 @@ def create_app(args):
425
  except Exception as e:
426
  raise HTTPException(status_code=500, detail=str(e))
427
 
428
- @app.get("/health")
429
  async def get_status():
430
  """Get current system status"""
431
  return {
 
21
  import json
22
  from fastapi.responses import StreamingResponse
23
 
24
+ from fastapi import Depends, Security
25
+ from fastapi.security import APIKeyHeader
26
+ from fastapi.middleware.cors import CORSMiddleware
27
+
28
+ from starlette.status import HTTP_403_FORBIDDEN
29
+
30
  load_dotenv()
31
 
32
  AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION")
 
99
  help="Logging level (default: INFO)",
100
  )
101
 
102
+ parser.add_argument(
103
+ "--key",
104
+ type=str,
105
+ help="API key for authentication. This protects lightrag server against unauthorized access",
106
+ default=None,
107
+ )
108
+
109
  return parser.parse_args()
110
 
111
 
 
168
  document_count: int
169
 
170
 
171
+ def get_api_key_dependency(api_key: Optional[str]):
172
+ if not api_key:
173
+ # If no API key is configured, return a dummy dependency that always succeeds
174
+ async def no_auth():
175
+ return None
176
+
177
+ return no_auth
178
+
179
+ # If API key is configured, use proper authentication
180
+ api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
181
+
182
+ async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
183
+ if not api_key_header_value:
184
+ raise HTTPException(
185
+ status_code=HTTP_403_FORBIDDEN, detail="API Key required"
186
+ )
187
+ if api_key_header_value != api_key:
188
+ raise HTTPException(
189
+ status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
190
+ )
191
+ return api_key_header_value
192
+
193
+ return api_key_auth
194
+
195
+
196
  async def get_embedding_dim(embedding_model: str) -> int:
197
  """Get embedding dimensions for the specified model"""
198
  test_text = ["This is a test sentence."]
 
206
  format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
207
  )
208
 
209
+ # Check if API key is provided either through env var or args
210
+ api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
211
+
212
+ # Initialize FastAPI
213
  app = FastAPI(
214
  title="LightRAG API",
215
+ description="API for querying text using LightRAG with separate storage and input directories"
216
+ + "(With authentication)"
217
+ if api_key
218
+ else "",
219
+ version="1.0.0",
220
+ openapi_tags=[{"name": "api"}],
221
  )
222
 
223
+ # Add CORS middleware
224
+ app.add_middleware(
225
+ CORSMiddleware,
226
+ allow_origins=["*"],
227
+ allow_credentials=True,
228
+ allow_methods=["*"],
229
+ allow_headers=["*"],
230
+ )
231
+
232
+ # Create the optional API key dependency
233
+ optional_api_key = get_api_key_dependency(api_key)
234
+
235
  # Create working directory if it doesn't exist
236
  Path(args.working_dir).mkdir(parents=True, exist_ok=True)
237
 
 
297
  except Exception as e:
298
  logging.error(f"Error during startup indexing: {str(e)}")
299
 
300
+ @app.post("/documents/scan", dependencies=[Depends(optional_api_key)])
301
  async def scan_for_new_documents():
302
  """Manually trigger scanning for new documents"""
303
  try:
 
322
  except Exception as e:
323
  raise HTTPException(status_code=500, detail=str(e))
324
 
325
+ @app.post("/resetcache", dependencies=[Depends(optional_api_key)])
326
  async def reset_cache():
327
  """Manually reset cache"""
328
  try:
 
334
  except Exception as e:
335
  raise HTTPException(status_code=500, detail=str(e))
336
 
337
+ @app.post("/documents/upload", dependencies=[Depends(optional_api_key)])
338
  async def upload_to_input_dir(file: UploadFile = File(...)):
339
  """Upload a file to the input directory"""
340
  try:
 
362
  except Exception as e:
363
  raise HTTPException(status_code=500, detail=str(e))
364
 
365
+ @app.post(
366
+ "/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
367
+ )
368
  async def query_text(request: QueryRequest):
369
  try:
370
  response = await rag.aquery(
 
379
  except Exception as e:
380
  raise HTTPException(status_code=500, detail=str(e))
381
 
382
+ @app.post("/query/stream", dependencies=[Depends(optional_api_key)])
383
  async def query_text_stream(request: QueryRequest):
384
  try:
385
  response = await rag.aquery(
 
405
  except Exception as e:
406
  raise HTTPException(status_code=500, detail=str(e))
407
 
408
+ @app.post(
409
+ "/documents/text",
410
+ response_model=InsertResponse,
411
+ dependencies=[Depends(optional_api_key)],
412
+ )
413
  async def insert_text(request: InsertTextRequest):
414
  try:
415
  await rag.ainsert(request.text)
 
421
  except Exception as e:
422
  raise HTTPException(status_code=500, detail=str(e))
423
 
424
+ @app.post(
425
+ "/documents/file",
426
+ response_model=InsertResponse,
427
+ dependencies=[Depends(optional_api_key)],
428
+ )
429
  async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
430
  try:
431
  content = await file.read()
 
449
  except Exception as e:
450
  raise HTTPException(status_code=500, detail=str(e))
451
 
452
+ @app.post(
453
+ "/documents/batch",
454
+ response_model=InsertResponse,
455
+ dependencies=[Depends(optional_api_key)],
456
+ )
457
  async def insert_batch(files: List[UploadFile] = File(...)):
458
  try:
459
  inserted_count = 0
 
483
  except Exception as e:
484
  raise HTTPException(status_code=500, detail=str(e))
485
 
486
+ @app.delete(
487
+ "/documents",
488
+ response_model=InsertResponse,
489
+ dependencies=[Depends(optional_api_key)],
490
+ )
491
  async def clear_documents():
492
  try:
493
  rag.text_chunks = []
 
501
  except Exception as e:
502
  raise HTTPException(status_code=500, detail=str(e))
503
 
504
+ @app.get("/health", dependencies=[Depends(optional_api_key)])
505
  async def get_status():
506
  """Get current system status"""
507
  return {
lightrag/api/lollms_lightrag_server.py CHANGED
@@ -11,6 +11,13 @@ from pathlib import Path
11
  import shutil
12
  import aiofiles
13
  from ascii_colors import trace_exception
 
 
 
 
 
 
 
14
 
15
 
16
  def parse_args():
@@ -86,6 +93,13 @@ def parse_args():
86
  help="Logging level (default: INFO)",
87
  )
88
 
 
 
 
 
 
 
 
89
  return parser.parse_args()
90
 
91
 
@@ -148,18 +162,63 @@ class InsertResponse(BaseModel):
148
  document_count: int
149
 
150
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
151
  def create_app(args):
152
  # Setup logging
153
  logging.basicConfig(
154
  format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
155
  )
156
 
157
- # Initialize FastAPI app
 
 
 
158
  app = FastAPI(
159
  title="LightRAG API",
160
- description="API for querying text using LightRAG with separate storage and input directories",
 
 
 
 
 
161
  )
162
 
 
 
 
 
 
 
 
 
 
 
 
 
163
  # Create working directory if it doesn't exist
164
  Path(args.working_dir).mkdir(parents=True, exist_ok=True)
165
 
@@ -209,7 +268,7 @@ def create_app(args):
209
  except Exception as e:
210
  logging.error(f"Error during startup indexing: {str(e)}")
211
 
212
- @app.post("/documents/scan")
213
  async def scan_for_new_documents():
214
  """Manually trigger scanning for new documents"""
215
  try:
@@ -234,7 +293,7 @@ def create_app(args):
234
  except Exception as e:
235
  raise HTTPException(status_code=500, detail=str(e))
236
 
237
- @app.post("/documents/upload")
238
  async def upload_to_input_dir(file: UploadFile = File(...)):
239
  """Upload a file to the input directory"""
240
  try:
@@ -262,7 +321,9 @@ def create_app(args):
262
  except Exception as e:
263
  raise HTTPException(status_code=500, detail=str(e))
264
 
265
- @app.post("/query", response_model=QueryResponse)
 
 
266
  async def query_text(request: QueryRequest):
267
  try:
268
  response = await rag.aquery(
@@ -284,7 +345,7 @@ def create_app(args):
284
  except Exception as e:
285
  raise HTTPException(status_code=500, detail=str(e))
286
 
287
- @app.post("/query/stream")
288
  async def query_text_stream(request: QueryRequest):
289
  try:
290
  response = rag.query(
@@ -304,7 +365,11 @@ def create_app(args):
304
  except Exception as e:
305
  raise HTTPException(status_code=500, detail=str(e))
306
 
307
- @app.post("/documents/text", response_model=InsertResponse)
 
 
 
 
308
  async def insert_text(request: InsertTextRequest):
309
  try:
310
  rag.insert(request.text)
@@ -316,7 +381,11 @@ def create_app(args):
316
  except Exception as e:
317
  raise HTTPException(status_code=500, detail=str(e))
318
 
319
- @app.post("/documents/file", response_model=InsertResponse)
 
 
 
 
320
  async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
321
  try:
322
  content = await file.read()
@@ -340,7 +409,11 @@ def create_app(args):
340
  except Exception as e:
341
  raise HTTPException(status_code=500, detail=str(e))
342
 
343
- @app.post("/documents/batch", response_model=InsertResponse)
 
 
 
 
344
  async def insert_batch(files: List[UploadFile] = File(...)):
345
  try:
346
  inserted_count = 0
@@ -370,7 +443,11 @@ def create_app(args):
370
  except Exception as e:
371
  raise HTTPException(status_code=500, detail=str(e))
372
 
373
- @app.delete("/documents", response_model=InsertResponse)
 
 
 
 
374
  async def clear_documents():
375
  try:
376
  rag.text_chunks = []
@@ -384,7 +461,7 @@ def create_app(args):
384
  except Exception as e:
385
  raise HTTPException(status_code=500, detail=str(e))
386
 
387
- @app.get("/health")
388
  async def get_status():
389
  """Get current system status"""
390
  return {
 
11
  import shutil
12
  import aiofiles
13
  from ascii_colors import trace_exception
14
+ import os
15
+
16
+ from fastapi import Depends, Security
17
+ from fastapi.security import APIKeyHeader
18
+ from fastapi.middleware.cors import CORSMiddleware
19
+
20
+ from starlette.status import HTTP_403_FORBIDDEN
21
 
22
 
23
  def parse_args():
 
93
  help="Logging level (default: INFO)",
94
  )
95
 
96
+ parser.add_argument(
97
+ "--key",
98
+ type=str,
99
+ help="API key for authentication. This protects lightrag server against unauthorized access",
100
+ default=None,
101
+ )
102
+
103
  return parser.parse_args()
104
 
105
 
 
162
  document_count: int
163
 
164
 
165
+ def get_api_key_dependency(api_key: Optional[str]):
166
+ if not api_key:
167
+ # If no API key is configured, return a dummy dependency that always succeeds
168
+ async def no_auth():
169
+ return None
170
+
171
+ return no_auth
172
+
173
+ # If API key is configured, use proper authentication
174
+ api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
175
+
176
+ async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
177
+ if not api_key_header_value:
178
+ raise HTTPException(
179
+ status_code=HTTP_403_FORBIDDEN, detail="API Key required"
180
+ )
181
+ if api_key_header_value != api_key:
182
+ raise HTTPException(
183
+ status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
184
+ )
185
+ return api_key_header_value
186
+
187
+ return api_key_auth
188
+
189
+
190
  def create_app(args):
191
  # Setup logging
192
  logging.basicConfig(
193
  format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
194
  )
195
 
196
+ # Check if API key is provided either through env var or args
197
+ api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
198
+
199
+ # Initialize FastAPI
200
  app = FastAPI(
201
  title="LightRAG API",
202
+ description="API for querying text using LightRAG with separate storage and input directories"
203
+ + "(With authentication)"
204
+ if api_key
205
+ else "",
206
+ version="1.0.0",
207
+ openapi_tags=[{"name": "api"}],
208
  )
209
 
210
+ # Add CORS middleware
211
+ app.add_middleware(
212
+ CORSMiddleware,
213
+ allow_origins=["*"],
214
+ allow_credentials=True,
215
+ allow_methods=["*"],
216
+ allow_headers=["*"],
217
+ )
218
+
219
+ # Create the optional API key dependency
220
+ optional_api_key = get_api_key_dependency(api_key)
221
+
222
  # Create working directory if it doesn't exist
223
  Path(args.working_dir).mkdir(parents=True, exist_ok=True)
224
 
 
268
  except Exception as e:
269
  logging.error(f"Error during startup indexing: {str(e)}")
270
 
271
+ @app.post("/documents/scan", dependencies=[Depends(optional_api_key)])
272
  async def scan_for_new_documents():
273
  """Manually trigger scanning for new documents"""
274
  try:
 
293
  except Exception as e:
294
  raise HTTPException(status_code=500, detail=str(e))
295
 
296
+ @app.post("/documents/upload", dependencies=[Depends(optional_api_key)])
297
  async def upload_to_input_dir(file: UploadFile = File(...)):
298
  """Upload a file to the input directory"""
299
  try:
 
321
  except Exception as e:
322
  raise HTTPException(status_code=500, detail=str(e))
323
 
324
+ @app.post(
325
+ "/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
326
+ )
327
  async def query_text(request: QueryRequest):
328
  try:
329
  response = await rag.aquery(
 
345
  except Exception as e:
346
  raise HTTPException(status_code=500, detail=str(e))
347
 
348
+ @app.post("/query/stream", dependencies=[Depends(optional_api_key)])
349
  async def query_text_stream(request: QueryRequest):
350
  try:
351
  response = rag.query(
 
365
  except Exception as e:
366
  raise HTTPException(status_code=500, detail=str(e))
367
 
368
+ @app.post(
369
+ "/documents/text",
370
+ response_model=InsertResponse,
371
+ dependencies=[Depends(optional_api_key)],
372
+ )
373
  async def insert_text(request: InsertTextRequest):
374
  try:
375
  rag.insert(request.text)
 
381
  except Exception as e:
382
  raise HTTPException(status_code=500, detail=str(e))
383
 
384
+ @app.post(
385
+ "/documents/file",
386
+ response_model=InsertResponse,
387
+ dependencies=[Depends(optional_api_key)],
388
+ )
389
  async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
390
  try:
391
  content = await file.read()
 
409
  except Exception as e:
410
  raise HTTPException(status_code=500, detail=str(e))
411
 
412
+ @app.post(
413
+ "/documents/batch",
414
+ response_model=InsertResponse,
415
+ dependencies=[Depends(optional_api_key)],
416
+ )
417
  async def insert_batch(files: List[UploadFile] = File(...)):
418
  try:
419
  inserted_count = 0
 
443
  except Exception as e:
444
  raise HTTPException(status_code=500, detail=str(e))
445
 
446
+ @app.delete(
447
+ "/documents",
448
+ response_model=InsertResponse,
449
+ dependencies=[Depends(optional_api_key)],
450
+ )
451
  async def clear_documents():
452
  try:
453
  rag.text_chunks = []
 
461
  except Exception as e:
462
  raise HTTPException(status_code=500, detail=str(e))
463
 
464
+ @app.get("/health", dependencies=[Depends(optional_api_key)])
465
  async def get_status():
466
  """Get current system status"""
467
  return {
lightrag/api/ollama_lightrag_server.py CHANGED
@@ -11,6 +11,13 @@ from pathlib import Path
11
  import shutil
12
  import aiofiles
13
  from ascii_colors import trace_exception
 
 
 
 
 
 
 
14
 
15
 
16
  def parse_args():
@@ -85,6 +92,12 @@ def parse_args():
85
  choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
86
  help="Logging level (default: INFO)",
87
  )
 
 
 
 
 
 
88
 
89
  return parser.parse_args()
90
 
@@ -148,18 +161,63 @@ class InsertResponse(BaseModel):
148
  document_count: int
149
 
150
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
151
  def create_app(args):
152
  # Setup logging
153
  logging.basicConfig(
154
  format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
155
  )
156
 
157
- # Initialize FastAPI app
 
 
 
158
  app = FastAPI(
159
  title="LightRAG API",
160
- description="API for querying text using LightRAG with separate storage and input directories",
 
 
 
 
 
161
  )
162
 
 
 
 
 
 
 
 
 
 
 
 
 
163
  # Create working directory if it doesn't exist
164
  Path(args.working_dir).mkdir(parents=True, exist_ok=True)
165
 
@@ -209,7 +267,7 @@ def create_app(args):
209
  except Exception as e:
210
  logging.error(f"Error during startup indexing: {str(e)}")
211
 
212
- @app.post("/documents/scan")
213
  async def scan_for_new_documents():
214
  """Manually trigger scanning for new documents"""
215
  try:
@@ -234,7 +292,7 @@ def create_app(args):
234
  except Exception as e:
235
  raise HTTPException(status_code=500, detail=str(e))
236
 
237
- @app.post("/documents/upload")
238
  async def upload_to_input_dir(file: UploadFile = File(...)):
239
  """Upload a file to the input directory"""
240
  try:
@@ -262,7 +320,9 @@ def create_app(args):
262
  except Exception as e:
263
  raise HTTPException(status_code=500, detail=str(e))
264
 
265
- @app.post("/query", response_model=QueryResponse)
 
 
266
  async def query_text(request: QueryRequest):
267
  try:
268
  response = await rag.aquery(
@@ -284,7 +344,7 @@ def create_app(args):
284
  except Exception as e:
285
  raise HTTPException(status_code=500, detail=str(e))
286
 
287
- @app.post("/query/stream")
288
  async def query_text_stream(request: QueryRequest):
289
  try:
290
  response = rag.query(
@@ -304,7 +364,11 @@ def create_app(args):
304
  except Exception as e:
305
  raise HTTPException(status_code=500, detail=str(e))
306
 
307
- @app.post("/documents/text", response_model=InsertResponse)
 
 
 
 
308
  async def insert_text(request: InsertTextRequest):
309
  try:
310
  await rag.ainsert(request.text)
@@ -316,7 +380,11 @@ def create_app(args):
316
  except Exception as e:
317
  raise HTTPException(status_code=500, detail=str(e))
318
 
319
- @app.post("/documents/file", response_model=InsertResponse)
 
 
 
 
320
  async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
321
  try:
322
  content = await file.read()
@@ -340,7 +408,11 @@ def create_app(args):
340
  except Exception as e:
341
  raise HTTPException(status_code=500, detail=str(e))
342
 
343
- @app.post("/documents/batch", response_model=InsertResponse)
 
 
 
 
344
  async def insert_batch(files: List[UploadFile] = File(...)):
345
  try:
346
  inserted_count = 0
@@ -370,7 +442,11 @@ def create_app(args):
370
  except Exception as e:
371
  raise HTTPException(status_code=500, detail=str(e))
372
 
373
- @app.delete("/documents", response_model=InsertResponse)
 
 
 
 
374
  async def clear_documents():
375
  try:
376
  rag.text_chunks = []
@@ -384,7 +460,7 @@ def create_app(args):
384
  except Exception as e:
385
  raise HTTPException(status_code=500, detail=str(e))
386
 
387
- @app.get("/health")
388
  async def get_status():
389
  """Get current system status"""
390
  return {
 
11
  import shutil
12
  import aiofiles
13
  from ascii_colors import trace_exception
14
+ import os
15
+
16
+ from fastapi import Depends, Security
17
+ from fastapi.security import APIKeyHeader
18
+ from fastapi.middleware.cors import CORSMiddleware
19
+
20
+ from starlette.status import HTTP_403_FORBIDDEN
21
 
22
 
23
  def parse_args():
 
92
  choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
93
  help="Logging level (default: INFO)",
94
  )
95
+ parser.add_argument(
96
+ "--key",
97
+ type=str,
98
+ help="API key for authentication. This protects lightrag server against unauthorized access",
99
+ default=None,
100
+ )
101
 
102
  return parser.parse_args()
103
 
 
161
  document_count: int
162
 
163
 
164
+ def get_api_key_dependency(api_key: Optional[str]):
165
+ if not api_key:
166
+ # If no API key is configured, return a dummy dependency that always succeeds
167
+ async def no_auth():
168
+ return None
169
+
170
+ return no_auth
171
+
172
+ # If API key is configured, use proper authentication
173
+ api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
174
+
175
+ async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
176
+ if not api_key_header_value:
177
+ raise HTTPException(
178
+ status_code=HTTP_403_FORBIDDEN, detail="API Key required"
179
+ )
180
+ if api_key_header_value != api_key:
181
+ raise HTTPException(
182
+ status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
183
+ )
184
+ return api_key_header_value
185
+
186
+ return api_key_auth
187
+
188
+
189
  def create_app(args):
190
  # Setup logging
191
  logging.basicConfig(
192
  format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
193
  )
194
 
195
+ # Check if API key is provided either through env var or args
196
+ api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
197
+
198
+ # Initialize FastAPI
199
  app = FastAPI(
200
  title="LightRAG API",
201
+ description="API for querying text using LightRAG with separate storage and input directories"
202
+ + "(With authentication)"
203
+ if api_key
204
+ else "",
205
+ version="1.0.0",
206
+ openapi_tags=[{"name": "api"}],
207
  )
208
 
209
+ # Add CORS middleware
210
+ app.add_middleware(
211
+ CORSMiddleware,
212
+ allow_origins=["*"],
213
+ allow_credentials=True,
214
+ allow_methods=["*"],
215
+ allow_headers=["*"],
216
+ )
217
+
218
+ # Create the optional API key dependency
219
+ optional_api_key = get_api_key_dependency(api_key)
220
+
221
  # Create working directory if it doesn't exist
222
  Path(args.working_dir).mkdir(parents=True, exist_ok=True)
223
 
 
267
  except Exception as e:
268
  logging.error(f"Error during startup indexing: {str(e)}")
269
 
270
+ @app.post("/documents/scan", dependencies=[Depends(optional_api_key)])
271
  async def scan_for_new_documents():
272
  """Manually trigger scanning for new documents"""
273
  try:
 
292
  except Exception as e:
293
  raise HTTPException(status_code=500, detail=str(e))
294
 
295
+ @app.post("/documents/upload", dependencies=[Depends(optional_api_key)])
296
  async def upload_to_input_dir(file: UploadFile = File(...)):
297
  """Upload a file to the input directory"""
298
  try:
 
320
  except Exception as e:
321
  raise HTTPException(status_code=500, detail=str(e))
322
 
323
+ @app.post(
324
+ "/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
325
+ )
326
  async def query_text(request: QueryRequest):
327
  try:
328
  response = await rag.aquery(
 
344
  except Exception as e:
345
  raise HTTPException(status_code=500, detail=str(e))
346
 
347
+ @app.post("/query/stream", dependencies=[Depends(optional_api_key)])
348
  async def query_text_stream(request: QueryRequest):
349
  try:
350
  response = rag.query(
 
364
  except Exception as e:
365
  raise HTTPException(status_code=500, detail=str(e))
366
 
367
+ @app.post(
368
+ "/documents/text",
369
+ response_model=InsertResponse,
370
+ dependencies=[Depends(optional_api_key)],
371
+ )
372
  async def insert_text(request: InsertTextRequest):
373
  try:
374
  await rag.ainsert(request.text)
 
380
  except Exception as e:
381
  raise HTTPException(status_code=500, detail=str(e))
382
 
383
+ @app.post(
384
+ "/documents/file",
385
+ response_model=InsertResponse,
386
+ dependencies=[Depends(optional_api_key)],
387
+ )
388
  async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
389
  try:
390
  content = await file.read()
 
408
  except Exception as e:
409
  raise HTTPException(status_code=500, detail=str(e))
410
 
411
+ @app.post(
412
+ "/documents/batch",
413
+ response_model=InsertResponse,
414
+ dependencies=[Depends(optional_api_key)],
415
+ )
416
  async def insert_batch(files: List[UploadFile] = File(...)):
417
  try:
418
  inserted_count = 0
 
442
  except Exception as e:
443
  raise HTTPException(status_code=500, detail=str(e))
444
 
445
+ @app.delete(
446
+ "/documents",
447
+ response_model=InsertResponse,
448
+ dependencies=[Depends(optional_api_key)],
449
+ )
450
  async def clear_documents():
451
  try:
452
  rag.text_chunks = []
 
460
  except Exception as e:
461
  raise HTTPException(status_code=500, detail=str(e))
462
 
463
+ @app.get("/health", dependencies=[Depends(optional_api_key)])
464
  async def get_status():
465
  """Get current system status"""
466
  return {
lightrag/api/openai_lightrag_server.py CHANGED
@@ -14,6 +14,14 @@ import aiofiles
14
  from ascii_colors import trace_exception
15
  import nest_asyncio
16
 
 
 
 
 
 
 
 
 
17
  # Apply nest_asyncio to solve event loop issues
18
  nest_asyncio.apply()
19
 
@@ -75,6 +83,13 @@ def parse_args():
75
  help="Logging level (default: INFO)",
76
  )
77
 
 
 
 
 
 
 
 
78
  return parser.parse_args()
79
 
80
 
@@ -137,6 +152,31 @@ class InsertResponse(BaseModel):
137
  document_count: int
138
 
139
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
140
  async def get_embedding_dim(embedding_model: str) -> int:
141
  """Get embedding dimensions for the specified model"""
142
  test_text = ["This is a test sentence."]
@@ -150,10 +190,39 @@ def create_app(args):
150
  format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
151
  )
152
 
153
- # Initialize FastAPI app
 
 
 
154
  app = FastAPI(
155
  title="LightRAG API",
156
- description="API for querying text using LightRAG with OpenAI integration",
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
157
  )
158
 
159
  # Create working directory if it doesn't exist
@@ -213,7 +282,7 @@ def create_app(args):
213
  except Exception as e:
214
  logging.error(f"Error during startup indexing: {str(e)}")
215
 
216
- @app.post("/documents/scan")
217
  async def scan_for_new_documents():
218
  """Manually trigger scanning for new documents"""
219
  try:
@@ -238,7 +307,7 @@ def create_app(args):
238
  except Exception as e:
239
  raise HTTPException(status_code=500, detail=str(e))
240
 
241
- @app.post("/documents/upload")
242
  async def upload_to_input_dir(file: UploadFile = File(...)):
243
  """Upload a file to the input directory"""
244
  try:
@@ -266,7 +335,9 @@ def create_app(args):
266
  except Exception as e:
267
  raise HTTPException(status_code=500, detail=str(e))
268
 
269
- @app.post("/query", response_model=QueryResponse)
 
 
270
  async def query_text(request: QueryRequest):
271
  try:
272
  response = await rag.aquery(
@@ -288,7 +359,7 @@ def create_app(args):
288
  except Exception as e:
289
  raise HTTPException(status_code=500, detail=str(e))
290
 
291
- @app.post("/query/stream")
292
  async def query_text_stream(request: QueryRequest):
293
  try:
294
  response = rag.query(
@@ -308,7 +379,11 @@ def create_app(args):
308
  except Exception as e:
309
  raise HTTPException(status_code=500, detail=str(e))
310
 
311
- @app.post("/documents/text", response_model=InsertResponse)
 
 
 
 
312
  async def insert_text(request: InsertTextRequest):
313
  try:
314
  rag.insert(request.text)
@@ -320,7 +395,11 @@ def create_app(args):
320
  except Exception as e:
321
  raise HTTPException(status_code=500, detail=str(e))
322
 
323
- @app.post("/documents/file", response_model=InsertResponse)
 
 
 
 
324
  async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
325
  try:
326
  content = await file.read()
@@ -344,7 +423,11 @@ def create_app(args):
344
  except Exception as e:
345
  raise HTTPException(status_code=500, detail=str(e))
346
 
347
- @app.post("/documents/batch", response_model=InsertResponse)
 
 
 
 
348
  async def insert_batch(files: List[UploadFile] = File(...)):
349
  try:
350
  inserted_count = 0
@@ -374,7 +457,11 @@ def create_app(args):
374
  except Exception as e:
375
  raise HTTPException(status_code=500, detail=str(e))
376
 
377
- @app.delete("/documents", response_model=InsertResponse)
 
 
 
 
378
  async def clear_documents():
379
  try:
380
  rag.text_chunks = []
@@ -388,7 +475,7 @@ def create_app(args):
388
  except Exception as e:
389
  raise HTTPException(status_code=500, detail=str(e))
390
 
391
- @app.get("/health")
392
  async def get_status():
393
  """Get current system status"""
394
  return {
 
14
  from ascii_colors import trace_exception
15
  import nest_asyncio
16
 
17
+ import os
18
+
19
+ from fastapi import Depends, Security
20
+ from fastapi.security import APIKeyHeader
21
+ from fastapi.middleware.cors import CORSMiddleware
22
+
23
+ from starlette.status import HTTP_403_FORBIDDEN
24
+
25
  # Apply nest_asyncio to solve event loop issues
26
  nest_asyncio.apply()
27
 
 
83
  help="Logging level (default: INFO)",
84
  )
85
 
86
+ parser.add_argument(
87
+ "--key",
88
+ type=str,
89
+ help="API key for authentication. This protects lightrag server against unauthorized access",
90
+ default=None,
91
+ )
92
+
93
  return parser.parse_args()
94
 
95
 
 
152
  document_count: int
153
 
154
 
155
+ def get_api_key_dependency(api_key: Optional[str]):
156
+ if not api_key:
157
+ # If no API key is configured, return a dummy dependency that always succeeds
158
+ async def no_auth():
159
+ return None
160
+
161
+ return no_auth
162
+
163
+ # If API key is configured, use proper authentication
164
+ api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
165
+
166
+ async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
167
+ if not api_key_header_value:
168
+ raise HTTPException(
169
+ status_code=HTTP_403_FORBIDDEN, detail="API Key required"
170
+ )
171
+ if api_key_header_value != api_key:
172
+ raise HTTPException(
173
+ status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
174
+ )
175
+ return api_key_header_value
176
+
177
+ return api_key_auth
178
+
179
+
180
  async def get_embedding_dim(embedding_model: str) -> int:
181
  """Get embedding dimensions for the specified model"""
182
  test_text = ["This is a test sentence."]
 
190
  format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
191
  )
192
 
193
+ # Check if API key is provided either through env var or args
194
+ api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
195
+
196
+ # Initialize FastAPI
197
  app = FastAPI(
198
  title="LightRAG API",
199
+ description="API for querying text using LightRAG with separate storage and input directories"
200
+ + "(With authentication)"
201
+ if api_key
202
+ else "",
203
+ version="1.0.0",
204
+ openapi_tags=[{"name": "api"}],
205
+ )
206
+
207
+ # Add CORS middleware
208
+ app.add_middleware(
209
+ CORSMiddleware,
210
+ allow_origins=["*"],
211
+ allow_credentials=True,
212
+ allow_methods=["*"],
213
+ allow_headers=["*"],
214
+ )
215
+
216
+ # Create the optional API key dependency
217
+ optional_api_key = get_api_key_dependency(api_key)
218
+
219
+ # Add CORS middleware
220
+ app.add_middleware(
221
+ CORSMiddleware,
222
+ allow_origins=["*"],
223
+ allow_credentials=True,
224
+ allow_methods=["*"],
225
+ allow_headers=["*"],
226
  )
227
 
228
  # Create working directory if it doesn't exist
 
282
  except Exception as e:
283
  logging.error(f"Error during startup indexing: {str(e)}")
284
 
285
+ @app.post("/documents/scan", dependencies=[Depends(optional_api_key)])
286
  async def scan_for_new_documents():
287
  """Manually trigger scanning for new documents"""
288
  try:
 
307
  except Exception as e:
308
  raise HTTPException(status_code=500, detail=str(e))
309
 
310
+ @app.post("/documents/upload", dependencies=[Depends(optional_api_key)])
311
  async def upload_to_input_dir(file: UploadFile = File(...)):
312
  """Upload a file to the input directory"""
313
  try:
 
335
  except Exception as e:
336
  raise HTTPException(status_code=500, detail=str(e))
337
 
338
+ @app.post(
339
+ "/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
340
+ )
341
  async def query_text(request: QueryRequest):
342
  try:
343
  response = await rag.aquery(
 
359
  except Exception as e:
360
  raise HTTPException(status_code=500, detail=str(e))
361
 
362
+ @app.post("/query/stream", dependencies=[Depends(optional_api_key)])
363
  async def query_text_stream(request: QueryRequest):
364
  try:
365
  response = rag.query(
 
379
  except Exception as e:
380
  raise HTTPException(status_code=500, detail=str(e))
381
 
382
+ @app.post(
383
+ "/documents/text",
384
+ response_model=InsertResponse,
385
+ dependencies=[Depends(optional_api_key)],
386
+ )
387
  async def insert_text(request: InsertTextRequest):
388
  try:
389
  rag.insert(request.text)
 
395
  except Exception as e:
396
  raise HTTPException(status_code=500, detail=str(e))
397
 
398
+ @app.post(
399
+ "/documents/file",
400
+ response_model=InsertResponse,
401
+ dependencies=[Depends(optional_api_key)],
402
+ )
403
  async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
404
  try:
405
  content = await file.read()
 
423
  except Exception as e:
424
  raise HTTPException(status_code=500, detail=str(e))
425
 
426
+ @app.post(
427
+ "/documents/batch",
428
+ response_model=InsertResponse,
429
+ dependencies=[Depends(optional_api_key)],
430
+ )
431
  async def insert_batch(files: List[UploadFile] = File(...)):
432
  try:
433
  inserted_count = 0
 
457
  except Exception as e:
458
  raise HTTPException(status_code=500, detail=str(e))
459
 
460
+ @app.delete(
461
+ "/documents",
462
+ response_model=InsertResponse,
463
+ dependencies=[Depends(optional_api_key)],
464
+ )
465
  async def clear_documents():
466
  try:
467
  rag.text_chunks = []
 
475
  except Exception as e:
476
  raise HTTPException(status_code=500, detail=str(e))
477
 
478
+ @app.get("/health", dependencies=[Depends(optional_api_key)])
479
  async def get_status():
480
  """Get current system status"""
481
  return {