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, Request
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, 'value'):
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, Depends, Request, status
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 if api_key_header is None else Security(api_key_header),
 
 
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 api_key_configured and api_key_header_value and api_key_header_value == api_key:
 
 
 
 
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(api_key_header, description="API Key for authentication"),
 
 
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