9923821d75
This parameter is no longer used. Its removal simplifies the API and clarifies that token length management is handled by upstream text chunking logic rather than the embedding wrapper.
411 lines
14 KiB
Python
411 lines
14 KiB
Python
"""
|
|
Configs for the LightRAG API.
|
|
"""
|
|
|
|
import os
|
|
import argparse
|
|
import logging
|
|
from dotenv import load_dotenv
|
|
from lightrag.utils import get_env_value
|
|
from lightrag.llm.binding_options import OllamaEmbeddingOptions, OllamaLLMOptions
|
|
import sys
|
|
|
|
from lightrag.constants import (
|
|
DEFAULT_WOKERS,
|
|
DEFAULT_TIMEOUT,
|
|
DEFAULT_TOP_K,
|
|
DEFAULT_CHUNK_TOP_K,
|
|
DEFAULT_HISTORY_TURNS,
|
|
DEFAULT_MAX_ENTITY_TOKENS,
|
|
DEFAULT_MAX_RELATION_TOKENS,
|
|
DEFAULT_MAX_TOTAL_TOKENS,
|
|
DEFAULT_COSINE_THRESHOLD,
|
|
DEFAULT_RELATED_CHUNK_NUMBER,
|
|
DEFAULT_MIN_RERANK_SCORE,
|
|
DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE,
|
|
DEFAULT_MAX_ASYNC,
|
|
DEFAULT_SUMMARY_MAX_TOKENS,
|
|
DEFAULT_SUMMARY_LANGUAGE,
|
|
DEFAULT_EMBEDDING_FUNC_MAX_ASYNC,
|
|
DEFAULT_EMBEDDING_BATCH_NUM,
|
|
DEFAULT_OLLAMA_MODEL_NAME,
|
|
DEFAULT_OLLAMA_MODEL_TAG,
|
|
DEFAULT_OLLAMA_MODEL_SIZE,
|
|
DEFAULT_OLLAMA_CREATED_AT,
|
|
DEFAULT_OLLAMA_DIGEST,
|
|
)
|
|
|
|
# use the .env that is inside the current folder
|
|
# allows to use different .env file for each lightrag instance
|
|
# the OS environment variables take precedence over the .env file
|
|
load_dotenv(dotenv_path=".env", override=False)
|
|
|
|
|
|
class OllamaServerInfos:
|
|
def __init__(self, name=None, tag=None):
|
|
self._lightrag_name = name or os.getenv(
|
|
"OLLAMA_EMULATING_MODEL_NAME", DEFAULT_OLLAMA_MODEL_NAME
|
|
)
|
|
self._lightrag_tag = tag or os.getenv(
|
|
"OLLAMA_EMULATING_MODEL_TAG", DEFAULT_OLLAMA_MODEL_TAG
|
|
)
|
|
self.LIGHTRAG_SIZE = DEFAULT_OLLAMA_MODEL_SIZE
|
|
self.LIGHTRAG_CREATED_AT = DEFAULT_OLLAMA_CREATED_AT
|
|
self.LIGHTRAG_DIGEST = DEFAULT_OLLAMA_DIGEST
|
|
|
|
@property
|
|
def LIGHTRAG_NAME(self):
|
|
return self._lightrag_name
|
|
|
|
@LIGHTRAG_NAME.setter
|
|
def LIGHTRAG_NAME(self, value):
|
|
self._lightrag_name = value
|
|
|
|
@property
|
|
def LIGHTRAG_TAG(self):
|
|
return self._lightrag_tag
|
|
|
|
@LIGHTRAG_TAG.setter
|
|
def LIGHTRAG_TAG(self, value):
|
|
self._lightrag_tag = value
|
|
|
|
@property
|
|
def LIGHTRAG_MODEL(self):
|
|
return f"{self._lightrag_name}:{self._lightrag_tag}"
|
|
|
|
|
|
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 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)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--timeout",
|
|
default=get_env_value("TIMEOUT", DEFAULT_TIMEOUT, int, special_none=True),
|
|
type=int,
|
|
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", DEFAULT_MAX_ASYNC, int),
|
|
help=f"Maximum async operations (default: from env or {DEFAULT_MAX_ASYNC})",
|
|
)
|
|
parser.add_argument(
|
|
"--max-tokens",
|
|
type=int,
|
|
default=get_env_value("MAX_TOKENS", DEFAULT_SUMMARY_MAX_TOKENS, int),
|
|
help=f"Maximum token size (default: from env or {DEFAULT_SUMMARY_MAX_TOKENS})",
|
|
)
|
|
|
|
# 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)",
|
|
)
|
|
|
|
# Ollama model configuration
|
|
parser.add_argument(
|
|
"--simulated-model-name",
|
|
type=str,
|
|
default=get_env_value("OLLAMA_EMULATING_MODEL_NAME", DEFAULT_OLLAMA_MODEL_NAME),
|
|
help="Name for the simulated Ollama model (default: from env or lightrag)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--simulated-model-tag",
|
|
type=str,
|
|
default=get_env_value("OLLAMA_EMULATING_MODEL_TAG", DEFAULT_OLLAMA_MODEL_TAG),
|
|
help="Tag for the simulated Ollama model (default: from env or latest)",
|
|
)
|
|
|
|
# Namespace
|
|
parser.add_argument(
|
|
"--workspace",
|
|
type=str,
|
|
default=get_env_value("WORKSPACE", ""),
|
|
help="Default workspace for all storage",
|
|
)
|
|
|
|
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", DEFAULT_WOKERS, 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)",
|
|
)
|
|
|
|
# Conditionally add binding options defined in binding_options module
|
|
# This will add command line arguments for all binding options (e.g., --ollama-embedding-num_ctx)
|
|
# and corresponding environment variables (e.g., OLLAMA_EMBEDDING_NUM_CTX)
|
|
if "--llm-binding" in sys.argv:
|
|
try:
|
|
idx = sys.argv.index("--llm-binding")
|
|
if idx + 1 < len(sys.argv) and sys.argv[idx + 1] == "ollama":
|
|
OllamaLLMOptions.add_args(parser)
|
|
except IndexError:
|
|
pass
|
|
elif os.environ.get("LLM_BINDING") == "ollama":
|
|
OllamaLLMOptions.add_args(parser)
|
|
|
|
if "--embedding-binding" in sys.argv:
|
|
try:
|
|
idx = sys.argv.index("--embedding-binding")
|
|
if idx + 1 < len(sys.argv) and sys.argv[idx + 1] == "ollama":
|
|
OllamaEmbeddingOptions.add_args(parser)
|
|
except IndexError:
|
|
pass
|
|
elif os.environ.get("EMBEDDING_BINDING") == "ollama":
|
|
OllamaEmbeddingOptions.add_args(parser)
|
|
|
|
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)
|
|
|
|
# Get MAX_GRAPH_NODES from environment
|
|
args.max_graph_nodes = get_env_value("MAX_GRAPH_NODES", 1000, int)
|
|
|
|
# Handle openai-ollama special case
|
|
if args.llm_binding == "openai-ollama":
|
|
args.llm_binding = "openai"
|
|
args.embedding_binding = "ollama"
|
|
|
|
# Ollama ctx_num
|
|
args.ollama_num_ctx = get_env_value("OLLAMA_NUM_CTX", 32768, int)
|
|
|
|
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)
|
|
|
|
# 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
|
|
)
|
|
args.enable_llm_cache = get_env_value("ENABLE_LLM_CACHE", 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")
|
|
|
|
# Add environment variables that were previously read directly
|
|
args.cors_origins = get_env_value("CORS_ORIGINS", "*")
|
|
args.summary_language = get_env_value("SUMMARY_LANGUAGE", DEFAULT_SUMMARY_LANGUAGE)
|
|
args.whitelist_paths = get_env_value("WHITELIST_PATHS", "/health,/api/*")
|
|
|
|
# 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", 48, int)
|
|
args.guest_token_expire_hours = get_env_value("GUEST_TOKEN_EXPIRE_HOURS", 24, int)
|
|
args.jwt_algorithm = get_env_value("JWT_ALGORITHM", "HS256")
|
|
|
|
# Rerank model configuration
|
|
args.rerank_model = get_env_value("RERANK_MODEL", "BAAI/bge-reranker-v2-m3")
|
|
args.rerank_binding_host = get_env_value("RERANK_BINDING_HOST", None)
|
|
args.rerank_binding_api_key = get_env_value("RERANK_BINDING_API_KEY", None)
|
|
|
|
# Min rerank score configuration
|
|
args.min_rerank_score = get_env_value(
|
|
"MIN_RERANK_SCORE", DEFAULT_MIN_RERANK_SCORE, float
|
|
)
|
|
|
|
# Query configuration
|
|
args.history_turns = get_env_value("HISTORY_TURNS", DEFAULT_HISTORY_TURNS, int)
|
|
args.top_k = get_env_value("TOP_K", DEFAULT_TOP_K, int)
|
|
args.chunk_top_k = get_env_value("CHUNK_TOP_K", DEFAULT_CHUNK_TOP_K, int)
|
|
args.max_entity_tokens = get_env_value(
|
|
"MAX_ENTITY_TOKENS", DEFAULT_MAX_ENTITY_TOKENS, int
|
|
)
|
|
args.max_relation_tokens = get_env_value(
|
|
"MAX_RELATION_TOKENS", DEFAULT_MAX_RELATION_TOKENS, int
|
|
)
|
|
args.max_total_tokens = get_env_value(
|
|
"MAX_TOTAL_TOKENS", DEFAULT_MAX_TOTAL_TOKENS, int
|
|
)
|
|
args.cosine_threshold = get_env_value(
|
|
"COSINE_THRESHOLD", DEFAULT_COSINE_THRESHOLD, float
|
|
)
|
|
args.related_chunk_number = get_env_value(
|
|
"RELATED_CHUNK_NUMBER", DEFAULT_RELATED_CHUNK_NUMBER, int
|
|
)
|
|
|
|
# Add missing environment variables for health endpoint
|
|
args.force_llm_summary_on_merge = get_env_value(
|
|
"FORCE_LLM_SUMMARY_ON_MERGE", DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE, int
|
|
)
|
|
args.embedding_func_max_async = get_env_value(
|
|
"EMBEDDING_FUNC_MAX_ASYNC", DEFAULT_EMBEDDING_FUNC_MAX_ASYNC, int
|
|
)
|
|
args.embedding_batch_num = get_env_value(
|
|
"EMBEDDING_BATCH_NUM", DEFAULT_EMBEDDING_BATCH_NUM, int
|
|
)
|
|
|
|
ollama_server_infos.LIGHTRAG_NAME = args.simulated_model_name
|
|
ollama_server_infos.LIGHTRAG_TAG = args.simulated_model_tag
|
|
|
|
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()
|