gzdaniel commited on
Commit
6a4842e
·
1 Parent(s): 107389b

feat: centralize environment variable defaults in constants.py

Browse files
README-zh.md CHANGED
@@ -294,7 +294,7 @@ class QueryParam:
294
  top_k: int = int(os.getenv("TOP_K", "60"))
295
  """Number of top items to retrieve. Represents entities in 'local' mode and relationships in 'global' mode."""
296
 
297
- chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", "5"))
298
  """Number of text chunks to retrieve initially from vector search and keep after reranking.
299
  If None, defaults to top_k value.
300
  """
 
294
  top_k: int = int(os.getenv("TOP_K", "60"))
295
  """Number of top items to retrieve. Represents entities in 'local' mode and relationships in 'global' mode."""
296
 
297
+ chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", "10"))
298
  """Number of text chunks to retrieve initially from vector search and keep after reranking.
299
  If None, defaults to top_k value.
300
  """
README.md CHANGED
@@ -301,7 +301,7 @@ class QueryParam:
301
  top_k: int = int(os.getenv("TOP_K", "60"))
302
  """Number of top items to retrieve. Represents entities in 'local' mode and relationships in 'global' mode."""
303
 
304
- chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", "5"))
305
  """Number of text chunks to retrieve initially from vector search and keep after reranking.
306
  If None, defaults to top_k value.
307
  """
 
301
  top_k: int = int(os.getenv("TOP_K", "60"))
302
  """Number of top items to retrieve. Represents entities in 'local' mode and relationships in 'global' mode."""
303
 
304
+ chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", "10"))
305
  """Number of text chunks to retrieve initially from vector search and keep after reranking.
306
  If None, defaults to top_k value.
307
  """
env.example CHANGED
@@ -57,9 +57,9 @@ OLLAMA_EMULATING_MODEL_TAG=latest
57
 
58
  # COSINE_THRESHOLD=0.2
59
  ### Number of entities or relations to retrieve from KG
60
- # TOP_K=60
61
  ### Number of text chunks to retrieve initially from vector search and keep after reranking
62
- # CHUNK_TOP_K=5
63
 
64
  ### Enable reranking for retrieved text chunks (default: true)
65
  # ENABLE_RERANK=true
 
57
 
58
  # COSINE_THRESHOLD=0.2
59
  ### Number of entities or relations to retrieve from KG
60
+ # TOP_K=40
61
  ### Number of text chunks to retrieve initially from vector search and keep after reranking
62
+ # CHUNK_TOP_K=10
63
 
64
  ### Enable reranking for retrieved text chunks (default: true)
65
  # ENABLE_RERANK=true
lightrag/api/config.py CHANGED
@@ -11,6 +11,9 @@ from lightrag.utils import get_env_value
11
  from lightrag.constants import (
12
  DEFAULT_WOKERS,
13
  DEFAULT_TIMEOUT,
 
 
 
14
  )
15
 
16
  # use the .env that is inside the current folder
@@ -154,7 +157,7 @@ def parse_args() -> argparse.Namespace:
154
  parser.add_argument(
155
  "--history-turns",
156
  type=int,
157
- default=get_env_value("HISTORY_TURNS", 3, int),
158
  help="Number of conversation history turns to include (default: from env or 3)",
159
  )
160
 
@@ -162,13 +165,13 @@ def parse_args() -> argparse.Namespace:
162
  parser.add_argument(
163
  "--top-k",
164
  type=int,
165
- default=get_env_value("TOP_K", 60, int),
166
  help="Number of most similar results to return (default: from env or 60)",
167
  )
168
  parser.add_argument(
169
  "--chunk-top-k",
170
  type=int,
171
- default=get_env_value("CHUNK_TOP_K", 5, int),
172
  help="Number of text chunks to retrieve initially from vector search and keep after reranking (default: from env or 5)",
173
  )
174
  parser.add_argument(
 
11
  from lightrag.constants import (
12
  DEFAULT_WOKERS,
13
  DEFAULT_TIMEOUT,
14
+ DEFAULT_TOP_K,
15
+ DEFAULT_CHUNK_TOP_K,
16
+ DEFAULT_HISTORY_TURNS,
17
  )
18
 
19
  # use the .env that is inside the current folder
 
157
  parser.add_argument(
158
  "--history-turns",
159
  type=int,
160
+ default=get_env_value("HISTORY_TURNS", DEFAULT_HISTORY_TURNS, int),
161
  help="Number of conversation history turns to include (default: from env or 3)",
162
  )
163
 
 
165
  parser.add_argument(
166
  "--top-k",
167
  type=int,
168
+ default=get_env_value("TOP_K", DEFAULT_TOP_K, int),
169
  help="Number of most similar results to return (default: from env or 60)",
170
  )
171
  parser.add_argument(
172
  "--chunk-top-k",
173
  type=int,
174
+ default=get_env_value("CHUNK_TOP_K", DEFAULT_CHUNK_TOP_K, int),
175
  help="Number of text chunks to retrieve initially from vector search and keep after reranking (default: from env or 5)",
176
  )
177
  parser.add_argument(
lightrag/base.py CHANGED
@@ -14,7 +14,16 @@ from typing import (
14
  )
15
  from .utils import EmbeddingFunc
16
  from .types import KnowledgeGraph
17
- from .constants import GRAPH_FIELD_SEP
 
 
 
 
 
 
 
 
 
18
 
19
  # use the .env that is inside the current folder
20
  # allows to use different .env file for each lightrag instance
@@ -57,21 +66,21 @@ class QueryParam:
57
  stream: bool = False
58
  """If True, enables streaming output for real-time responses."""
59
 
60
- top_k: int = int(os.getenv("TOP_K", "60"))
61
  """Number of top items to retrieve. Represents entities in 'local' mode and relationships in 'global' mode."""
62
 
63
- chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", "5"))
64
  """Number of text chunks to retrieve initially from vector search and keep after reranking.
65
  If None, defaults to top_k value.
66
  """
67
 
68
- max_entity_tokens: int = int(os.getenv("MAX_ENTITY_TOKENS", "10000"))
69
  """Maximum number of tokens allocated for entity context in unified token control system."""
70
 
71
- max_relation_tokens: int = int(os.getenv("MAX_RELATION_TOKENS", "10000"))
72
  """Maximum number of tokens allocated for relationship context in unified token control system."""
73
 
74
- max_total_tokens: int = int(os.getenv("MAX_TOTAL_TOKENS", "32000"))
75
  """Maximum total tokens budget for the entire query context (entities + relations + chunks + system prompt)."""
76
 
77
  hl_keywords: list[str] = field(default_factory=list)
@@ -85,7 +94,7 @@ class QueryParam:
85
  Format: [{"role": "user/assistant", "content": "message"}].
86
  """
87
 
88
- history_turns: int = int(os.getenv("HISTORY_TURNS", "3"))
89
  """Number of complete conversation turns (user-assistant pairs) to consider in the response context."""
90
 
91
  ids: list[str] | None = None
@@ -102,7 +111,7 @@ class QueryParam:
102
  If proivded, this will be use instead of the default vaulue from prompt template.
103
  """
104
 
105
- enable_rerank: bool = os.getenv("ENABLE_RERANK", "true").lower() == "true"
106
  """Enable reranking for retrieved text chunks. If True but no rerank model is configured, a warning will be issued.
107
  Default is True to enable reranking when rerank model is available.
108
  """
 
14
  )
15
  from .utils import EmbeddingFunc
16
  from .types import KnowledgeGraph
17
+ from .constants import (
18
+ GRAPH_FIELD_SEP,
19
+ DEFAULT_TOP_K,
20
+ DEFAULT_CHUNK_TOP_K,
21
+ DEFAULT_MAX_ENTITY_TOKENS,
22
+ DEFAULT_MAX_RELATION_TOKENS,
23
+ DEFAULT_MAX_TOTAL_TOKENS,
24
+ DEFAULT_HISTORY_TURNS,
25
+ DEFAULT_ENABLE_RERANK,
26
+ )
27
 
28
  # use the .env that is inside the current folder
29
  # allows to use different .env file for each lightrag instance
 
66
  stream: bool = False
67
  """If True, enables streaming output for real-time responses."""
68
 
69
+ top_k: int = int(os.getenv("TOP_K", str(DEFAULT_TOP_K)))
70
  """Number of top items to retrieve. Represents entities in 'local' mode and relationships in 'global' mode."""
71
 
72
+ chunk_top_k: int = int(os.getenv("CHUNK_TOP_K", str(DEFAULT_CHUNK_TOP_K)))
73
  """Number of text chunks to retrieve initially from vector search and keep after reranking.
74
  If None, defaults to top_k value.
75
  """
76
 
77
+ max_entity_tokens: int = int(os.getenv("MAX_ENTITY_TOKENS", str(DEFAULT_MAX_ENTITY_TOKENS)))
78
  """Maximum number of tokens allocated for entity context in unified token control system."""
79
 
80
+ max_relation_tokens: int = int(os.getenv("MAX_RELATION_TOKENS", str(DEFAULT_MAX_RELATION_TOKENS)))
81
  """Maximum number of tokens allocated for relationship context in unified token control system."""
82
 
83
+ max_total_tokens: int = int(os.getenv("MAX_TOTAL_TOKENS", str(DEFAULT_MAX_TOTAL_TOKENS)))
84
  """Maximum total tokens budget for the entire query context (entities + relations + chunks + system prompt)."""
85
 
86
  hl_keywords: list[str] = field(default_factory=list)
 
94
  Format: [{"role": "user/assistant", "content": "message"}].
95
  """
96
 
97
+ history_turns: int = int(os.getenv("HISTORY_TURNS", str(DEFAULT_HISTORY_TURNS)))
98
  """Number of complete conversation turns (user-assistant pairs) to consider in the response context."""
99
 
100
  ids: list[str] | None = None
 
111
  If proivded, this will be use instead of the default vaulue from prompt template.
112
  """
113
 
114
+ enable_rerank: bool = os.getenv("ENABLE_RERANK", str(DEFAULT_ENABLE_RERANK).lower()).lower() == "true"
115
  """Enable reranking for retrieved text chunks. If True but no rerank model is configured, a warning will be issued.
116
  Default is True to enable reranking when rerank model is available.
117
  """
lightrag/constants.py CHANGED
@@ -13,6 +13,15 @@ DEFAULT_FORCE_LLM_SUMMARY_ON_MERGE = 6
13
  DEFAULT_WOKERS = 2
14
  DEFAULT_TIMEOUT = 150
15
 
 
 
 
 
 
 
 
 
 
16
  # Separator for graph fields
17
  GRAPH_FIELD_SEP = "<SEP>"
18
 
 
13
  DEFAULT_WOKERS = 2
14
  DEFAULT_TIMEOUT = 150
15
 
16
+ # Query and retrieval configuration defaults
17
+ DEFAULT_TOP_K = 40
18
+ DEFAULT_CHUNK_TOP_K = 10
19
+ DEFAULT_MAX_ENTITY_TOKENS = 10000
20
+ DEFAULT_MAX_RELATION_TOKENS = 10000
21
+ DEFAULT_MAX_TOTAL_TOKENS = 32000
22
+ DEFAULT_HISTORY_TURNS = 3
23
+ DEFAULT_ENABLE_RERANK = True
24
+
25
  # Separator for graph fields
26
  GRAPH_FIELD_SEP = "<SEP>"
27
 
lightrag/operate.py CHANGED
@@ -36,7 +36,12 @@ from .base import (
36
  QueryParam,
37
  )
38
  from .prompt import PROMPTS
39
- from .constants import GRAPH_FIELD_SEP
 
 
 
 
 
40
  from .kg.shared_storage import get_storage_keyed_lock
41
  import time
42
  from dotenv import load_dotenv
@@ -1960,17 +1965,17 @@ async def _build_query_context(
1960
  max_entity_tokens = getattr(
1961
  query_param,
1962
  "max_entity_tokens",
1963
- text_chunks_db.global_config.get("MAX_ENTITY_TOKENS", 10000),
1964
  )
1965
  max_relation_tokens = getattr(
1966
  query_param,
1967
  "max_relation_tokens",
1968
- text_chunks_db.global_config.get("MAX_RELATION_TOKENS", 10000),
1969
  )
1970
  max_total_tokens = getattr(
1971
  query_param,
1972
  "max_total_tokens",
1973
- text_chunks_db.global_config.get("MAX_TOTAL_TOKENS", 32000),
1974
  )
1975
 
1976
  # Truncate entities based on complete JSON serialization
@@ -2688,7 +2693,7 @@ async def naive_query(
2688
  # Calculate dynamic token limit for chunks
2689
  # Get token limits from query_param (with fallback to global_config)
2690
  max_total_tokens = getattr(
2691
- query_param, "max_total_tokens", global_config.get("MAX_TOTAL_TOKENS", 32000)
2692
  )
2693
 
2694
  # Calculate conversation history tokens
 
36
  QueryParam,
37
  )
38
  from .prompt import PROMPTS
39
+ from .constants import (
40
+ GRAPH_FIELD_SEP,
41
+ DEFAULT_MAX_ENTITY_TOKENS,
42
+ DEFAULT_MAX_RELATION_TOKENS,
43
+ DEFAULT_MAX_TOTAL_TOKENS,
44
+ )
45
  from .kg.shared_storage import get_storage_keyed_lock
46
  import time
47
  from dotenv import load_dotenv
 
1965
  max_entity_tokens = getattr(
1966
  query_param,
1967
  "max_entity_tokens",
1968
+ text_chunks_db.global_config.get("MAX_ENTITY_TOKENS", DEFAULT_MAX_ENTITY_TOKENS),
1969
  )
1970
  max_relation_tokens = getattr(
1971
  query_param,
1972
  "max_relation_tokens",
1973
+ text_chunks_db.global_config.get("MAX_RELATION_TOKENS", DEFAULT_MAX_RELATION_TOKENS),
1974
  )
1975
  max_total_tokens = getattr(
1976
  query_param,
1977
  "max_total_tokens",
1978
+ text_chunks_db.global_config.get("MAX_TOTAL_TOKENS", DEFAULT_MAX_TOTAL_TOKENS),
1979
  )
1980
 
1981
  # Truncate entities based on complete JSON serialization
 
2693
  # Calculate dynamic token limit for chunks
2694
  # Get token limits from query_param (with fallback to global_config)
2695
  max_total_tokens = getattr(
2696
+ query_param, "max_total_tokens", global_config.get("MAX_TOTAL_TOKENS", DEFAULT_MAX_TOTAL_TOKENS)
2697
  )
2698
 
2699
  # Calculate conversation history tokens