From 088fc19318430aa43ef9c7ff76ad61296f8c8666 Mon Sep 17 00:00:00 2001 From: Milin Date: Mon, 31 Mar 2025 11:19:47 +0800 Subject: [PATCH 1/2] feat(config): Refactor configuration management - Optimize JWT Auth module to load configuration via `global_args`. - Decouple configuration-related code from `utils_api.py`, and add a new `config.py` file for unified configuration management. - Adjust configuration import in `lightrag_server.py`, `auth.py`, and `document_routes.py` to be introduced through `global_args`. --- env.example | 4 +- lightrag/api/auth.py | 17 +- lightrag/api/config.py | 324 +++++++++++++++++++++++ lightrag/api/lightrag_server.py | 30 ++- lightrag/api/routers/document_routes.py | 18 +- lightrag/api/utils_api.py | 329 ++---------------------- 6 files changed, 376 insertions(+), 346 deletions(-) create mode 100644 lightrag/api/config.py diff --git a/env.example b/env.example index 8e54c812..4bf9a340 100644 --- a/env.example +++ b/env.example @@ -160,7 +160,9 @@ REDIS_URI=redis://localhost:6379 ### For JWT Auth AUTH_ACCOUNTS='admin:admin123,user1:pass456' # username:password,username:password TOKEN_SECRET=Your-Key-For-LightRAG-API-Server # JWT key -TOKEN_EXPIRE_HOURS=4 # expire duration +#TOKEN_EXPIRE_HOURS=4 # Expire duration, default 4 +#GUEST_TOKEN_EXPIRE_HOURS=2 # Guest expire duration, default 2 +#JWT_ALGORITHM=HS256 # JWT encode algorithm, default HS256 ### API-Key to access LightRAG Server API # LIGHTRAG_API_KEY=your-secure-api-key-here diff --git a/lightrag/api/auth.py b/lightrag/api/auth.py index 50648504..de02721b 100644 --- a/lightrag/api/auth.py +++ b/lightrag/api/auth.py @@ -1,9 +1,11 @@ -import os from datetime import datetime, timedelta + import jwt +from dotenv import load_dotenv from fastapi import HTTPException, status from pydantic import BaseModel -from dotenv import load_dotenv + +from .config import global_args load_dotenv() @@ -17,13 +19,12 @@ class TokenPayload(BaseModel): class AuthHandler: def __init__(self): - self.secret = os.getenv("TOKEN_SECRET", "4f85ds4f56dsf46") - self.algorithm = "HS256" - self.expire_hours = int(os.getenv("TOKEN_EXPIRE_HOURS", 4)) - self.guest_expire_hours = int(os.getenv("GUEST_TOKEN_EXPIRE_HOURS", 2)) - + self.secret = global_args.token_secret + self.algorithm = global_args.jwt_algorithm + self.expire_hours = global_args.token_expire_hours + self.guest_expire_hours = global_args.guest_token_expire_hours self.accounts = {} - auth_accounts = os.getenv("AUTH_ACCOUNTS") + auth_accounts = global_args.auth_accounts if auth_accounts: for account in auth_accounts.split(","): username, password = account.split(":", 1) diff --git a/lightrag/api/config.py b/lightrag/api/config.py new file mode 100644 index 00000000..bacf1b7f --- /dev/null +++ b/lightrag/api/config.py @@ -0,0 +1,324 @@ +""" +Configs for the LightRAG API. +""" + +import os +import argparse +import logging + + +class OllamaServerInfos: + # Constants for emulated Ollama model information + LIGHTRAG_NAME = "lightrag" + LIGHTRAG_TAG = os.getenv("OLLAMA_EMULATING_MODEL_TAG", "latest") + LIGHTRAG_MODEL = f"{LIGHTRAG_NAME}:{LIGHTRAG_TAG}" + LIGHTRAG_SIZE = 7365960935 # it's a dummy value + LIGHTRAG_CREATED_AT = "2024-01-15T00:00:00Z" + LIGHTRAG_DIGEST = "sha256:lightrag" + + +ollama_server_infos = OllamaServerInfos() + + +class DefaultRAGStorageConfig: + KV_STORAGE = "JsonKVStorage" + VECTOR_STORAGE = "NanoVectorDBStorage" + GRAPH_STORAGE = "NetworkXStorage" + DOC_STATUS_STORAGE = "JsonDocStatusStorage" + + +def get_default_host(binding_type: str) -> str: + default_hosts = { + "ollama": os.getenv("LLM_BINDING_HOST", "http://localhost:11434"), + "lollms": os.getenv("LLM_BINDING_HOST", "http://localhost:9600"), + "azure_openai": os.getenv("AZURE_OPENAI_ENDPOINT", "https://api.openai.com/v1"), + "openai": os.getenv("LLM_BINDING_HOST", "https://api.openai.com/v1"), + } + return default_hosts.get( + binding_type, os.getenv("LLM_BINDING_HOST", "http://localhost:11434") + ) # fallback to ollama if unknown + + +def get_env_value(env_key: str, default: any, value_type: type = str) -> any: + """ + Get value from environment variable with type conversion + + Args: + env_key (str): Environment variable key + default (any): Default value if env variable is not set + value_type (type): Type to convert the value to + + Returns: + any: Converted value from environment or default + """ + value = os.getenv(env_key) + if value is None: + return default + + if value_type is bool: + return value.lower() in ("true", "1", "yes", "t", "on") + try: + return value_type(value) + except ValueError: + return default + + +def parse_args() -> argparse.Namespace: + """ + Parse command line arguments with environment variable fallback + + Args: + is_uvicorn_mode: Whether running under uvicorn mode + + Returns: + argparse.Namespace: Parsed arguments + """ + + parser = argparse.ArgumentParser( + description="LightRAG FastAPI Server with separate working and input directories" + ) + + # Server configuration + parser.add_argument( + "--host", + default=get_env_value("HOST", "0.0.0.0"), + help="Server host (default: from env or 0.0.0.0)", + ) + parser.add_argument( + "--port", + type=int, + default=get_env_value("PORT", 9621, int), + help="Server port (default: from env or 9621)", + ) + + # Directory configuration + parser.add_argument( + "--working-dir", + default=get_env_value("WORKING_DIR", "./rag_storage"), + help="Working directory for RAG storage (default: from env or ./rag_storage)", + ) + parser.add_argument( + "--input-dir", + default=get_env_value("INPUT_DIR", "./inputs"), + help="Directory containing input documents (default: from env or ./inputs)", + ) + + def timeout_type(value): + if value is None: + return 150 + if value is None or value == "None": + return None + return int(value) + + parser.add_argument( + "--timeout", + default=get_env_value("TIMEOUT", None, timeout_type), + type=timeout_type, + help="Timeout in seconds (useful when using slow AI). Use None for infinite timeout", + ) + + # RAG configuration + parser.add_argument( + "--max-async", + type=int, + default=get_env_value("MAX_ASYNC", 4, int), + help="Maximum async operations (default: from env or 4)", + ) + parser.add_argument( + "--max-tokens", + type=int, + default=get_env_value("MAX_TOKENS", 32768, int), + help="Maximum token size (default: from env or 32768)", + ) + + # Logging configuration + parser.add_argument( + "--log-level", + default=get_env_value("LOG_LEVEL", "INFO"), + choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], + help="Logging level (default: from env or INFO)", + ) + parser.add_argument( + "--verbose", + action="store_true", + default=get_env_value("VERBOSE", False, bool), + help="Enable verbose debug output(only valid for DEBUG log-level)", + ) + + parser.add_argument( + "--key", + type=str, + default=get_env_value("LIGHTRAG_API_KEY", None), + help="API key for authentication. This protects lightrag server against unauthorized access", + ) + + # Optional https parameters + parser.add_argument( + "--ssl", + action="store_true", + default=get_env_value("SSL", False, bool), + help="Enable HTTPS (default: from env or False)", + ) + parser.add_argument( + "--ssl-certfile", + default=get_env_value("SSL_CERTFILE", None), + help="Path to SSL certificate file (required if --ssl is enabled)", + ) + parser.add_argument( + "--ssl-keyfile", + default=get_env_value("SSL_KEYFILE", None), + help="Path to SSL private key file (required if --ssl is enabled)", + ) + + parser.add_argument( + "--history-turns", + type=int, + default=get_env_value("HISTORY_TURNS", 3, int), + help="Number of conversation history turns to include (default: from env or 3)", + ) + + # Search parameters + parser.add_argument( + "--top-k", + type=int, + default=get_env_value("TOP_K", 60, int), + help="Number of most similar results to return (default: from env or 60)", + ) + parser.add_argument( + "--cosine-threshold", + type=float, + default=get_env_value("COSINE_THRESHOLD", 0.2, float), + help="Cosine similarity threshold (default: from env or 0.4)", + ) + + # Ollama model name + parser.add_argument( + "--simulated-model-name", + type=str, + default=get_env_value( + "SIMULATED_MODEL_NAME", ollama_server_infos.LIGHTRAG_MODEL + ), + help="Number of conversation history turns to include (default: from env or 3)", + ) + + # Namespace + parser.add_argument( + "--namespace-prefix", + type=str, + default=get_env_value("NAMESPACE_PREFIX", ""), + help="Prefix of the namespace", + ) + + parser.add_argument( + "--auto-scan-at-startup", + action="store_true", + default=False, + help="Enable automatic scanning when the program starts", + ) + + # Server workers configuration + parser.add_argument( + "--workers", + type=int, + default=get_env_value("WORKERS", 1, int), + help="Number of worker processes (default: from env or 1)", + ) + + # LLM and embedding bindings + parser.add_argument( + "--llm-binding", + type=str, + default=get_env_value("LLM_BINDING", "ollama"), + choices=["lollms", "ollama", "openai", "openai-ollama", "azure_openai"], + help="LLM binding type (default: from env or ollama)", + ) + parser.add_argument( + "--embedding-binding", + type=str, + default=get_env_value("EMBEDDING_BINDING", "ollama"), + choices=["lollms", "ollama", "openai", "azure_openai"], + help="Embedding binding type (default: from env or ollama)", + ) + + args = parser.parse_args() + + # convert relative path to absolute path + args.working_dir = os.path.abspath(args.working_dir) + args.input_dir = os.path.abspath(args.input_dir) + + # Inject storage configuration from environment variables + args.kv_storage = get_env_value( + "LIGHTRAG_KV_STORAGE", DefaultRAGStorageConfig.KV_STORAGE + ) + args.doc_status_storage = get_env_value( + "LIGHTRAG_DOC_STATUS_STORAGE", DefaultRAGStorageConfig.DOC_STATUS_STORAGE + ) + args.graph_storage = get_env_value( + "LIGHTRAG_GRAPH_STORAGE", DefaultRAGStorageConfig.GRAPH_STORAGE + ) + args.vector_storage = get_env_value( + "LIGHTRAG_VECTOR_STORAGE", DefaultRAGStorageConfig.VECTOR_STORAGE + ) + + # Get MAX_PARALLEL_INSERT from environment + args.max_parallel_insert = get_env_value("MAX_PARALLEL_INSERT", 2, int) + + # Handle openai-ollama special case + if args.llm_binding == "openai-ollama": + args.llm_binding = "openai" + args.embedding_binding = "ollama" + + args.llm_binding_host = get_env_value( + "LLM_BINDING_HOST", get_default_host(args.llm_binding) + ) + args.embedding_binding_host = get_env_value( + "EMBEDDING_BINDING_HOST", get_default_host(args.embedding_binding) + ) + args.llm_binding_api_key = get_env_value("LLM_BINDING_API_KEY", None) + args.embedding_binding_api_key = get_env_value("EMBEDDING_BINDING_API_KEY", "") + + # Inject model configuration + args.llm_model = get_env_value("LLM_MODEL", "mistral-nemo:latest") + args.embedding_model = get_env_value("EMBEDDING_MODEL", "bge-m3:latest") + args.embedding_dim = get_env_value("EMBEDDING_DIM", 1024, int) + args.max_embed_tokens = get_env_value("MAX_EMBED_TOKENS", 8192, int) + + # Inject chunk configuration + args.chunk_size = get_env_value("CHUNK_SIZE", 1200, int) + args.chunk_overlap_size = get_env_value("CHUNK_OVERLAP_SIZE", 100, int) + + # Inject LLM cache configuration + args.enable_llm_cache_for_extract = get_env_value( + "ENABLE_LLM_CACHE_FOR_EXTRACT", True, bool + ) + + # Inject LLM temperature configuration + args.temperature = get_env_value("TEMPERATURE", 0.5, float) + + # Select Document loading tool (DOCLING, DEFAULT) + args.document_loading_engine = get_env_value("DOCUMENT_LOADING_ENGINE", "DEFAULT") + + # For JWT Auth + args.auth_accounts = get_env_value("AUTH_ACCOUNTS", "") + args.token_secret = get_env_value("TOKEN_SECRET", "lightrag-jwt-default-secret") + args.token_expire_hours = get_env_value("TOKEN_EXPIRE_HOURS", 4, int) + args.guest_token_expire_hours = get_env_value("GUEST_TOKEN_EXPIRE_HOURS", 2, int) + args.jwt_algorithm = get_env_value("JWT_ALGORITHM", "HS256") + + ollama_server_infos.LIGHTRAG_MODEL = args.simulated_model_name + + return args + + +def update_uvicorn_mode_config(): + # If in uvicorn mode and workers > 1, force it to 1 and log warning + if global_args.workers > 1: + original_workers = global_args.workers + global_args.workers = 1 + # Log warning directly here + logging.warning( + f"In uvicorn mode, workers parameter was set to {original_workers}. Forcing workers=1" + ) + + +global_args = parse_args() diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py index fe0e416e..a00eab6c 100644 --- a/lightrag/api/lightrag_server.py +++ b/lightrag/api/lightrag_server.py @@ -19,11 +19,14 @@ from contextlib import asynccontextmanager from dotenv import load_dotenv from lightrag.api.utils_api import ( get_combined_auth_dependency, - parse_args, - get_default_host, display_splash_screen, check_env_file, ) +from .config import ( + global_args, + update_uvicorn_mode_config, + get_default_host, +) import sys from lightrag import LightRAG, __version__ as core_version from lightrag.api import __api_version__ @@ -489,7 +492,7 @@ def create_app(args): def get_application(args=None): """Factory function for creating the FastAPI application""" if args is None: - args = parse_args() + args = global_args return create_app(args) @@ -610,30 +613,31 @@ def main(): # Configure logging before parsing args configure_logging() - - args = parse_args(is_uvicorn_mode=True) - display_splash_screen(args) + update_uvicorn_mode_config() + display_splash_screen(global_args) # Create application instance directly instead of using factory function - app = create_app(args) + app = create_app(global_args) # Start Uvicorn in single process mode uvicorn_config = { "app": app, # Pass application instance directly instead of string path - "host": args.host, - "port": args.port, + "host": global_args.host, + "port": global_args.port, "log_config": None, # Disable default config } - if args.ssl: + if global_args.ssl: uvicorn_config.update( { - "ssl_certfile": args.ssl_certfile, - "ssl_keyfile": args.ssl_keyfile, + "ssl_certfile": global_args.ssl_certfile, + "ssl_keyfile": global_args.ssl_keyfile, } ) - print(f"Starting Uvicorn server in single-process mode on {args.host}:{args.port}") + print( + f"Starting Uvicorn server in single-process mode on {global_args.host}:{global_args.port}" + ) uvicorn.run(**uvicorn_config) diff --git a/lightrag/api/routers/document_routes.py b/lightrag/api/routers/document_routes.py index d6434e08..2a7873b8 100644 --- a/lightrag/api/routers/document_routes.py +++ b/lightrag/api/routers/document_routes.py @@ -16,10 +16,8 @@ from pydantic import BaseModel, Field, field_validator from lightrag import LightRAG from lightrag.base import DocProcessingStatus, DocStatus -from lightrag.api.utils_api import ( - get_combined_auth_dependency, - global_args, -) +from lightrag.api.utils_api import get_combined_auth_dependency +from ..config import global_args router = APIRouter( prefix="/documents", @@ -276,7 +274,7 @@ async def pipeline_enqueue_file(rag: LightRAG, file_path: Path) -> bool: ) return False case ".pdf": - if global_args["main_args"].document_loading_engine == "DOCLING": + if global_args.document_loading_engine == "DOCLING": if not pm.is_installed("docling"): # type: ignore pm.install("docling") from docling.document_converter import DocumentConverter # type: ignore @@ -295,7 +293,7 @@ async def pipeline_enqueue_file(rag: LightRAG, file_path: Path) -> bool: for page in reader.pages: content += page.extract_text() + "\n" case ".docx": - if global_args["main_args"].document_loading_engine == "DOCLING": + if global_args.document_loading_engine == "DOCLING": if not pm.is_installed("docling"): # type: ignore pm.install("docling") from docling.document_converter import DocumentConverter # type: ignore @@ -315,7 +313,7 @@ async def pipeline_enqueue_file(rag: LightRAG, file_path: Path) -> bool: [paragraph.text for paragraph in doc.paragraphs] ) case ".pptx": - if global_args["main_args"].document_loading_engine == "DOCLING": + if global_args.document_loading_engine == "DOCLING": if not pm.is_installed("docling"): # type: ignore pm.install("docling") from docling.document_converter import DocumentConverter # type: ignore @@ -336,7 +334,7 @@ async def pipeline_enqueue_file(rag: LightRAG, file_path: Path) -> bool: if hasattr(shape, "text"): content += shape.text + "\n" case ".xlsx": - if global_args["main_args"].document_loading_engine == "DOCLING": + if global_args.document_loading_engine == "DOCLING": if not pm.is_installed("docling"): # type: ignore pm.install("docling") from docling.document_converter import DocumentConverter # type: ignore @@ -476,8 +474,8 @@ async def run_scanning_process(rag: LightRAG, doc_manager: DocumentManager): if not new_files: return - # Get MAX_PARALLEL_INSERT from global_args["main_args"] - max_parallel = global_args["main_args"].max_parallel_insert + # Get MAX_PARALLEL_INSERT from global_args + max_parallel = global_args.max_parallel_insert # Calculate batch size as 2 * MAX_PARALLEL_INSERT batch_size = 2 * max_parallel diff --git a/lightrag/api/utils_api.py b/lightrag/api/utils_api.py index 32a453bf..62ceecce 100644 --- a/lightrag/api/utils_api.py +++ b/lightrag/api/utils_api.py @@ -2,18 +2,20 @@ Utility functions for the LightRAG API. """ -import os import argparse -from typing import Optional, List, Tuple +import os import sys +from typing import Optional, List, Tuple + from ascii_colors import ASCIIColors -import logging -from lightrag.api import __api_version__ -from fastapi import HTTPException, Security, Request, status from dotenv import load_dotenv +from fastapi import HTTPException, Security, Request, status from fastapi.security import APIKeyHeader, OAuth2PasswordBearer from starlette.status import HTTP_403_FORBIDDEN + +from lightrag.api import __api_version__ from .auth import auth_handler +from .config import ollama_server_infos from ..prompt import PROMPTS @@ -40,8 +42,6 @@ def check_env_file(): # Load environment variables load_dotenv() -global_args = {"main_args": None} - # Get whitelist paths from environment variable, only once during initialization default_whitelist = "/health,/api/*" whitelist_paths = os.getenv("WHITELIST_PATHS", default_whitelist).split(",") @@ -62,19 +62,6 @@ for path in whitelist_paths: auth_configured = bool(auth_handler.accounts) -class OllamaServerInfos: - # Constants for emulated Ollama model information - LIGHTRAG_NAME = "lightrag" - LIGHTRAG_TAG = os.getenv("OLLAMA_EMULATING_MODEL_TAG", "latest") - LIGHTRAG_MODEL = f"{LIGHTRAG_NAME}:{LIGHTRAG_TAG}" - LIGHTRAG_SIZE = 7365960935 # it's a dummy value - LIGHTRAG_CREATED_AT = "2024-01-15T00:00:00Z" - LIGHTRAG_DIGEST = "sha256:lightrag" - - -ollama_server_infos = OllamaServerInfos() - - def get_combined_auth_dependency(api_key: Optional[str] = None): """ Create a combined authentication dependency that implements authentication logic @@ -185,299 +172,6 @@ def get_combined_auth_dependency(api_key: Optional[str] = None): return combined_dependency -class DefaultRAGStorageConfig: - KV_STORAGE = "JsonKVStorage" - VECTOR_STORAGE = "NanoVectorDBStorage" - GRAPH_STORAGE = "NetworkXStorage" - DOC_STATUS_STORAGE = "JsonDocStatusStorage" - - -def get_default_host(binding_type: str) -> str: - default_hosts = { - "ollama": os.getenv("LLM_BINDING_HOST", "http://localhost:11434"), - "lollms": os.getenv("LLM_BINDING_HOST", "http://localhost:9600"), - "azure_openai": os.getenv("AZURE_OPENAI_ENDPOINT", "https://api.openai.com/v1"), - "openai": os.getenv("LLM_BINDING_HOST", "https://api.openai.com/v1"), - } - return default_hosts.get( - binding_type, os.getenv("LLM_BINDING_HOST", "http://localhost:11434") - ) # fallback to ollama if unknown - - -def get_env_value(env_key: str, default: any, value_type: type = str) -> any: - """ - Get value from environment variable with type conversion - - Args: - env_key (str): Environment variable key - default (any): Default value if env variable is not set - value_type (type): Type to convert the value to - - Returns: - any: Converted value from environment or default - """ - value = os.getenv(env_key) - if value is None: - return default - - if value_type is bool: - return value.lower() in ("true", "1", "yes", "t", "on") - try: - return value_type(value) - except ValueError: - return default - - -def parse_args(is_uvicorn_mode: bool = False) -> argparse.Namespace: - """ - Parse command line arguments with environment variable fallback - - Args: - is_uvicorn_mode: Whether running under uvicorn mode - - Returns: - argparse.Namespace: Parsed arguments - """ - - parser = argparse.ArgumentParser( - description="LightRAG FastAPI Server with separate working and input directories" - ) - - # Server configuration - parser.add_argument( - "--host", - default=get_env_value("HOST", "0.0.0.0"), - help="Server host (default: from env or 0.0.0.0)", - ) - parser.add_argument( - "--port", - type=int, - default=get_env_value("PORT", 9621, int), - help="Server port (default: from env or 9621)", - ) - - # Directory configuration - parser.add_argument( - "--working-dir", - default=get_env_value("WORKING_DIR", "./rag_storage"), - help="Working directory for RAG storage (default: from env or ./rag_storage)", - ) - parser.add_argument( - "--input-dir", - default=get_env_value("INPUT_DIR", "./inputs"), - help="Directory containing input documents (default: from env or ./inputs)", - ) - - def timeout_type(value): - if value is None: - return 150 - if value is None or value == "None": - return None - return int(value) - - parser.add_argument( - "--timeout", - default=get_env_value("TIMEOUT", None, timeout_type), - type=timeout_type, - help="Timeout in seconds (useful when using slow AI). Use None for infinite timeout", - ) - - # RAG configuration - parser.add_argument( - "--max-async", - type=int, - default=get_env_value("MAX_ASYNC", 4, int), - help="Maximum async operations (default: from env or 4)", - ) - parser.add_argument( - "--max-tokens", - type=int, - default=get_env_value("MAX_TOKENS", 32768, int), - help="Maximum token size (default: from env or 32768)", - ) - - # Logging configuration - parser.add_argument( - "--log-level", - default=get_env_value("LOG_LEVEL", "INFO"), - choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], - help="Logging level (default: from env or INFO)", - ) - parser.add_argument( - "--verbose", - action="store_true", - default=get_env_value("VERBOSE", False, bool), - help="Enable verbose debug output(only valid for DEBUG log-level)", - ) - - parser.add_argument( - "--key", - type=str, - default=get_env_value("LIGHTRAG_API_KEY", None), - help="API key for authentication. This protects lightrag server against unauthorized access", - ) - - # Optional https parameters - parser.add_argument( - "--ssl", - action="store_true", - default=get_env_value("SSL", False, bool), - help="Enable HTTPS (default: from env or False)", - ) - parser.add_argument( - "--ssl-certfile", - default=get_env_value("SSL_CERTFILE", None), - help="Path to SSL certificate file (required if --ssl is enabled)", - ) - parser.add_argument( - "--ssl-keyfile", - default=get_env_value("SSL_KEYFILE", None), - help="Path to SSL private key file (required if --ssl is enabled)", - ) - - parser.add_argument( - "--history-turns", - type=int, - default=get_env_value("HISTORY_TURNS", 3, int), - help="Number of conversation history turns to include (default: from env or 3)", - ) - - # Search parameters - parser.add_argument( - "--top-k", - type=int, - default=get_env_value("TOP_K", 60, int), - help="Number of most similar results to return (default: from env or 60)", - ) - parser.add_argument( - "--cosine-threshold", - type=float, - default=get_env_value("COSINE_THRESHOLD", 0.2, float), - help="Cosine similarity threshold (default: from env or 0.4)", - ) - - # Ollama model name - parser.add_argument( - "--simulated-model-name", - type=str, - default=get_env_value( - "SIMULATED_MODEL_NAME", ollama_server_infos.LIGHTRAG_MODEL - ), - help="Number of conversation history turns to include (default: from env or 3)", - ) - - # Namespace - parser.add_argument( - "--namespace-prefix", - type=str, - default=get_env_value("NAMESPACE_PREFIX", ""), - help="Prefix of the namespace", - ) - - parser.add_argument( - "--auto-scan-at-startup", - action="store_true", - default=False, - help="Enable automatic scanning when the program starts", - ) - - # Server workers configuration - parser.add_argument( - "--workers", - type=int, - default=get_env_value("WORKERS", 1, int), - help="Number of worker processes (default: from env or 1)", - ) - - # LLM and embedding bindings - parser.add_argument( - "--llm-binding", - type=str, - default=get_env_value("LLM_BINDING", "ollama"), - choices=["lollms", "ollama", "openai", "openai-ollama", "azure_openai"], - help="LLM binding type (default: from env or ollama)", - ) - parser.add_argument( - "--embedding-binding", - type=str, - default=get_env_value("EMBEDDING_BINDING", "ollama"), - choices=["lollms", "ollama", "openai", "azure_openai"], - help="Embedding binding type (default: from env or ollama)", - ) - - args = parser.parse_args() - - # If in uvicorn mode and workers > 1, force it to 1 and log warning - if is_uvicorn_mode and args.workers > 1: - original_workers = args.workers - args.workers = 1 - # Log warning directly here - logging.warning( - f"In uvicorn mode, workers parameter was set to {original_workers}. Forcing workers=1" - ) - - # convert relative path to absolute path - args.working_dir = os.path.abspath(args.working_dir) - args.input_dir = os.path.abspath(args.input_dir) - - # Inject storage configuration from environment variables - args.kv_storage = get_env_value( - "LIGHTRAG_KV_STORAGE", DefaultRAGStorageConfig.KV_STORAGE - ) - args.doc_status_storage = get_env_value( - "LIGHTRAG_DOC_STATUS_STORAGE", DefaultRAGStorageConfig.DOC_STATUS_STORAGE - ) - args.graph_storage = get_env_value( - "LIGHTRAG_GRAPH_STORAGE", DefaultRAGStorageConfig.GRAPH_STORAGE - ) - args.vector_storage = get_env_value( - "LIGHTRAG_VECTOR_STORAGE", DefaultRAGStorageConfig.VECTOR_STORAGE - ) - - # Get MAX_PARALLEL_INSERT from environment - args.max_parallel_insert = get_env_value("MAX_PARALLEL_INSERT", 2, int) - - # Handle openai-ollama special case - if args.llm_binding == "openai-ollama": - args.llm_binding = "openai" - args.embedding_binding = "ollama" - - args.llm_binding_host = get_env_value( - "LLM_BINDING_HOST", get_default_host(args.llm_binding) - ) - args.embedding_binding_host = get_env_value( - "EMBEDDING_BINDING_HOST", get_default_host(args.embedding_binding) - ) - args.llm_binding_api_key = get_env_value("LLM_BINDING_API_KEY", None) - args.embedding_binding_api_key = get_env_value("EMBEDDING_BINDING_API_KEY", "") - - # Inject model configuration - args.llm_model = get_env_value("LLM_MODEL", "mistral-nemo:latest") - args.embedding_model = get_env_value("EMBEDDING_MODEL", "bge-m3:latest") - args.embedding_dim = get_env_value("EMBEDDING_DIM", 1024, int) - args.max_embed_tokens = get_env_value("MAX_EMBED_TOKENS", 8192, int) - - # Inject chunk configuration - args.chunk_size = get_env_value("CHUNK_SIZE", 1200, int) - args.chunk_overlap_size = get_env_value("CHUNK_OVERLAP_SIZE", 100, int) - - # Inject LLM cache configuration - args.enable_llm_cache_for_extract = get_env_value( - "ENABLE_LLM_CACHE_FOR_EXTRACT", True, bool - ) - - # Inject LLM temperature configuration - args.temperature = get_env_value("TEMPERATURE", 0.5, float) - - # Select Document loading tool (DOCLING, DEFAULT) - args.document_loading_engine = get_env_value("DOCUMENT_LOADING_ENGINE", "DEFAULT") - - ollama_server_infos.LIGHTRAG_MODEL = args.simulated_model_name - - global_args["main_args"] = args - return args - - def display_splash_screen(args: argparse.Namespace) -> None: """ Display a colorful splash screen showing LightRAG server configuration @@ -518,8 +212,10 @@ def display_splash_screen(args: argparse.Namespace) -> None: ASCIIColors.yellow(f"{args.verbose}") ASCIIColors.white(" ├─ History Turns: ", end="") ASCIIColors.yellow(f"{args.history_turns}") - ASCIIColors.white(" └─ API Key: ", end="") + ASCIIColors.white(" ├─ API Key: ", end="") ASCIIColors.yellow("Set" if args.key else "Not Set") + ASCIIColors.white(" └─ JWT Auth: ", end="") + ASCIIColors.yellow("Enabled" if args.auth_accounts else "Disabled") # Directory Configuration ASCIIColors.magenta("\n📂 Directory Configuration:") @@ -652,6 +348,11 @@ def display_splash_screen(args: argparse.Namespace) -> None: ASCIIColors.white(""" API Key authentication is enabled. Make sure to include the X-API-Key header in all your requests. """) + if args.auth_accounts: + ASCIIColors.yellow("\n⚠️ Security Notice:") + ASCIIColors.white(""" JWT authentication is enabled. + Make sure to login before making the request, and include the 'Authorization' in the header. + """) # Ensure splash output flush to system log sys.stdout.flush() From 787d91cb26fe5cafbf7548956fe0594a82b823ea Mon Sep 17 00:00:00 2001 From: Milin Date: Mon, 31 Mar 2025 11:34:49 +0800 Subject: [PATCH 2/2] refactor(env): optimize env.example --- env.example | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/env.example b/env.example index b81d82b9..adbd82e2 100644 --- a/env.example +++ b/env.example @@ -175,11 +175,11 @@ QDRANT_URL=http://localhost:16333 REDIS_URI=redis://localhost:6379 ### For JWT Auth -#AUTH_ACCOUNTS='admin:admin123,user1:pass456' -#TOKEN_SECRET=Your-Key-For-LightRAG-API-Server -#TOKEN_EXPIRE_HOURS=4 -#GUEST_TOKEN_EXPIRE_HOURS=2 -#JWT_ALGORITHM=HS256 +# AUTH_ACCOUNTS='admin:admin123,user1:pass456' +# TOKEN_SECRET=Your-Key-For-LightRAG-API-Server +# TOKEN_EXPIRE_HOURS=4 +# GUEST_TOKEN_EXPIRE_HOURS=2 +# JWT_ALGORITHM=HS256 ### API-Key to access LightRAG Server API # LIGHTRAG_API_KEY=your-secure-api-key-here