From 5f7cb437e852ea73387fde9a567d1d42c92bbb3b Mon Sep 17 00:00:00 2001 From: yangdx Date: Tue, 15 Jul 2025 23:56:49 +0800 Subject: [PATCH] Centralize query parameters into LightRAG class This commit refactors query parameter management by consolidating settings like `top_k`, token limits, and thresholds into the `LightRAG` class, and consistently sourcing parameters from a single location. --- env.example | 2 ++ lightrag/api/config.py | 52 ++++++++++++++++++++---------------------- lightrag/constants.py | 5 ++-- lightrag/lightrag.py | 43 ++++++++++++++++++++++++++++++++++ lightrag/operate.py | 26 +++++---------------- 5 files changed, 78 insertions(+), 50 deletions(-) diff --git a/env.example b/env.example index 5bbcfac3..38ac249e 100644 --- a/env.example +++ b/env.example @@ -62,6 +62,8 @@ ENABLE_LLM_CACHE=true # MAX_RELATION_TOKENS=10000 ### control the maximum tokens send to LLM (include entities, raltions and chunks) # MAX_TOTAL_TOKENS=32000 +### maxumium related chunks grab from single entity or relations +# RELATED_CHUNK_NUMBER=5 ### Reranker configuration (Set ENABLE_RERANK to true in reranking model is configed) ENABLE_RERANK=False diff --git a/lightrag/api/config.py b/lightrag/api/config.py index cd498c1d..98b817c1 100644 --- a/lightrag/api/config.py +++ b/lightrag/api/config.py @@ -14,6 +14,11 @@ from lightrag.constants import ( 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, ) # use the .env that is inside the current folder @@ -154,33 +159,6 @@ def parse_args() -> argparse.Namespace: 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", DEFAULT_HISTORY_TURNS, 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", DEFAULT_TOP_K, int), - help="Number of most similar results to return (default: from env or 60)", - ) - parser.add_argument( - "--chunk-top-k", - type=int, - default=get_env_value("CHUNK_TOP_K", DEFAULT_CHUNK_TOP_K, int), - help="Number of text chunks to retrieve initially from vector search and keep after reranking (default: from env or 5)", - ) - 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", @@ -312,6 +290,26 @@ def parse_args() -> argparse.Namespace: args.rerank_binding_host = get_env_value("RERANK_BINDING_HOST", None) args.rerank_binding_api_key = get_env_value("RERANK_BINDING_API_KEY", None) + # 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 + ) + ollama_server_infos.LIGHTRAG_MODEL = args.simulated_model_name return args diff --git a/lightrag/constants.py b/lightrag/constants.py index e975f2bf..fc25a855 100644 --- a/lightrag/constants.py +++ b/lightrag/constants.py @@ -20,6 +20,8 @@ DEFAULT_MAX_RELATION_TOKENS = 10000 DEFAULT_MAX_TOTAL_TOKENS = 32000 DEFAULT_HISTORY_TURNS = 3 DEFAULT_ENABLE_RERANK = True +DEFAULT_COSINE_THRESHOLD = 0.2 +DEFAULT_RELATED_CHUNK_NUMBER = 5 # Separator for graph fields GRAPH_FIELD_SEP = "" @@ -28,6 +30,3 @@ GRAPH_FIELD_SEP = "" DEFAULT_LOG_MAX_BYTES = 10485760 # Default 10MB DEFAULT_LOG_BACKUP_COUNT = 5 # Default 5 backups DEFAULT_LOG_FILENAME = "lightrag.log" # Default log filename - -# Related Chunk Number for Single Entity or Relation -DEFAULT_RELATED_CHUNK_NUMBER = 5 diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index df3929f6..092e06eb 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -24,6 +24,13 @@ from typing import ( from lightrag.constants import ( DEFAULT_MAX_GLEANING, DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE, + DEFAULT_TOP_K, + DEFAULT_CHUNK_TOP_K, + DEFAULT_MAX_ENTITY_TOKENS, + DEFAULT_MAX_RELATION_TOKENS, + DEFAULT_MAX_TOTAL_TOKENS, + DEFAULT_COSINE_THRESHOLD, + DEFAULT_RELATED_CHUNK_NUMBER, ) from lightrag.utils import get_env_value @@ -125,6 +132,42 @@ class LightRAG: log_level: int | None = field(default=None) log_file_path: str | None = field(default=None) + # Query parameters + # --- + + top_k: int = field(default=get_env_value("TOP_K", DEFAULT_TOP_K, int)) + """Number of entities/relations to retrieve for each query.""" + + chunk_top_k: int = field( + default=get_env_value("CHUNK_TOP_K", DEFAULT_CHUNK_TOP_K, int) + ) + """Maximum number of chunks in context.""" + + max_entity_tokens: int = field( + default=get_env_value("MAX_ENTITY_TOKENS", DEFAULT_MAX_ENTITY_TOKENS, int) + ) + """Maximum number of tokens for entity in context.""" + + max_relation_tokens: int = field( + default=get_env_value("MAX_RELATION_TOKENS", DEFAULT_MAX_RELATION_TOKENS, int) + ) + """Maximum number of tokens for relation in context.""" + + max_total_tokens: int = field( + default=get_env_value("MAX_TOTAL_TOKENS", DEFAULT_MAX_TOTAL_TOKENS, int) + ) + """Maximum total tokens in context (including system prompt, entities, relations and chunks).""" + + cosine_threshold: int = field( + default=get_env_value("COSINE_THRESHOLD", DEFAULT_COSINE_THRESHOLD, int) + ) + """Cosine threshold of vector DB retrieval for entities, relations and chunks.""" + + related_chunk_number: int = field( + default=get_env_value("RELATED_CHUNK_NUMBER", DEFAULT_RELATED_CHUNK_NUMBER, int) + ) + """Number of related chunks to grab from single entity or relation.""" + # Entity extraction # --- diff --git a/lightrag/operate.py b/lightrag/operate.py index 073db4d4..f45b8075 100644 --- a/lightrag/operate.py +++ b/lightrag/operate.py @@ -1908,7 +1908,6 @@ async def _build_query_context( ll_keywords, knowledge_graph_inst, entities_vdb, - text_chunks_db, query_param, ) original_node_datas = node_datas @@ -1924,7 +1923,6 @@ async def _build_query_context( hl_keywords, knowledge_graph_inst, relationships_vdb, - text_chunks_db, query_param, ) original_edge_datas = edge_datas @@ -1935,14 +1933,12 @@ async def _build_query_context( ll_keywords, knowledge_graph_inst, entities_vdb, - text_chunks_db, query_param, ) hl_data = await _get_edge_data( hl_keywords, knowledge_graph_inst, relationships_vdb, - text_chunks_db, query_param, ) @@ -1985,23 +1981,17 @@ async def _build_query_context( max_entity_tokens = getattr( query_param, "max_entity_tokens", - text_chunks_db.global_config.get( - "MAX_ENTITY_TOKENS", DEFAULT_MAX_ENTITY_TOKENS - ), + text_chunks_db.global_config.get("max_entity_tokens", DEFAULT_MAX_ENTITY_TOKENS), ) max_relation_tokens = getattr( query_param, "max_relation_tokens", - text_chunks_db.global_config.get( - "MAX_RELATION_TOKENS", DEFAULT_MAX_RELATION_TOKENS - ), + text_chunks_db.global_config.get("max_relation_tokens", DEFAULT_MAX_RELATION_TOKENS), ) max_total_tokens = getattr( query_param, "max_total_tokens", - text_chunks_db.global_config.get( - "MAX_TOTAL_TOKENS", DEFAULT_MAX_TOTAL_TOKENS - ), + text_chunks_db.global_config.get("max_total_tokens", DEFAULT_MAX_TOTAL_TOKENS), ) # Truncate entities based on complete JSON serialization @@ -2095,7 +2085,6 @@ async def _build_query_context( final_edge_datas, query_param, text_chunks_db, - knowledge_graph_inst, ) ) @@ -2255,7 +2244,6 @@ async def _get_node_data( query: str, knowledge_graph_inst: BaseGraphStorage, entities_vdb: BaseVectorStorage, - text_chunks_db: BaseKVStorage, query_param: QueryParam, ): # get similar entities @@ -2362,7 +2350,7 @@ async def _find_most_related_text_unit_from_entities( text_units = [ split_string_by_multi_markers(dp["source_id"], [GRAPH_FIELD_SEP])[ - :DEFAULT_RELATED_CHUNK_NUMBER + : text_chunks_db.global_config.get("related_chunk_number", DEFAULT_RELATED_CHUNK_NUMBER) ] for dp in node_datas if dp["source_id"] is not None @@ -2519,7 +2507,6 @@ async def _get_edge_data( keywords, knowledge_graph_inst: BaseGraphStorage, relationships_vdb: BaseVectorStorage, - text_chunks_db: BaseKVStorage, query_param: QueryParam, ): logger.info( @@ -2668,13 +2655,12 @@ async def _find_related_text_unit_from_relationships( edge_datas: list[dict], query_param: QueryParam, text_chunks_db: BaseKVStorage, - knowledge_graph_inst: BaseGraphStorage, ): logger.debug(f"Searching text chunks for {len(edge_datas)} relationships") text_units = [ split_string_by_multi_markers(dp["source_id"], [GRAPH_FIELD_SEP])[ - :DEFAULT_RELATED_CHUNK_NUMBER + : text_chunks_db.global_config.get("related_chunk_number", DEFAULT_RELATED_CHUNK_NUMBER) ] for dp in edge_datas if dp["source_id"] is not None @@ -2761,7 +2747,7 @@ async def naive_query( max_total_tokens = getattr( query_param, "max_total_tokens", - global_config.get("MAX_TOTAL_TOKENS", DEFAULT_MAX_TOTAL_TOKENS), + global_config.get("max_total_tokens", DEFAULT_MAX_TOTAL_TOKENS), ) # Calculate conversation history tokens