applyed linting
This commit is contained in:
@@ -1077,7 +1077,7 @@ Each server has its own specific configuration options:
|
|||||||
| --key | none | Access Key to protect the lightrag service |
|
| --key | none | Access Key to protect the lightrag service |
|
||||||
|
|
||||||
|
|
||||||
For protecting the server using an authentication key, you can also use an environment variable named `LIGHTRAG_API_KEY`.
|
For protecting the server using an authentication key, you can also use an environment variable named `LIGHTRAG_API_KEY`.
|
||||||
### Example Usage
|
### Example Usage
|
||||||
|
|
||||||
#### LoLLMs RAG Server
|
#### LoLLMs RAG Server
|
||||||
@@ -1088,6 +1088,10 @@ lollms-lightrag-server --model mistral-nemo --port 8080 --working-dir ./custom_r
|
|||||||
|
|
||||||
# Using specific models (ensure they are installed in your LoLLMs instance)
|
# Using specific models (ensure they are installed in your LoLLMs instance)
|
||||||
lollms-lightrag-server --model mistral-nemo:latest --embedding-model bge-m3 --embedding-dim 1024
|
lollms-lightrag-server --model mistral-nemo:latest --embedding-model bge-m3 --embedding-dim 1024
|
||||||
|
|
||||||
|
# Using specific models and an authentication key
|
||||||
|
lollms-lightrag-server --model mistral-nemo:latest --embedding-model bge-m3 --embedding-dim 1024 --key ky-mykey
|
||||||
|
|
||||||
```
|
```
|
||||||
|
|
||||||
#### Ollama RAG Server
|
#### Ollama RAG Server
|
||||||
|
@@ -20,19 +20,12 @@ from dotenv import load_dotenv
|
|||||||
import inspect
|
import inspect
|
||||||
import json
|
import json
|
||||||
from fastapi.responses import StreamingResponse
|
from fastapi.responses import StreamingResponse
|
||||||
from fastapi import FastAPI, HTTPException
|
|
||||||
import os
|
|
||||||
from typing import Optional
|
|
||||||
|
|
||||||
from fastapi import FastAPI, Depends, HTTPException, Security
|
from fastapi import Depends, Security
|
||||||
from fastapi.security import APIKeyHeader
|
from fastapi.security import APIKeyHeader
|
||||||
import os
|
|
||||||
import argparse
|
|
||||||
from typing import Optional
|
|
||||||
from fastapi.middleware.cors import CORSMiddleware
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
|
|
||||||
from starlette.status import HTTP_403_FORBIDDEN
|
from starlette.status import HTTP_403_FORBIDDEN
|
||||||
from fastapi import HTTPException
|
|
||||||
|
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
|
|
||||||
@@ -106,8 +99,12 @@ def parse_args():
|
|||||||
help="Logging level (default: INFO)",
|
help="Logging level (default: INFO)",
|
||||||
)
|
)
|
||||||
|
|
||||||
parser.add_argument('--key', type=str, help='API key for authentication. This protects lightrag server against unauthorized access', default=None)
|
parser.add_argument(
|
||||||
|
"--key",
|
||||||
|
type=str,
|
||||||
|
help="API key for authentication. This protects lightrag server against unauthorized access",
|
||||||
|
default=None,
|
||||||
|
)
|
||||||
|
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
@@ -170,29 +167,29 @@ class InsertResponse(BaseModel):
|
|||||||
message: str
|
message: str
|
||||||
document_count: int
|
document_count: int
|
||||||
|
|
||||||
|
|
||||||
def get_api_key_dependency(api_key: Optional[str]):
|
def get_api_key_dependency(api_key: Optional[str]):
|
||||||
if not api_key:
|
if not api_key:
|
||||||
# If no API key is configured, return a dummy dependency that always succeeds
|
# If no API key is configured, return a dummy dependency that always succeeds
|
||||||
async def no_auth():
|
async def no_auth():
|
||||||
return None
|
return None
|
||||||
|
|
||||||
return no_auth
|
return no_auth
|
||||||
|
|
||||||
# If API key is configured, use proper authentication
|
# If API key is configured, use proper authentication
|
||||||
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
|
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
|
||||||
|
|
||||||
async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
|
async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
|
||||||
if not api_key_header_value:
|
if not api_key_header_value:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=HTTP_403_FORBIDDEN,
|
status_code=HTTP_403_FORBIDDEN, detail="API Key required"
|
||||||
detail="API Key required"
|
|
||||||
)
|
)
|
||||||
if api_key_header_value != api_key:
|
if api_key_header_value != api_key:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=HTTP_403_FORBIDDEN,
|
status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
|
||||||
detail="Invalid API Key"
|
|
||||||
)
|
)
|
||||||
return api_key_header_value
|
return api_key_header_value
|
||||||
|
|
||||||
return api_key_auth
|
return api_key_auth
|
||||||
|
|
||||||
|
|
||||||
@@ -209,18 +206,20 @@ def create_app(args):
|
|||||||
format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
|
format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
# Check if API key is provided either through env var or args
|
# Check if API key is provided either through env var or args
|
||||||
api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
|
api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
|
||||||
|
|
||||||
# Initialize FastAPI
|
# Initialize FastAPI
|
||||||
app = FastAPI(
|
app = FastAPI(
|
||||||
title="LightRAG API",
|
title="LightRAG API",
|
||||||
description="API for querying text using LightRAG with separate storage and input directories"+"(With authentication)" if api_key else "",
|
description="API for querying text using LightRAG with separate storage and input directories"
|
||||||
|
+ "(With authentication)"
|
||||||
|
if api_key
|
||||||
|
else "",
|
||||||
version="1.0.0",
|
version="1.0.0",
|
||||||
openapi_tags=[{"name": "api"}]
|
openapi_tags=[{"name": "api"}],
|
||||||
)
|
)
|
||||||
|
|
||||||
# Add CORS middleware
|
# Add CORS middleware
|
||||||
app.add_middleware(
|
app.add_middleware(
|
||||||
CORSMiddleware,
|
CORSMiddleware,
|
||||||
@@ -363,7 +362,9 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
|
||||||
|
)
|
||||||
async def query_text(request: QueryRequest):
|
async def query_text(request: QueryRequest):
|
||||||
try:
|
try:
|
||||||
response = await rag.aquery(
|
response = await rag.aquery(
|
||||||
@@ -404,7 +405,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/text", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/text",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_text(request: InsertTextRequest):
|
async def insert_text(request: InsertTextRequest):
|
||||||
try:
|
try:
|
||||||
await rag.ainsert(request.text)
|
await rag.ainsert(request.text)
|
||||||
@@ -416,7 +421,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/file", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/file",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
|
async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
|
||||||
try:
|
try:
|
||||||
content = await file.read()
|
content = await file.read()
|
||||||
@@ -440,7 +449,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/batch", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/batch",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_batch(files: List[UploadFile] = File(...)):
|
async def insert_batch(files: List[UploadFile] = File(...)):
|
||||||
try:
|
try:
|
||||||
inserted_count = 0
|
inserted_count = 0
|
||||||
@@ -470,7 +483,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.delete("/documents", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.delete(
|
||||||
|
"/documents",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def clear_documents():
|
async def clear_documents():
|
||||||
try:
|
try:
|
||||||
rag.text_chunks = []
|
rag.text_chunks = []
|
||||||
|
@@ -11,19 +11,14 @@ from pathlib import Path
|
|||||||
import shutil
|
import shutil
|
||||||
import aiofiles
|
import aiofiles
|
||||||
from ascii_colors import trace_exception
|
from ascii_colors import trace_exception
|
||||||
from fastapi import FastAPI, HTTPException
|
|
||||||
import os
|
import os
|
||||||
from typing import Optional
|
|
||||||
|
|
||||||
from fastapi import FastAPI, Depends, HTTPException, Security
|
from fastapi import Depends, Security
|
||||||
from fastapi.security import APIKeyHeader
|
from fastapi.security import APIKeyHeader
|
||||||
import os
|
|
||||||
import argparse
|
|
||||||
from typing import Optional
|
|
||||||
from fastapi.middleware.cors import CORSMiddleware
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
|
|
||||||
from starlette.status import HTTP_403_FORBIDDEN
|
from starlette.status import HTTP_403_FORBIDDEN
|
||||||
from fastapi import HTTPException
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = argparse.ArgumentParser(
|
parser = argparse.ArgumentParser(
|
||||||
@@ -98,8 +93,12 @@ def parse_args():
|
|||||||
help="Logging level (default: INFO)",
|
help="Logging level (default: INFO)",
|
||||||
)
|
)
|
||||||
|
|
||||||
parser.add_argument('--key', type=str, help='API key for authentication. This protects lightrag server against unauthorized access', default=None)
|
parser.add_argument(
|
||||||
|
"--key",
|
||||||
|
type=str,
|
||||||
|
help="API key for authentication. This protects lightrag server against unauthorized access",
|
||||||
|
default=None,
|
||||||
|
)
|
||||||
|
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
@@ -162,29 +161,29 @@ class InsertResponse(BaseModel):
|
|||||||
message: str
|
message: str
|
||||||
document_count: int
|
document_count: int
|
||||||
|
|
||||||
|
|
||||||
def get_api_key_dependency(api_key: Optional[str]):
|
def get_api_key_dependency(api_key: Optional[str]):
|
||||||
if not api_key:
|
if not api_key:
|
||||||
# If no API key is configured, return a dummy dependency that always succeeds
|
# If no API key is configured, return a dummy dependency that always succeeds
|
||||||
async def no_auth():
|
async def no_auth():
|
||||||
return None
|
return None
|
||||||
|
|
||||||
return no_auth
|
return no_auth
|
||||||
|
|
||||||
# If API key is configured, use proper authentication
|
# If API key is configured, use proper authentication
|
||||||
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
|
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
|
||||||
|
|
||||||
async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
|
async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
|
||||||
if not api_key_header_value:
|
if not api_key_header_value:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=HTTP_403_FORBIDDEN,
|
status_code=HTTP_403_FORBIDDEN, detail="API Key required"
|
||||||
detail="API Key required"
|
|
||||||
)
|
)
|
||||||
if api_key_header_value != api_key:
|
if api_key_header_value != api_key:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=HTTP_403_FORBIDDEN,
|
status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
|
||||||
detail="Invalid API Key"
|
|
||||||
)
|
)
|
||||||
return api_key_header_value
|
return api_key_header_value
|
||||||
|
|
||||||
return api_key_auth
|
return api_key_auth
|
||||||
|
|
||||||
|
|
||||||
@@ -196,15 +195,18 @@ def create_app(args):
|
|||||||
|
|
||||||
# Check if API key is provided either through env var or args
|
# Check if API key is provided either through env var or args
|
||||||
api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
|
api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
|
||||||
|
|
||||||
# Initialize FastAPI
|
# Initialize FastAPI
|
||||||
app = FastAPI(
|
app = FastAPI(
|
||||||
title="LightRAG API",
|
title="LightRAG API",
|
||||||
description="API for querying text using LightRAG with separate storage and input directories"+"(With authentication)" if api_key else "",
|
description="API for querying text using LightRAG with separate storage and input directories"
|
||||||
|
+ "(With authentication)"
|
||||||
|
if api_key
|
||||||
|
else "",
|
||||||
version="1.0.0",
|
version="1.0.0",
|
||||||
openapi_tags=[{"name": "api"}]
|
openapi_tags=[{"name": "api"}],
|
||||||
)
|
)
|
||||||
|
|
||||||
# Add CORS middleware
|
# Add CORS middleware
|
||||||
app.add_middleware(
|
app.add_middleware(
|
||||||
CORSMiddleware,
|
CORSMiddleware,
|
||||||
@@ -319,7 +321,9 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
|
||||||
|
)
|
||||||
async def query_text(request: QueryRequest):
|
async def query_text(request: QueryRequest):
|
||||||
try:
|
try:
|
||||||
response = await rag.aquery(
|
response = await rag.aquery(
|
||||||
@@ -361,7 +365,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/text", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/text",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_text(request: InsertTextRequest):
|
async def insert_text(request: InsertTextRequest):
|
||||||
try:
|
try:
|
||||||
rag.insert(request.text)
|
rag.insert(request.text)
|
||||||
@@ -373,7 +381,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/file", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/file",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
|
async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
|
||||||
try:
|
try:
|
||||||
content = await file.read()
|
content = await file.read()
|
||||||
@@ -397,7 +409,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/batch", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/batch",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_batch(files: List[UploadFile] = File(...)):
|
async def insert_batch(files: List[UploadFile] = File(...)):
|
||||||
try:
|
try:
|
||||||
inserted_count = 0
|
inserted_count = 0
|
||||||
@@ -427,7 +443,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.delete("/documents", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.delete(
|
||||||
|
"/documents",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def clear_documents():
|
async def clear_documents():
|
||||||
try:
|
try:
|
||||||
rag.text_chunks = []
|
rag.text_chunks = []
|
||||||
|
@@ -11,19 +11,13 @@ from pathlib import Path
|
|||||||
import shutil
|
import shutil
|
||||||
import aiofiles
|
import aiofiles
|
||||||
from ascii_colors import trace_exception
|
from ascii_colors import trace_exception
|
||||||
from fastapi import FastAPI, HTTPException
|
|
||||||
import os
|
import os
|
||||||
from typing import Optional
|
|
||||||
|
|
||||||
from fastapi import FastAPI, Depends, HTTPException, Security
|
from fastapi import Depends, Security
|
||||||
from fastapi.security import APIKeyHeader
|
from fastapi.security import APIKeyHeader
|
||||||
import os
|
|
||||||
import argparse
|
|
||||||
from typing import Optional
|
|
||||||
from fastapi.middleware.cors import CORSMiddleware
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
|
|
||||||
from starlette.status import HTTP_403_FORBIDDEN
|
from starlette.status import HTTP_403_FORBIDDEN
|
||||||
from fastapi import HTTPException
|
|
||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
@@ -98,7 +92,12 @@ def parse_args():
|
|||||||
choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
|
choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
|
||||||
help="Logging level (default: INFO)",
|
help="Logging level (default: INFO)",
|
||||||
)
|
)
|
||||||
parser.add_argument('--key', type=str, help='API key for authentication. This protects lightrag server against unauthorized access', default=None)
|
parser.add_argument(
|
||||||
|
"--key",
|
||||||
|
type=str,
|
||||||
|
help="API key for authentication. This protects lightrag server against unauthorized access",
|
||||||
|
default=None,
|
||||||
|
)
|
||||||
|
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
@@ -161,29 +160,29 @@ class InsertResponse(BaseModel):
|
|||||||
message: str
|
message: str
|
||||||
document_count: int
|
document_count: int
|
||||||
|
|
||||||
|
|
||||||
def get_api_key_dependency(api_key: Optional[str]):
|
def get_api_key_dependency(api_key: Optional[str]):
|
||||||
if not api_key:
|
if not api_key:
|
||||||
# If no API key is configured, return a dummy dependency that always succeeds
|
# If no API key is configured, return a dummy dependency that always succeeds
|
||||||
async def no_auth():
|
async def no_auth():
|
||||||
return None
|
return None
|
||||||
|
|
||||||
return no_auth
|
return no_auth
|
||||||
|
|
||||||
# If API key is configured, use proper authentication
|
# If API key is configured, use proper authentication
|
||||||
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
|
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
|
||||||
|
|
||||||
async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
|
async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
|
||||||
if not api_key_header_value:
|
if not api_key_header_value:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=HTTP_403_FORBIDDEN,
|
status_code=HTTP_403_FORBIDDEN, detail="API Key required"
|
||||||
detail="API Key required"
|
|
||||||
)
|
)
|
||||||
if api_key_header_value != api_key:
|
if api_key_header_value != api_key:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=HTTP_403_FORBIDDEN,
|
status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
|
||||||
detail="Invalid API Key"
|
|
||||||
)
|
)
|
||||||
return api_key_header_value
|
return api_key_header_value
|
||||||
|
|
||||||
return api_key_auth
|
return api_key_auth
|
||||||
|
|
||||||
|
|
||||||
@@ -195,15 +194,18 @@ def create_app(args):
|
|||||||
|
|
||||||
# Check if API key is provided either through env var or args
|
# Check if API key is provided either through env var or args
|
||||||
api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
|
api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
|
||||||
|
|
||||||
# Initialize FastAPI
|
# Initialize FastAPI
|
||||||
app = FastAPI(
|
app = FastAPI(
|
||||||
title="LightRAG API",
|
title="LightRAG API",
|
||||||
description="API for querying text using LightRAG with separate storage and input directories"+"(With authentication)" if api_key else "",
|
description="API for querying text using LightRAG with separate storage and input directories"
|
||||||
|
+ "(With authentication)"
|
||||||
|
if api_key
|
||||||
|
else "",
|
||||||
version="1.0.0",
|
version="1.0.0",
|
||||||
openapi_tags=[{"name": "api"}]
|
openapi_tags=[{"name": "api"}],
|
||||||
)
|
)
|
||||||
|
|
||||||
# Add CORS middleware
|
# Add CORS middleware
|
||||||
app.add_middleware(
|
app.add_middleware(
|
||||||
CORSMiddleware,
|
CORSMiddleware,
|
||||||
@@ -216,7 +218,6 @@ def create_app(args):
|
|||||||
# Create the optional API key dependency
|
# Create the optional API key dependency
|
||||||
optional_api_key = get_api_key_dependency(api_key)
|
optional_api_key = get_api_key_dependency(api_key)
|
||||||
|
|
||||||
|
|
||||||
# Create working directory if it doesn't exist
|
# Create working directory if it doesn't exist
|
||||||
Path(args.working_dir).mkdir(parents=True, exist_ok=True)
|
Path(args.working_dir).mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
@@ -319,7 +320,9 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
|
||||||
|
)
|
||||||
async def query_text(request: QueryRequest):
|
async def query_text(request: QueryRequest):
|
||||||
try:
|
try:
|
||||||
response = await rag.aquery(
|
response = await rag.aquery(
|
||||||
@@ -361,7 +364,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/text", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/text",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_text(request: InsertTextRequest):
|
async def insert_text(request: InsertTextRequest):
|
||||||
try:
|
try:
|
||||||
await rag.ainsert(request.text)
|
await rag.ainsert(request.text)
|
||||||
@@ -373,7 +380,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/file", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/file",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
|
async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
|
||||||
try:
|
try:
|
||||||
content = await file.read()
|
content = await file.read()
|
||||||
@@ -397,7 +408,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/batch", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/batch",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_batch(files: List[UploadFile] = File(...)):
|
async def insert_batch(files: List[UploadFile] = File(...)):
|
||||||
try:
|
try:
|
||||||
inserted_count = 0
|
inserted_count = 0
|
||||||
@@ -427,7 +442,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.delete("/documents", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.delete(
|
||||||
|
"/documents",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def clear_documents():
|
async def clear_documents():
|
||||||
try:
|
try:
|
||||||
rag.text_chunks = []
|
rag.text_chunks = []
|
||||||
|
@@ -14,19 +14,13 @@ import aiofiles
|
|||||||
from ascii_colors import trace_exception
|
from ascii_colors import trace_exception
|
||||||
import nest_asyncio
|
import nest_asyncio
|
||||||
|
|
||||||
from fastapi import FastAPI, HTTPException
|
|
||||||
import os
|
import os
|
||||||
from typing import Optional
|
|
||||||
|
|
||||||
from fastapi import FastAPI, Depends, HTTPException, Security
|
from fastapi import Depends, Security
|
||||||
from fastapi.security import APIKeyHeader
|
from fastapi.security import APIKeyHeader
|
||||||
import os
|
|
||||||
import argparse
|
|
||||||
from typing import Optional
|
|
||||||
from fastapi.middleware.cors import CORSMiddleware
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
|
|
||||||
from starlette.status import HTTP_403_FORBIDDEN
|
from starlette.status import HTTP_403_FORBIDDEN
|
||||||
from fastapi import HTTPException
|
|
||||||
|
|
||||||
# Apply nest_asyncio to solve event loop issues
|
# Apply nest_asyncio to solve event loop issues
|
||||||
nest_asyncio.apply()
|
nest_asyncio.apply()
|
||||||
@@ -89,8 +83,12 @@ def parse_args():
|
|||||||
help="Logging level (default: INFO)",
|
help="Logging level (default: INFO)",
|
||||||
)
|
)
|
||||||
|
|
||||||
parser.add_argument('--key', type=str, help='API key for authentication. This protects lightrag server against unauthorized access', default=None)
|
parser.add_argument(
|
||||||
|
"--key",
|
||||||
|
type=str,
|
||||||
|
help="API key for authentication. This protects lightrag server against unauthorized access",
|
||||||
|
default=None,
|
||||||
|
)
|
||||||
|
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
@@ -153,29 +151,29 @@ class InsertResponse(BaseModel):
|
|||||||
message: str
|
message: str
|
||||||
document_count: int
|
document_count: int
|
||||||
|
|
||||||
|
|
||||||
def get_api_key_dependency(api_key: Optional[str]):
|
def get_api_key_dependency(api_key: Optional[str]):
|
||||||
if not api_key:
|
if not api_key:
|
||||||
# If no API key is configured, return a dummy dependency that always succeeds
|
# If no API key is configured, return a dummy dependency that always succeeds
|
||||||
async def no_auth():
|
async def no_auth():
|
||||||
return None
|
return None
|
||||||
|
|
||||||
return no_auth
|
return no_auth
|
||||||
|
|
||||||
# If API key is configured, use proper authentication
|
# If API key is configured, use proper authentication
|
||||||
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
|
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
|
||||||
|
|
||||||
async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
|
async def api_key_auth(api_key_header_value: str | None = Security(api_key_header)):
|
||||||
if not api_key_header_value:
|
if not api_key_header_value:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=HTTP_403_FORBIDDEN,
|
status_code=HTTP_403_FORBIDDEN, detail="API Key required"
|
||||||
detail="API Key required"
|
|
||||||
)
|
)
|
||||||
if api_key_header_value != api_key:
|
if api_key_header_value != api_key:
|
||||||
raise HTTPException(
|
raise HTTPException(
|
||||||
status_code=HTTP_403_FORBIDDEN,
|
status_code=HTTP_403_FORBIDDEN, detail="Invalid API Key"
|
||||||
detail="Invalid API Key"
|
|
||||||
)
|
)
|
||||||
return api_key_header_value
|
return api_key_header_value
|
||||||
|
|
||||||
return api_key_auth
|
return api_key_auth
|
||||||
|
|
||||||
|
|
||||||
@@ -192,18 +190,20 @@ def create_app(args):
|
|||||||
format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
|
format="%(levelname)s:%(message)s", level=getattr(logging, args.log_level)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
# Check if API key is provided either through env var or args
|
# Check if API key is provided either through env var or args
|
||||||
api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
|
api_key = os.getenv("LIGHTRAG_API_KEY") or args.key
|
||||||
|
|
||||||
# Initialize FastAPI
|
# Initialize FastAPI
|
||||||
app = FastAPI(
|
app = FastAPI(
|
||||||
title="LightRAG API",
|
title="LightRAG API",
|
||||||
description="API for querying text using LightRAG with separate storage and input directories"+"(With authentication)" if api_key else "",
|
description="API for querying text using LightRAG with separate storage and input directories"
|
||||||
|
+ "(With authentication)"
|
||||||
|
if api_key
|
||||||
|
else "",
|
||||||
version="1.0.0",
|
version="1.0.0",
|
||||||
openapi_tags=[{"name": "api"}]
|
openapi_tags=[{"name": "api"}],
|
||||||
)
|
)
|
||||||
|
|
||||||
# Add CORS middleware
|
# Add CORS middleware
|
||||||
app.add_middleware(
|
app.add_middleware(
|
||||||
CORSMiddleware,
|
CORSMiddleware,
|
||||||
@@ -335,7 +335,9 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/query", response_model=QueryResponse, dependencies=[Depends(optional_api_key)]
|
||||||
|
)
|
||||||
async def query_text(request: QueryRequest):
|
async def query_text(request: QueryRequest):
|
||||||
try:
|
try:
|
||||||
response = await rag.aquery(
|
response = await rag.aquery(
|
||||||
@@ -377,7 +379,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/text", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/text",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_text(request: InsertTextRequest):
|
async def insert_text(request: InsertTextRequest):
|
||||||
try:
|
try:
|
||||||
rag.insert(request.text)
|
rag.insert(request.text)
|
||||||
@@ -389,7 +395,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/file", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/file",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
|
async def insert_file(file: UploadFile = File(...), description: str = Form(None)):
|
||||||
try:
|
try:
|
||||||
content = await file.read()
|
content = await file.read()
|
||||||
@@ -413,7 +423,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.post("/documents/batch", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.post(
|
||||||
|
"/documents/batch",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def insert_batch(files: List[UploadFile] = File(...)):
|
async def insert_batch(files: List[UploadFile] = File(...)):
|
||||||
try:
|
try:
|
||||||
inserted_count = 0
|
inserted_count = 0
|
||||||
@@ -443,7 +457,11 @@ def create_app(args):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
raise HTTPException(status_code=500, detail=str(e))
|
raise HTTPException(status_code=500, detail=str(e))
|
||||||
|
|
||||||
@app.delete("/documents", response_model=InsertResponse, dependencies=[Depends(optional_api_key)])
|
@app.delete(
|
||||||
|
"/documents",
|
||||||
|
response_model=InsertResponse,
|
||||||
|
dependencies=[Depends(optional_api_key)],
|
||||||
|
)
|
||||||
async def clear_documents():
|
async def clear_documents():
|
||||||
try:
|
try:
|
||||||
rag.text_chunks = []
|
rag.text_chunks = []
|
||||||
|
Reference in New Issue
Block a user