gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor LLM cache config to use argparse and add status display

Browse Source
This commit is contained in:
yangdx
2025-03-09 16:45:57 +08:00
parent c854aabde0
commit 020a6b5ae0
2 changed files with 13 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
lightrag/api/lightrag_server.py
View File

@@ -50,9 +50,6 @@ from .auth import auth_handler
# This update allows the user to put a different.env file for each lightrag folder # This update allows the user to put a different.env file for each lightrag folder
load_dotenv(".env", override=True) load_dotenv(".env", override=True)
# Read entity extraction cache config
enable_llm_cache = os.getenv("ENABLE_LLM_CACHE_FOR_EXTRACT", "false").lower() == "true"
# Initialize config parser # Initialize config parser
config = configparser.ConfigParser() config = configparser.ConfigParser()
config.read("config.ini") config.read("config.ini")
@@ -326,7 +323,7 @@ def create_app(args):
vector_db_storage_cls_kwargs={ vector_db_storage_cls_kwargs={
"cosine_better_than_threshold": args.cosine_threshold "cosine_better_than_threshold": args.cosine_threshold
}, },
enable_llm_cache_for_entity_extract=enable_llm_cache, # Read from environment variable enable_llm_cache_for_entity_extract=args.enable_llm_cache, # Read from args
embedding_cache_config={ embedding_cache_config={
"enabled": True, "enabled": True,
"similarity_threshold": 0.95, "similarity_threshold": 0.95,
@@ -355,7 +352,7 @@ def create_app(args):
vector_db_storage_cls_kwargs={ vector_db_storage_cls_kwargs={
"cosine_better_than_threshold": args.cosine_threshold "cosine_better_than_threshold": args.cosine_threshold
}, },
enable_llm_cache_for_entity_extract=enable_llm_cache, # Read from environment variable enable_llm_cache_for_entity_extract=args.enable_llm_cache, # Read from args
embedding_cache_config={ embedding_cache_config={
"enabled": True, "enabled": True,
"similarity_threshold": 0.95, "similarity_threshold": 0.95,
@@ -419,6 +416,7 @@ def create_app(args):
"doc_status_storage": args.doc_status_storage, "doc_status_storage": args.doc_status_storage,
"graph_storage": args.graph_storage, "graph_storage": args.graph_storage,
"vector_storage": args.vector_storage, "vector_storage": args.vector_storage,
"enable_llm_cache": args.enable_llm_cache,
}, },
"update_status": update_status, "update_status": update_status,
} }

11
lightrag/api/utils_api.py
View File

@@ -360,6 +360,13 @@ def parse_args(is_uvicorn_mode: bool = False) -> argparse.Namespace:
args.chunk_size = get_env_value("CHUNK_SIZE", 1200, int) args.chunk_size = get_env_value("CHUNK_SIZE", 1200, int)
args.chunk_overlap_size = get_env_value("CHUNK_OVERLAP_SIZE", 100, int) args.chunk_overlap_size = get_env_value("CHUNK_OVERLAP_SIZE", 100, int)
# Inject LLM cache configuration
args.enable_llm_cache = get_env_value(
"ENABLE_LLM_CACHE_FOR_EXTRACT",
False,
bool
)
ollama_server_infos.LIGHTRAG_MODEL = args.simulated_model_name ollama_server_infos.LIGHTRAG_MODEL = args.simulated_model_name
return args return args
@@ -451,8 +458,10 @@ def display_splash_screen(args: argparse.Namespace) -> None:
ASCIIColors.yellow(f"{args.history_turns}") ASCIIColors.yellow(f"{args.history_turns}")
ASCIIColors.white(" ├─ Cosine Threshold: ", end="") ASCIIColors.white(" ├─ Cosine Threshold: ", end="")
ASCIIColors.yellow(f"{args.cosine_threshold}") ASCIIColors.yellow(f"{args.cosine_threshold}")
ASCIIColors.white(" └─ Top-K: ", end="") ASCIIColors.white(" ├─ Top-K: ", end="")
ASCIIColors.yellow(f"{args.top_k}") ASCIIColors.yellow(f"{args.top_k}")
ASCIIColors.white(" └─ LLM Cache Enabled: ", end="")
ASCIIColors.yellow(f"{args.enable_llm_cache}")
# System Configuration # System Configuration
ASCIIColors.magenta("\n💾 Storage Configuration:") ASCIIColors.magenta("\n💾 Storage Configuration:")