Merge pull request #1782 from HKUDS/rerank

README-zh.md

@@ -293,26 +293,19 @@ class QueryParam:
     top_k: int = int(os.getenv("TOP_K", "60"))
     """Number of top items to retrieve. Represents entities in 'local' mode and relationships in 'global' mode."""
 
-    chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", "5"))
-    """Number of text chunks to retrieve initially from vector search.
+    chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", "10"))
+    """Number of text chunks to retrieve initially from vector search and keep after reranking.
     If None, defaults to top_k value.
     """
 
-    chunk_rerank_top_k: int = int(os.getenv("CHUNK_RERANK_TOP_K", "5"))
-    """Number of text chunks to keep after reranking.
-    If None, keeps all chunks returned from initial retrieval.
-    """
-
-    max_token_for_text_unit: int = int(os.getenv("MAX_TOKEN_TEXT_CHUNK", "4000"))
-    """Maximum number of tokens allowed for each retrieved text chunk."""
+    max_entity_tokens: int = int(os.getenv("MAX_ENTITY_TOKENS", "10000"))
+    """Maximum number of tokens allocated for entity context in unified token control system."""
 
-    max_token_for_global_context: int = int(
-        os.getenv("MAX_TOKEN_RELATION_DESC", "4000")
-    )
-    """Maximum number of tokens allocated for relationship descriptions in global retrieval."""
+    max_relation_tokens: int = int(os.getenv("MAX_RELATION_TOKENS", "10000"))
+    """Maximum number of tokens allocated for relationship context in unified token control system."""
 
-    max_token_for_local_context: int = int(os.getenv("MAX_TOKEN_ENTITY_DESC", "4000"))
-    """Maximum number of tokens allocated for entity descriptions in local retrieval."""
+    max_total_tokens: int = int(os.getenv("MAX_TOTAL_TOKENS", "32000"))
+    """Maximum total tokens budget for the entire query context (entities + relations + chunks + system prompt)."""
 
     hl_keywords: list[str] = field(default_factory=list)
     """List of high-level keywords to prioritize in retrieval."""
@@ -341,6 +334,11 @@ class QueryParam:
     """User-provided prompt for the query.
     If proivded, this will be use instead of the default vaulue from prompt template.
     """
+
+    enable_rerank: bool = True
+    """Enable reranking for retrieved text chunks. If True but no rerank model is configured, a warning will be issued.
+    Default is True to enable reranking when rerank model is available.
+    """
 ```
 
 > top_k的默认值可以通过环境变量TOP_K更改。

README.md

@@ -300,26 +300,19 @@ class QueryParam:
     top_k: int = int(os.getenv("TOP_K", "60"))
     """Number of top items to retrieve. Represents entities in 'local' mode and relationships in 'global' mode."""
 
-    chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", "5"))
-    """Number of text chunks to retrieve initially from vector search.
+    chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", "10"))
+    """Number of text chunks to retrieve initially from vector search and keep after reranking.
     If None, defaults to top_k value.
     """
 
-    chunk_rerank_top_k: int = int(os.getenv("CHUNK_RERANK_TOP_K", "5"))
-    """Number of text chunks to keep after reranking.
-    If None, keeps all chunks returned from initial retrieval.
-    """
-
-    max_token_for_text_unit: int = int(os.getenv("MAX_TOKEN_TEXT_CHUNK", "4000"))
-    """Maximum number of tokens allowed for each retrieved text chunk."""
+    max_entity_tokens: int = int(os.getenv("MAX_ENTITY_TOKENS", "10000"))
+    """Maximum number of tokens allocated for entity context in unified token control system."""
 
-    max_token_for_global_context: int = int(
-        os.getenv("MAX_TOKEN_RELATION_DESC", "4000")
-    )
-    """Maximum number of tokens allocated for relationship descriptions in global retrieval."""
+    max_relation_tokens: int = int(os.getenv("MAX_RELATION_TOKENS", "10000"))
+    """Maximum number of tokens allocated for relationship context in unified token control system."""
 
-    max_token_for_local_context: int = int(os.getenv("MAX_TOKEN_ENTITY_DESC", "4000"))
-    """Maximum number of tokens allocated for entity descriptions in local retrieval."""
+    max_total_tokens: int = int(os.getenv("MAX_TOTAL_TOKENS", "32000"))
+    """Maximum total tokens budget for the entire query context (entities + relations + chunks + system prompt)."""
 
     conversation_history: list[dict[str, str]] = field(default_factory=list)
     """Stores past conversation history to maintain context.
@@ -342,6 +335,11 @@ class QueryParam:
     """User-provided prompt for the query.
     If proivded, this will be use instead of the default vaulue from prompt template.
     """
+
+    enable_rerank: bool = True
+    """Enable reranking for retrieved text chunks. If True but no rerank model is configured, a warning will be issued.
+    Default is True to enable reranking when rerank model is available.
+    """
 ```
 
 > default value of Top_k can be change by environment  variables  TOP_K.

docs/rerank_integration.md

@@ -1,36 +1,24 @@
-# Rerank Integration in LightRAG
+# Rerank Integration Guide
 
-This document explains how to configure and use the rerank functionality in LightRAG to improve retrieval quality.
+LightRAG supports reranking functionality to improve retrieval quality by re-ordering documents based on their relevance to the query. Reranking is now controlled per query via the `enable_rerank` parameter (default: True).
 
-## Overview
-
-Reranking is an optional feature that improves the quality of retrieved documents by re-ordering them based on their relevance to the query. This is particularly useful when you want higher precision in document retrieval across all query modes (naive, local, global, hybrid, mix).
-
-## Architecture
-
-The rerank integration follows a simplified design pattern:
-
-- **Single Function Configuration**: All rerank settings (model, API keys, top_k, etc.) are contained within the rerank function
-- **Async Processing**: Non-blocking rerank operations
-- **Error Handling**: Graceful fallback to original results
-- **Optional Feature**: Can be enabled/disabled via configuration
-- **Code Reuse**: Single generic implementation for Jina/Cohere compatible APIs
-
-## Configuration
+## Quick Start
 
 ### Environment Variables
 
-Set this variable in your `.env` file or environment:
+Set these variables in your `.env` file or environment for rerank model configuration:
 
 ```bash
-# Enable/disable reranking
-ENABLE_RERANK=True
+# Rerank model configuration (required when enable_rerank=True in queries)
+RERANK_MODEL=BAAI/bge-reranker-v2-m3
+RERANK_BINDING_HOST=https://api.your-provider.com/v1/rerank
+RERANK_BINDING_API_KEY=your_api_key_here
 ```
 
 ### Programmatic Configuration
 
 ```python
-from lightrag import LightRAG
+from lightrag import LightRAG, QueryParam
 from lightrag.rerank import custom_rerank, RerankModel
 
 # Method 1: Using a custom rerank function with all settings included
@@ -49,8 +37,19 @@ rag = LightRAG(
     working_dir="./rag_storage",
     llm_model_func=your_llm_func,
     embedding_func=your_embedding_func,
-    enable_rerank=True,
-    rerank_model_func=my_rerank_func,
+    rerank_model_func=my_rerank_func,  # Configure rerank function
+)
+
+# Query with rerank enabled (default)
+result = await rag.aquery(
+    "your query",
+    param=QueryParam(enable_rerank=True)  # Control rerank per query
+)
+
+# Query with rerank disabled
+result = await rag.aquery(
+    "your query",
+    param=QueryParam(enable_rerank=False)
 )
 
 # Method 2: Using RerankModel wrapper
@@ -67,9 +66,17 @@ rag = LightRAG(
     working_dir="./rag_storage",
     llm_model_func=your_llm_func,
     embedding_func=your_embedding_func,
-    enable_rerank=True,
     rerank_model_func=rerank_model.rerank,
 )
+
+# Control rerank per query
+result = await rag.aquery(
+    "your query",
+    param=QueryParam(
+        enable_rerank=True,  # Enable rerank for this query
+        chunk_top_k=5       # Number of chunks to keep after reranking
+    )
+)
 ```
 
 ## Supported Providers
@@ -164,7 +171,6 @@ async def main():
         working_dir="./rag_storage",
         llm_model_func=gpt_4o_mini_complete,
         embedding_func=openai_embedding,
-        enable_rerank=True,
         rerank_model_func=my_rerank_func,
     )
 
@@ -180,7 +186,7 @@ async def main():
     # Query with rerank (automatically applied)
     result = await rag.aquery(
         "Your question here",
-        param=QueryParam(mode="hybrid", top_k=5)  # This top_k is passed to rerank function
+        param=QueryParam(enable_rerank=True)  # This top_k is passed to rerank function
     )
 
     print(result)

env.example

@@ -1,6 +1,8 @@
 ### This is sample file of .env
 
+###########################
 ### Server Configuration
+###########################
 HOST=0.0.0.0
 PORT=9621
 WEBUI_TITLE='My Graph KB'
@@ -9,29 +11,17 @@ OLLAMA_EMULATING_MODEL_TAG=latest
 # WORKERS=2
 # CORS_ORIGINS=http://localhost:3000,http://localhost:8080
 
-### Login Configuration
-# AUTH_ACCOUNTS='admin:admin123,user1:pass456'
-# TOKEN_SECRET=Your-Key-For-LightRAG-API-Server
-# TOKEN_EXPIRE_HOURS=48
-# GUEST_TOKEN_EXPIRE_HOURS=24
-# JWT_ALGORITHM=HS256
-
-### API-Key to access LightRAG Server API
-# LIGHTRAG_API_KEY=your-secure-api-key-here
-# WHITELIST_PATHS=/health,/api/*
-
 ### Optional SSL Configuration
 # SSL=true
 # SSL_CERTFILE=/path/to/cert.pem
 # SSL_KEYFILE=/path/to/key.pem
 
 ### Directory Configuration (defaults to current working directory)
-### Should not be set if deploy by docker (Set by Dockerfile instead of .env)
 ### Default value is ./inputs and ./rag_storage
 # INPUT_DIR=<absolute_path_for_doc_input_dir>
 # WORKING_DIR=<absolute_path_for_working_dir>
 
-### Max nodes return from grap retrieval
+### Max nodes return from grap retrieval in webui
 # MAX_GRAPH_NODES=1000
 
 ### Logging level
@@ -42,65 +32,97 @@ OLLAMA_EMULATING_MODEL_TAG=latest
 ### Logfile location (defaults to current working directory)
 # LOG_DIR=/path/to/log/directory
 
-### RAG Configuration
-### Chunk size for document splitting, 500~1500 is recommended
-# CHUNK_SIZE=1200
-# CHUNK_OVERLAP_SIZE=100
+#####################################
+### Login and API-Key Configuration
+#####################################
+# AUTH_ACCOUNTS='admin:admin123,user1:pass456'
+# TOKEN_SECRET=Your-Key-For-LightRAG-API-Server
+# TOKEN_EXPIRE_HOURS=48
+# GUEST_TOKEN_EXPIRE_HOURS=24
+# JWT_ALGORITHM=HS256
 
-### RAG Query Configuration
+### API-Key to access LightRAG Server API
+# LIGHTRAG_API_KEY=your-secure-api-key-here
+# WHITELIST_PATHS=/health,/api/*
+
+########################
+### Query Configuration
+########################
+# LLM responde cache for query (Not valid for streaming response
+ENABLE_LLM_CACHE=true
 # HISTORY_TURNS=3
-# MAX_TOKEN_TEXT_CHUNK=6000
-# MAX_TOKEN_RELATION_DESC=4000
-# MAX_TOKEN_ENTITY_DESC=4000
 # COSINE_THRESHOLD=0.2
-### Number of entities or relations to retrieve from KG
-# TOP_K=60
-### Number of text chunks to retrieve initially from vector search
-# CHUNK_TOP_K=5
-
-### Rerank Configuration
-# ENABLE_RERANK=False
-### Number of text chunks to keep after reranking (should be <= CHUNK_TOP_K)
-# CHUNK_RERANK_TOP_K=5
-### Rerank model configuration (required when ENABLE_RERANK=True)
+### Number of entities or relations retrieved from KG
+# TOP_K=40
+### Maxmium number or chunks plan to send to LLM
+# CHUNK_TOP_K=10
+### control the actual enties send to LLM
+# MAX_ENTITY_TOKENS=10000
+### control the actual relations send to LLM
+# 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=10
+
+### Reranker configuration (Set ENABLE_RERANK to true in reranking model is configed)
+ENABLE_RERANK=False
 # RERANK_MODEL=BAAI/bge-reranker-v2-m3
 # RERANK_BINDING_HOST=https://api.your-rerank-provider.com/v1/rerank
 # RERANK_BINDING_API_KEY=your_rerank_api_key_here
 
-### Entity and relation summarization configuration
+########################################
+### Document processing configuration
+########################################
 ### Language: English, Chinese, French, German ...
 SUMMARY_LANGUAGE=English
+ENABLE_LLM_CACHE_FOR_EXTRACT=true
+### MAX_TOKENS: max tokens send to LLM for entity relation summaries (less than context size of the model)
+MAX_TOKENS=32000
+### Chunk size for document splitting, 500~1500 is recommended
+# CHUNK_SIZE=1200
+# CHUNK_OVERLAP_SIZE=100
+### Entity and relation summarization configuration
 ### Number of duplicated entities/edges to trigger LLM re-summary on merge ( at least 3 is recommented)
-# FORCE_LLM_SUMMARY_ON_MERGE=6
+# FORCE_LLM_SUMMARY_ON_MERGE=4
 ### Maximum number of entity extraction attempts for ambiguous content
 # MAX_GLEANING=1
 
-### Number of parallel processing documents(Less than MAX_ASYNC/2 is recommended)
-# MAX_PARALLEL_INSERT=2
+###############################
+### Concurrency Configuration
+###############################
+### Max concurrency requests of LLM (for both query and document processing)
+MAX_ASYNC=4
+### Number of parallel processing documents(between 2~10, MAX_ASYNC/4 is recommended)
+MAX_PARALLEL_INSERT=2
+### Max concurrency requests for Embedding
+# EMBEDDING_FUNC_MAX_ASYNC=8
+### Num of chunks send to Embedding in single request
+# EMBEDDING_BATCH_NUM=10
 
+#######################
 ### LLM Configuration
-ENABLE_LLM_CACHE=true
-ENABLE_LLM_CACHE_FOR_EXTRACT=true
+#######################
 ### Time out in seconds for LLM, None for infinite timeout
 TIMEOUT=240
 ### Some models like o1-mini require temperature to be set to 1
 TEMPERATURE=0
-### Max concurrency requests of LLM
-MAX_ASYNC=4
-### MAX_TOKENS: max tokens send to LLM for entity relation summaries (less than context size of the model)
-MAX_TOKENS=32000
 ### LLM Binding type: openai, ollama, lollms, azure_openai
 LLM_BINDING=openai
 LLM_MODEL=gpt-4o
 LLM_BINDING_HOST=https://api.openai.com/v1
 LLM_BINDING_API_KEY=your_api_key
+
+### Set as num_ctx option for Ollama LLM
+# OLLAMA_NUM_CTX=32768
+
 ### Optional for Azure
 # AZURE_OPENAI_API_VERSION=2024-08-01-preview
 # AZURE_OPENAI_DEPLOYMENT=gpt-4o
-### set as num_ctx option for Ollama LLM
-# OLLAMA_NUM_CTX=32768
 
-### Embedding Configuration
+####################################################################################
+### Embedding Configuration (Should not be changed after the first file processed)
+####################################################################################
 ### Embedding Binding type: openai, ollama, lollms, azure_openai
 EMBEDDING_BINDING=ollama
 EMBEDDING_MODEL=bge-m3:latest
@@ -108,51 +130,53 @@ EMBEDDING_DIM=1024
 EMBEDDING_BINDING_API_KEY=your_api_key
 # If the embedding service is deployed within the same Docker stack, use host.docker.internal instead of localhost
 EMBEDDING_BINDING_HOST=http://localhost:11434
-### Num of chunks send to Embedding in single request
-# EMBEDDING_BATCH_NUM=10
-### Max concurrency requests for Embedding
-# EMBEDDING_FUNC_MAX_ASYNC=8
 ### Maximum tokens sent to Embedding for each chunk (no longer in use?)
 # MAX_EMBED_TOKENS=8192
+
 ### Optional for Azure
 # AZURE_EMBEDDING_DEPLOYMENT=text-embedding-3-large
 # AZURE_EMBEDDING_API_VERSION=2023-05-15
 # AZURE_EMBEDDING_ENDPOINT=your_endpoint
 # AZURE_EMBEDDING_API_KEY=your_api_key
 
-###########################
+############################
 ### Data storage selection
-###########################
-### In-memory database with local file persistence(Recommended for small scale deployment)
+############################
+### Default storage (Recommended for small scale deployment)
 # LIGHTRAG_KV_STORAGE=JsonKVStorage
 # LIGHTRAG_DOC_STATUS_STORAGE=JsonDocStatusStorage
 # LIGHTRAG_GRAPH_STORAGE=NetworkXStorage
 # LIGHTRAG_VECTOR_STORAGE=NanoVectorDBStorage
+
+### Redis Storage (Recommended for production deployment)
+# LIGHTRAG_KV_STORAGE=RedisKVStorage
+# LIGHTRAG_DOC_STATUS_STORAGE=RedisDocStatusStorage
+
+### Vector Storage (Recommended for production deployment)
+# LIGHTRAG_VECTOR_STORAGE=MilvusVectorDBStorage
+# LIGHTRAG_VECTOR_STORAGE=QdrantVectorDBStorage
 # LIGHTRAG_VECTOR_STORAGE=FaissVectorDBStorage
+
+### Graph Storage (Recommended for production deployment)
+# LIGHTRAG_GRAPH_STORAGE=Neo4JStorage
+# LIGHTRAG_GRAPH_STORAGE=MemgraphStorage
+
 ### PostgreSQL
 # LIGHTRAG_KV_STORAGE=PGKVStorage
 # LIGHTRAG_DOC_STATUS_STORAGE=PGDocStatusStorage
 # LIGHTRAG_GRAPH_STORAGE=PGGraphStorage
 # LIGHTRAG_VECTOR_STORAGE=PGVectorStorage
+
 ### MongoDB (Vector storage only available on Atlas Cloud)
 # LIGHTRAG_KV_STORAGE=MongoKVStorage
 # LIGHTRAG_DOC_STATUS_STORAGE=MongoDocStatusStorage
 # LIGHTRAG_GRAPH_STORAGE=MongoGraphStorage
 # LIGHTRAG_VECTOR_STORAGE=MongoVectorDBStorage
-### Redis Storage (Recommended for production deployment)
-# LIGHTRAG_KV_STORAGE=RedisKVStorage
-# LIGHTRAG_DOC_STATUS_STORAGE=RedisDocStatusStorage
-### Vector Storage (Recommended for production deployment)
-# LIGHTRAG_VECTOR_STORAGE=MilvusVectorDBStorage
-# LIGHTRAG_VECTOR_STORAGE=QdrantVectorDBStorage
-### Graph Storage (Recommended for production deployment)
-# LIGHTRAG_GRAPH_STORAGE=Neo4JStorage
-# LIGHTRAG_GRAPH_STORAGE=MemgraphStorage
 
 ####################################################################
-### Default workspace for all storage types
-### For the purpose of isolation of data for each LightRAG instance
-### Valid characters: a-z, A-Z, 0-9, and _
+### WORKSPACE setting workspace name for all storage types
+### in the purpose of isolating data from LightRAG instances.
+### Valid workspace name constraints: a-z, A-Z, 0-9, and _
 ####################################################################
 # WORKSPACE=space1
 

examples/rerank_example.py

@@ -9,7 +9,11 @@ Configuration Required:
 2. Set your embedding API key and base URL in embedding_func()
 3. Set your rerank API key and base URL in the rerank configuration
 4. Or use environment variables (.env file):
-   - ENABLE_RERANK=True
+   - RERANK_MODEL=your_rerank_model
+   - RERANK_BINDING_HOST=your_rerank_endpoint
+   - RERANK_BINDING_API_KEY=your_rerank_api_key
+
+Note: Rerank is now controlled per query via the 'enable_rerank' parameter (default: True)
 """
 
 import asyncio
@@ -83,8 +87,7 @@ async def create_rag_with_rerank():
             max_token_size=8192,
             func=embedding_func,
         ),
-        # Simplified Rerank Configuration
-        enable_rerank=True,
+        # Rerank Configuration - provide the rerank function
         rerank_model_func=my_rerank_func,
     )
 
@@ -120,7 +123,6 @@ async def create_rag_with_rerank_model():
             max_token_size=8192,
             func=embedding_func,
         ),
-        enable_rerank=True,
         rerank_model_func=rerank_model.rerank,
     )
 
@@ -130,9 +132,9 @@ async def create_rag_with_rerank_model():
     return rag
 
 
-async def test_rerank_with_different_topk():
+async def test_rerank_with_different_settings():
     """
-    Test rerank functionality with different top_k settings
+    Test rerank functionality with different enable_rerank settings
     """
     print("🚀 Setting up LightRAG with Rerank functionality...")
 
@@ -154,16 +156,41 @@ async def test_rerank_with_different_topk():
     print(f"\n🔍 Testing query: '{query}'")
     print("=" * 80)
 
-    # Test different top_k values to show parameter priority
-    top_k_values = [2, 5, 10]
-
-    for top_k in top_k_values:
-        print(f"\n📊 Testing with QueryParam(top_k={top_k}):")
+    # Test with rerank enabled (default)
+    print("\n📊 Testing with enable_rerank=True (default):")
+    result_with_rerank = await rag.aquery(
+        query,
+        param=QueryParam(
+            mode="naive",
+            top_k=10,
+            chunk_top_k=5,
+            enable_rerank=True,  # Explicitly enable rerank
+        ),
+    )
+    print(f"   Result length: {len(result_with_rerank)} characters")
+    print(f"   Preview: {result_with_rerank[:100]}...")
+
+    # Test with rerank disabled
+    print("\n📊 Testing with enable_rerank=False:")
+    result_without_rerank = await rag.aquery(
+        query,
+        param=QueryParam(
+            mode="naive",
+            top_k=10,
+            chunk_top_k=5,
+            enable_rerank=False,  # Disable rerank
+        ),
+    )
+    print(f"   Result length: {len(result_without_rerank)} characters")
+    print(f"   Preview: {result_without_rerank[:100]}...")
 
-        # Test naive mode with specific top_k
-        result = await rag.aquery(query, param=QueryParam(mode="naive", top_k=top_k))
-        print(f"   Result length: {len(result)} characters")
-        print(f"   Preview: {result[:100]}...")
+    # Test with default settings (enable_rerank defaults to True)
+    print("\n📊 Testing with default settings (enable_rerank defaults to True):")
+    result_default = await rag.aquery(
+        query, param=QueryParam(mode="naive", top_k=10, chunk_top_k=5)
+    )
+    print(f"   Result length: {len(result_default)} characters")
+    print(f"   Preview: {result_default[:100]}...")
 
 
 async def test_direct_rerank():
@@ -209,17 +236,21 @@ async def main():
     print("=" * 60)
 
     try:
-        # Test rerank with different top_k values
-        await test_rerank_with_different_topk()
+        # Test rerank with different enable_rerank settings
+        await test_rerank_with_different_settings()
 
         # Test direct rerank
         await test_direct_rerank()
 
         print("\n✅ Example completed successfully!")
         print("\n💡 Key Points:")
-        print("   ✓ All rerank configurations are contained within rerank_model_func")
-        print("   ✓ Rerank improves document relevance ordering")
-        print("   ✓ Configure API keys within your rerank function")
+        print("   ✓ Rerank is now controlled per query via 'enable_rerank' parameter")
+        print("   ✓ Default value for enable_rerank is True")
+        print("   ✓ Rerank function is configured at LightRAG initialization")
+        print("   ✓ Per-query enable_rerank setting overrides default behavior")
+        print(
+            "   ✓ If enable_rerank=True but no rerank model is configured, a warning is issued"
+        )
         print("   ✓ Monitor API usage and costs when using rerank services")
 
     except Exception as e:

lightrag/api/config.py

@@ -11,6 +11,14 @@ from lightrag.utils import get_env_value
 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,
 )
 
 # use the .env that is inside the current folder
@@ -151,45 +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", 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(
-        "--chunk-top-k",
-        type=int,
-        default=get_env_value("CHUNK_TOP_K", 15, int),
-        help="Number of text chunks to retrieve initially from vector search (default: from env or 15)",
-    )
-    parser.add_argument(
-        "--chunk-rerank-top-k",
-        type=int,
-        default=get_env_value("CHUNK_RERANK_TOP_K", 5, int),
-        help="Number of text chunks to keep after reranking (default: from env or 5)",
-    )
-    parser.add_argument(
-        "--enable-rerank",
-        action="store_true",
-        default=get_env_value("ENABLE_RERANK", False, bool),
-        help="Enable rerank functionality (default: from env or False)",
-    )
-    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",
@@ -321,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

lightrag/api/lightrag_server.py

@@ -292,9 +292,9 @@ def create_app(args):
         ),
     )
 
-    # Configure rerank function if enabled
+    # Configure rerank function if model and API are configured
     rerank_model_func = None
-    if args.enable_rerank and args.rerank_binding_api_key and args.rerank_binding_host:
+    if args.rerank_binding_api_key and args.rerank_binding_host:
         from lightrag.rerank import custom_rerank
 
         async def server_rerank_func(
@@ -312,10 +312,12 @@ def create_app(args):
             )
 
         rerank_model_func = server_rerank_func
-        logger.info(f"Rerank enabled with model: {args.rerank_model}")
-    elif args.enable_rerank:
-        logger.warning(
-            "Rerank enabled but RERANK_BINDING_API_KEY or RERANK_BINDING_HOST not configured. Rerank will be disabled."
+        logger.info(
+            f"Rerank model configured: {args.rerank_model} (can be enabled per query)"
+        )
+    else:
+        logger.info(
+            "Rerank model not configured. Set RERANK_BINDING_API_KEY and RERANK_BINDING_HOST to enable reranking."
         )
 
     # Initialize RAG
@@ -351,7 +353,6 @@ def create_app(args):
             },
             enable_llm_cache_for_entity_extract=args.enable_llm_cache_for_extract,
             enable_llm_cache=args.enable_llm_cache,
-            enable_rerank=args.enable_rerank,
             rerank_model_func=rerank_model_func,
             auto_manage_storages_states=False,
             max_parallel_insert=args.max_parallel_insert,
@@ -381,7 +382,6 @@ def create_app(args):
             },
             enable_llm_cache_for_entity_extract=args.enable_llm_cache_for_extract,
             enable_llm_cache=args.enable_llm_cache,
-            enable_rerank=args.enable_rerank,
             rerank_model_func=rerank_model_func,
             auto_manage_storages_states=False,
             max_parallel_insert=args.max_parallel_insert,
@@ -512,11 +512,13 @@ def create_app(args):
                     "enable_llm_cache": args.enable_llm_cache,
                     "workspace": args.workspace,
                     "max_graph_nodes": args.max_graph_nodes,
-                    # Rerank configuration
-                    "enable_rerank": args.enable_rerank,
-                    "rerank_model": args.rerank_model if args.enable_rerank else None,
+                    # Rerank configuration (based on whether rerank model is configured)
+                    "enable_rerank": rerank_model_func is not None,
+                    "rerank_model": args.rerank_model
+                    if rerank_model_func is not None
+                    else None,
                     "rerank_binding_host": args.rerank_binding_host
-                    if args.enable_rerank
+                    if rerank_model_func is not None
                     else None,
                 },
                 "auth_mode": auth_mode,

lightrag/api/routers/query_routes.py

@@ -52,31 +52,25 @@ class QueryRequest(BaseModel):
     chunk_top_k: Optional[int] = Field(
         ge=1,
         default=None,
-        description="Number of text chunks to retrieve initially from vector search.",
+        description="Number of text chunks to retrieve initially from vector search and keep after reranking.",
     )
 
-    chunk_rerank_top_k: Optional[int] = Field(
-        ge=1,
+    max_entity_tokens: Optional[int] = Field(
         default=None,
-        description="Number of text chunks to keep after reranking.",
-    )
-
-    max_token_for_text_unit: Optional[int] = Field(
-        gt=1,
-        default=None,
-        description="Maximum number of tokens allowed for each retrieved text chunk.",
+        description="Maximum number of tokens allocated for entity context in unified token control system.",
+        ge=1,
     )
 
-    max_token_for_global_context: Optional[int] = Field(
-        gt=1,
+    max_relation_tokens: Optional[int] = Field(
         default=None,
-        description="Maximum number of tokens allocated for relationship descriptions in global retrieval.",
+        description="Maximum number of tokens allocated for relationship context in unified token control system.",
+        ge=1,
     )
 
-    max_token_for_local_context: Optional[int] = Field(
-        gt=1,
+    max_total_tokens: Optional[int] = Field(
         default=None,
-        description="Maximum number of tokens allocated for entity descriptions in local retrieval.",
+        description="Maximum total tokens budget for the entire query context (entities + relations + chunks + system prompt).",
+        ge=1,
     )
 
     conversation_history: Optional[List[Dict[str, Any]]] = Field(
@@ -99,6 +93,11 @@ class QueryRequest(BaseModel):
         description="User-provided prompt for the query. If provided, this will be used instead of the default value from prompt template.",
     )
 
+    enable_rerank: Optional[bool] = Field(
+        default=None,
+        description="Enable reranking for retrieved text chunks. If True but no rerank model is configured, a warning will be issued. Default is True.",
+    )
+
     @field_validator("query", mode="after")
     @classmethod
     def query_strip_after(cls, query: str) -> str: