yangdx
commited on
Commit
·
d2a7d19
1
Parent(s):
37b1fa9
Fix linting
Browse files
lightrag/api/lightrag_server.py
CHANGED
@@ -2,7 +2,7 @@
|
|
2 |
LightRAG FastAPI Server
|
3 |
"""
|
4 |
|
5 |
-
from fastapi import FastAPI, Depends, HTTPException, status
|
6 |
import asyncio
|
7 |
import os
|
8 |
import logging
|
@@ -18,7 +18,6 @@ from fastapi.middleware.cors import CORSMiddleware
|
|
18 |
from contextlib import asynccontextmanager
|
19 |
from dotenv import load_dotenv
|
20 |
from lightrag.api.utils_api import (
|
21 |
-
get_api_key_dependency,
|
22 |
get_combined_auth_dependency,
|
23 |
parse_args,
|
24 |
get_default_host,
|
@@ -149,14 +148,14 @@ def create_app(args):
|
|
149 |
"openapi_tags": [{"name": "api"}],
|
150 |
"lifespan": lifespan,
|
151 |
}
|
152 |
-
|
153 |
# Configure Swagger UI parameters
|
154 |
# Enable persistAuthorization and tryItOutEnabled for better user experience
|
155 |
app_kwargs["swagger_ui_parameters"] = {
|
156 |
"persistAuthorization": True,
|
157 |
"tryItOutEnabled": True,
|
158 |
}
|
159 |
-
|
160 |
app = FastAPI(**app_kwargs)
|
161 |
|
162 |
def get_cors_origins():
|
|
|
2 |
LightRAG FastAPI Server
|
3 |
"""
|
4 |
|
5 |
+
from fastapi import FastAPI, Depends, HTTPException, status
|
6 |
import asyncio
|
7 |
import os
|
8 |
import logging
|
|
|
18 |
from contextlib import asynccontextmanager
|
19 |
from dotenv import load_dotenv
|
20 |
from lightrag.api.utils_api import (
|
|
|
21 |
get_combined_auth_dependency,
|
22 |
parse_args,
|
23 |
get_default_host,
|
|
|
148 |
"openapi_tags": [{"name": "api"}],
|
149 |
"lifespan": lifespan,
|
150 |
}
|
151 |
+
|
152 |
# Configure Swagger UI parameters
|
153 |
# Enable persistAuthorization and tryItOutEnabled for better user experience
|
154 |
app_kwargs["swagger_ui_parameters"] = {
|
155 |
"persistAuthorization": True,
|
156 |
"tryItOutEnabled": True,
|
157 |
}
|
158 |
+
|
159 |
app = FastAPI(**app_kwargs)
|
160 |
|
161 |
def get_cors_origins():
|
lightrag/api/routers/document_routes.py
CHANGED
@@ -808,14 +808,14 @@ def create_document_routes(
|
|
808 |
|
809 |
# Get update flags status for all namespaces
|
810 |
update_status = await get_all_update_flags_status()
|
811 |
-
|
812 |
# Convert MutableBoolean objects to regular boolean values
|
813 |
processed_update_status = {}
|
814 |
for namespace, flags in update_status.items():
|
815 |
processed_flags = []
|
816 |
for flag in flags:
|
817 |
# Handle both multiprocess and single process cases
|
818 |
-
if hasattr(flag,
|
819 |
processed_flags.append(bool(flag.value))
|
820 |
else:
|
821 |
processed_flags.append(bool(flag))
|
|
|
808 |
|
809 |
# Get update flags status for all namespaces
|
810 |
update_status = await get_all_update_flags_status()
|
811 |
+
|
812 |
# Convert MutableBoolean objects to regular boolean values
|
813 |
processed_update_status = {}
|
814 |
for namespace, flags in update_status.items():
|
815 |
processed_flags = []
|
816 |
for flag in flags:
|
817 |
# Handle both multiprocess and single process cases
|
818 |
+
if hasattr(flag, "value"):
|
819 |
processed_flags.append(bool(flag.value))
|
820 |
else:
|
821 |
processed_flags.append(bool(flag))
|
lightrag/api/routers/ollama_api.py
CHANGED
@@ -135,9 +135,7 @@ class OllamaAPI:
|
|
135 |
# Create combined auth dependency for Ollama API routes
|
136 |
combined_auth = get_combined_auth_dependency(self.api_key)
|
137 |
|
138 |
-
@self.router.get(
|
139 |
-
"/version", dependencies=[Depends(combined_auth)]
|
140 |
-
)
|
141 |
async def get_version():
|
142 |
"""Get Ollama version information"""
|
143 |
return OllamaVersionResponse(version="0.5.4")
|
@@ -165,9 +163,7 @@ class OllamaAPI:
|
|
165 |
]
|
166 |
)
|
167 |
|
168 |
-
@self.router.post(
|
169 |
-
"/generate", dependencies=[Depends(combined_auth)]
|
170 |
-
)
|
171 |
async def generate(raw_request: Request, request: OllamaGenerateRequest):
|
172 |
"""Handle generate completion requests acting as an Ollama model
|
173 |
For compatibility purpose, the request is not processed by LightRAG,
|
|
|
135 |
# Create combined auth dependency for Ollama API routes
|
136 |
combined_auth = get_combined_auth_dependency(self.api_key)
|
137 |
|
138 |
+
@self.router.get("/version", dependencies=[Depends(combined_auth)])
|
|
|
|
|
139 |
async def get_version():
|
140 |
"""Get Ollama version information"""
|
141 |
return OllamaVersionResponse(version="0.5.4")
|
|
|
163 |
]
|
164 |
)
|
165 |
|
166 |
+
@self.router.post("/generate", dependencies=[Depends(combined_auth)])
|
|
|
|
|
167 |
async def generate(raw_request: Request, request: OllamaGenerateRequest):
|
168 |
"""Handle generate completion requests acting as an Ollama model
|
169 |
For compatibility purpose, the request is not processed by LightRAG,
|
lightrag/api/utils_api.py
CHANGED
@@ -9,7 +9,7 @@ import sys
|
|
9 |
import logging
|
10 |
from ascii_colors import ASCIIColors
|
11 |
from lightrag.api import __api_version__
|
12 |
-
from fastapi import HTTPException, Security,
|
13 |
from dotenv import load_dotenv
|
14 |
from fastapi.security import APIKeyHeader, OAuth2PasswordBearer
|
15 |
from starlette.status import HTTP_403_FORBIDDEN
|
@@ -72,27 +72,25 @@ def get_combined_auth_dependency(api_key: Optional[str] = None):
|
|
72 |
|
73 |
# Only calculate api_key_configured as it depends on the function parameter
|
74 |
api_key_configured = bool(api_key)
|
75 |
-
|
76 |
# Create security dependencies with proper descriptions for Swagger UI
|
77 |
oauth2_scheme = OAuth2PasswordBearer(
|
78 |
-
tokenUrl="login",
|
79 |
-
auto_error=False,
|
80 |
-
description="OAuth2 Password Authentication"
|
81 |
)
|
82 |
-
|
83 |
# If API key is configured, create an API key header security
|
84 |
api_key_header = None
|
85 |
if api_key_configured:
|
86 |
api_key_header = APIKeyHeader(
|
87 |
-
name="X-API-Key",
|
88 |
-
auto_error=False,
|
89 |
-
description="API Key Authentication"
|
90 |
)
|
91 |
|
92 |
async def combined_dependency(
|
93 |
request: Request,
|
94 |
token: str = Security(oauth2_scheme),
|
95 |
-
api_key_header_value: Optional[str] = None
|
|
|
|
|
96 |
):
|
97 |
# 1. Check if path is in whitelist
|
98 |
path = request.url.path
|
@@ -106,11 +104,15 @@ def get_combined_auth_dependency(api_key: Optional[str] = None):
|
|
106 |
is_special_endpoint = path == "/health" or path.startswith("/api/")
|
107 |
if is_special_endpoint and not api_key_configured:
|
108 |
return # Special endpoint and no API key configured, allow access
|
109 |
-
|
110 |
# 3. Validate API key
|
111 |
-
if
|
|
|
|
|
|
|
|
|
112 |
return # API key validation successful
|
113 |
-
|
114 |
# 4. Validate token
|
115 |
if token:
|
116 |
try:
|
@@ -121,34 +123,33 @@ def get_combined_auth_dependency(api_key: Optional[str] = None):
|
|
121 |
# Accept non-guest token if auth is configured
|
122 |
if auth_configured and token_info.get("role") != "guest":
|
123 |
return
|
124 |
-
|
125 |
# Token validation failed, immediately return 401 error
|
126 |
raise HTTPException(
|
127 |
status_code=status.HTTP_401_UNAUTHORIZED,
|
128 |
-
detail="Invalid token. Please login again."
|
129 |
)
|
130 |
except HTTPException as e:
|
131 |
# If already a 401 error, re-raise it
|
132 |
if e.status_code == status.HTTP_401_UNAUTHORIZED:
|
133 |
raise
|
134 |
# For other exceptions, continue processing
|
135 |
-
|
136 |
# If token exists but validation failed (didn't return above), return 401
|
137 |
raise HTTPException(
|
138 |
status_code=status.HTTP_401_UNAUTHORIZED,
|
139 |
-
detail="Invalid token. Please login again."
|
140 |
)
|
141 |
-
|
142 |
# 5. No token and API key validation failed, return 403 error
|
143 |
if api_key_configured:
|
144 |
raise HTTPException(
|
145 |
status_code=HTTP_403_FORBIDDEN,
|
146 |
-
detail="API Key required or login authentication required."
|
147 |
)
|
148 |
else:
|
149 |
raise HTTPException(
|
150 |
-
status_code=HTTP_403_FORBIDDEN,
|
151 |
-
detail="Login authentication required."
|
152 |
)
|
153 |
|
154 |
return combined_dependency
|
@@ -183,7 +184,9 @@ def get_api_key_dependency(api_key: Optional[str]):
|
|
183 |
|
184 |
async def api_key_auth(
|
185 |
request: Request,
|
186 |
-
api_key_header_value: Optional[str] = Security(
|
|
|
|
|
187 |
):
|
188 |
# Check if request path is in whitelist
|
189 |
path = request.url.path
|
|
|
9 |
import logging
|
10 |
from ascii_colors import ASCIIColors
|
11 |
from lightrag.api import __api_version__
|
12 |
+
from fastapi import HTTPException, Security, Request, status
|
13 |
from dotenv import load_dotenv
|
14 |
from fastapi.security import APIKeyHeader, OAuth2PasswordBearer
|
15 |
from starlette.status import HTTP_403_FORBIDDEN
|
|
|
72 |
|
73 |
# Only calculate api_key_configured as it depends on the function parameter
|
74 |
api_key_configured = bool(api_key)
|
75 |
+
|
76 |
# Create security dependencies with proper descriptions for Swagger UI
|
77 |
oauth2_scheme = OAuth2PasswordBearer(
|
78 |
+
tokenUrl="login", auto_error=False, description="OAuth2 Password Authentication"
|
|
|
|
|
79 |
)
|
80 |
+
|
81 |
# If API key is configured, create an API key header security
|
82 |
api_key_header = None
|
83 |
if api_key_configured:
|
84 |
api_key_header = APIKeyHeader(
|
85 |
+
name="X-API-Key", auto_error=False, description="API Key Authentication"
|
|
|
|
|
86 |
)
|
87 |
|
88 |
async def combined_dependency(
|
89 |
request: Request,
|
90 |
token: str = Security(oauth2_scheme),
|
91 |
+
api_key_header_value: Optional[str] = None
|
92 |
+
if api_key_header is None
|
93 |
+
else Security(api_key_header),
|
94 |
):
|
95 |
# 1. Check if path is in whitelist
|
96 |
path = request.url.path
|
|
|
104 |
is_special_endpoint = path == "/health" or path.startswith("/api/")
|
105 |
if is_special_endpoint and not api_key_configured:
|
106 |
return # Special endpoint and no API key configured, allow access
|
107 |
+
|
108 |
# 3. Validate API key
|
109 |
+
if (
|
110 |
+
api_key_configured
|
111 |
+
and api_key_header_value
|
112 |
+
and api_key_header_value == api_key
|
113 |
+
):
|
114 |
return # API key validation successful
|
115 |
+
|
116 |
# 4. Validate token
|
117 |
if token:
|
118 |
try:
|
|
|
123 |
# Accept non-guest token if auth is configured
|
124 |
if auth_configured and token_info.get("role") != "guest":
|
125 |
return
|
126 |
+
|
127 |
# Token validation failed, immediately return 401 error
|
128 |
raise HTTPException(
|
129 |
status_code=status.HTTP_401_UNAUTHORIZED,
|
130 |
+
detail="Invalid token. Please login again.",
|
131 |
)
|
132 |
except HTTPException as e:
|
133 |
# If already a 401 error, re-raise it
|
134 |
if e.status_code == status.HTTP_401_UNAUTHORIZED:
|
135 |
raise
|
136 |
# For other exceptions, continue processing
|
137 |
+
|
138 |
# If token exists but validation failed (didn't return above), return 401
|
139 |
raise HTTPException(
|
140 |
status_code=status.HTTP_401_UNAUTHORIZED,
|
141 |
+
detail="Invalid token. Please login again.",
|
142 |
)
|
143 |
+
|
144 |
# 5. No token and API key validation failed, return 403 error
|
145 |
if api_key_configured:
|
146 |
raise HTTPException(
|
147 |
status_code=HTTP_403_FORBIDDEN,
|
148 |
+
detail="API Key required or login authentication required.",
|
149 |
)
|
150 |
else:
|
151 |
raise HTTPException(
|
152 |
+
status_code=HTTP_403_FORBIDDEN, detail="Login authentication required."
|
|
|
153 |
)
|
154 |
|
155 |
return combined_dependency
|
|
|
184 |
|
185 |
async def api_key_auth(
|
186 |
request: Request,
|
187 |
+
api_key_header_value: Optional[str] = Security(
|
188 |
+
api_key_header, description="API Key for authentication"
|
189 |
+
),
|
190 |
):
|
191 |
# Check if request path is in whitelist
|
192 |
path = request.url.path
|