yangdx commited on
Commit
c952e16
·
1 Parent(s): 5bffd9b

Add token size limit truncation for node data retrieval

Browse files

• Add truncate_list_by_token_size function
• Limit max tokens for local context
• Add logging for truncation info
• Apply truncation to node_datas list

Files changed (1) hide show
  1. lightrag/operate.py +11 -0
lightrag/operate.py CHANGED
@@ -1108,6 +1108,17 @@ async def _get_node_data(
1108
  node_datas, query_param, knowledge_graph_inst
1109
  ),
1110
  )
 
 
 
 
 
 
 
 
 
 
 
1111
  logger.info(
1112
  f"Local query uses {len(node_datas)} entites, {len(use_relations)} relations, {len(use_text_units)} chunks"
1113
  )
 
1108
  node_datas, query_param, knowledge_graph_inst
1109
  ),
1110
  )
1111
+
1112
+ len_node_datas = len(node_datas)
1113
+ node_datas = truncate_list_by_token_size(
1114
+ node_datas,
1115
+ key=lambda x: x["description"],
1116
+ max_token_size=query_param.max_token_for_local_context,
1117
+ )
1118
+ logger.info(
1119
+ f"Truncate entities from {len_node_datas} to {len(node_datas)} (max tokens:{query_param.max_token_for_local_context})"
1120
+ )
1121
+
1122
  logger.info(
1123
  f"Local query uses {len(node_datas)} entites, {len(use_relations)} relations, {len(use_text_units)} chunks"
1124
  )