diff --git a/env.example b/env.example index f0b13bb6..7b09aa38 100644 --- a/env.example +++ b/env.example @@ -67,6 +67,9 @@ ENABLE_LLM_CACHE=true ### Reranker configuration (Set ENABLE_RERANK to true in reranking model is configed) # ENABLE_RERANK=True +### Minimum rerank score for document chunk exclusion (set to 0.0 to keep all chunks, 0.6 or above if LLM is not strong enought) +# MIN_RERANK_SCORE=0.0 +### Rerank model configuration (required when ENABLE_RERANK=True) # RERANK_MODEL=jina-reranker-v2-base-multilingual # RERANK_BINDING_HOST=https://api.jina.ai/v1/rerank # RERANK_BINDING_API_KEY=your_rerank_api_key_here diff --git a/lightrag/api/config.py b/lightrag/api/config.py index 98b817c1..77fbbf8b 100644 --- a/lightrag/api/config.py +++ b/lightrag/api/config.py @@ -19,6 +19,7 @@ from lightrag.constants import ( DEFAULT_MAX_TOTAL_TOKENS, DEFAULT_COSINE_THRESHOLD, DEFAULT_RELATED_CHUNK_NUMBER, + DEFAULT_MIN_RERANK_SCORE, ) # use the .env that is inside the current folder @@ -290,6 +291,11 @@ 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) + # 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) diff --git a/lightrag/constants.py b/lightrag/constants.py index 35d10825..f98eb937 100644 --- a/lightrag/constants.py +++ b/lightrag/constants.py @@ -23,6 +23,9 @@ DEFAULT_ENABLE_RERANK = True DEFAULT_COSINE_THRESHOLD = 0.2 DEFAULT_RELATED_CHUNK_NUMBER = 5 +# Rerank configuration defaults +DEFAULT_MIN_RERANK_SCORE = 0.0 + # Separator for graph fields GRAPH_FIELD_SEP = "" diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index 8bd6857d..9df35175 100644 --- a/lightrag/lightrag.py +++ b/lightrag/lightrag.py @@ -31,6 +31,7 @@ from lightrag.constants import ( DEFAULT_MAX_TOTAL_TOKENS, DEFAULT_COSINE_THRESHOLD, DEFAULT_RELATED_CHUNK_NUMBER, + DEFAULT_MIN_RERANK_SCORE, ) from lightrag.utils import get_env_value @@ -284,6 +285,11 @@ class LightRAG: rerank_model_func: Callable[..., object] | None = field(default=None) """Function for reranking retrieved documents. All rerank configurations (model name, API keys, top_k, etc.) should be included in this function. Optional.""" + min_rerank_score: float = field( + default=get_env_value("MIN_RERANK_SCORE", DEFAULT_MIN_RERANK_SCORE, float) + ) + """Minimum rerank score threshold for filtering chunks after reranking.""" + # Storage # --- diff --git a/lightrag/utils.py b/lightrag/utils.py index 4a91f229..621cddee 100644 --- a/lightrag/utils.py +++ b/lightrag/utils.py @@ -1871,7 +1871,29 @@ async def process_chunks_unified( top_n=rerank_top_k, ) - # 2. Apply chunk_top_k limiting if specified + # 2. Filter by minimum rerank score if reranking is enabled + if query_param.enable_rerank and unique_chunks: + min_rerank_score = global_config.get("min_rerank_score", 0.5) + original_count = len(unique_chunks) + + # Filter chunks with score below threshold + filtered_chunks = [] + for chunk in unique_chunks: + rerank_score = chunk.get("rerank_score", 1.0) # Default to 1.0 if no score + if rerank_score >= min_rerank_score: + filtered_chunks.append(chunk) + + unique_chunks = filtered_chunks + filtered_count = original_count - len(unique_chunks) + + if filtered_count > 0: + logger.info( + f"Rerank filtering remained: {len(unique_chunks)} chunks (min rerank score: {min_rerank_score})" + ) + if not unique_chunks: + return [] + + # 3. Apply chunk_top_k limiting if specified if query_param.chunk_top_k is not None and query_param.chunk_top_k > 0: if len(unique_chunks) > query_param.chunk_top_k: unique_chunks = unique_chunks[: query_param.chunk_top_k] @@ -1879,7 +1901,7 @@ async def process_chunks_unified( f"Kept chunk_top-k: {len(unique_chunks)} chunks (deduplicated original: {origin_count})" ) - # 3. Token-based final truncation + # 4. Token-based final truncation tokenizer = global_config.get("tokenizer") if tokenizer and unique_chunks: # Set default chunk_token_limit if not provided