Merge remote-tracking branch 'origin/main'
Browse files- README.md +9 -0
- lightrag/api/azure_openai_lightrag_server.py +88 -12
- lightrag/api/lollms_lightrag_server.py +88 -11
- lightrag/api/ollama_lightrag_server.py +87 -11
- lightrag/api/openai_lightrag_server.py +98 -11
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 |
-
#
|
|
|
|
|
|
|
172 |
app = FastAPI(
|
173 |
title="LightRAG API",
|
174 |
-
description="API for querying text using LightRAG with
|
|
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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 |
-
#
|
|
|
|
|
|
|
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(
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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 |
-
#
|
|
|
|
|
|
|
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(
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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 |
-
#
|
|
|
|
|
|
|
154 |
app = FastAPI(
|
155 |
title="LightRAG API",
|
156 |
-
description="API for querying text using LightRAG with
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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(
|
|
|
|
|
|
|
|
|
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 {
|