Magicyuan commited on
Commit
3acbbd4
·
1 Parent(s): d6ef863

feat(lightrag): Implement mix search mode combining knowledge graph and vector retrieval

Browse files

- Add 'mix' mode to QueryParam for hybrid search functionality
- Implement mix_kg_vector_query to combine knowledge graph and vector search results
- Update LightRAG class to handle 'mix' mode queries
- Enhance README with examples and explanations for the new mix search mode
- Introduce new prompt structure for generating responses based on combined search results

Files changed (5) hide show
  1. README.md +14 -1
  2. lightrag/base.py +1 -1
  3. lightrag/lightrag.py +20 -0
  4. lightrag/operate.py +192 -0
  5. lightrag/prompt.py +78 -0
README.md CHANGED
@@ -106,8 +106,21 @@ print(rag.query("What are the top themes in this story?", param=QueryParam(mode=
106
 
107
  # Perform hybrid search
108
  print(rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")))
 
 
 
 
 
 
 
 
 
 
109
  ```
110
 
 
 
 
111
  <details>
112
  <summary> Using Open AI-like APIs </summary>
113
 
@@ -262,7 +275,7 @@ In order to run this experiment on low RAM GPU you should select small model and
262
 
263
  ```python
264
  class QueryParam:
265
- mode: Literal["local", "global", "hybrid", "naive"] = "global"
266
  only_need_context: bool = False
267
  response_type: str = "Multiple Paragraphs"
268
  # Number of top-k items to retrieve; corresponds to entities in "local" mode and relationships in "global" mode.
 
106
 
107
  # Perform hybrid search
108
  print(rag.query("What are the top themes in this story?", param=QueryParam(mode="hybrid")))
109
+
110
+ # Perform mix search (Knowledge Graph + Vector Retrieval)
111
+ # Mix mode combines knowledge graph and vector search:
112
+ # - Uses both structured (KG) and unstructured (vector) information
113
+ # - Provides comprehensive answers by analyzing relationships and context
114
+ # - Supports image content through HTML img tags
115
+ # - Allows control over retrieval depth via top_k parameter
116
+ print(rag.query("What are the top themes in this story?", param=QueryParam(
117
+ mode="mix")))
118
+
119
  ```
120
 
121
+
122
+
123
+
124
  <details>
125
  <summary> Using Open AI-like APIs </summary>
126
 
 
275
 
276
  ```python
277
  class QueryParam:
278
+ mode: Literal["local", "global", "hybrid", "naive", "mix"] = "global"
279
  only_need_context: bool = False
280
  response_type: str = "Multiple Paragraphs"
281
  # Number of top-k items to retrieve; corresponds to entities in "local" mode and relationships in "global" mode.
lightrag/base.py CHANGED
@@ -16,7 +16,7 @@ T = TypeVar("T")
16
 
17
  @dataclass
18
  class QueryParam:
19
- mode: Literal["local", "global", "hybrid", "naive"] = "global"
20
  only_need_context: bool = False
21
  only_need_prompt: bool = False
22
  response_type: str = "Multiple Paragraphs"
 
16
 
17
  @dataclass
18
  class QueryParam:
19
+ mode: Literal["local", "global", "hybrid", "naive", "mix"] = "global"
20
  only_need_context: bool = False
21
  only_need_prompt: bool = False
22
  response_type: str = "Multiple Paragraphs"
lightrag/lightrag.py CHANGED
@@ -16,6 +16,7 @@ from .operate import (
16
  # local_query,global_query,hybrid_query,
17
  kg_query,
18
  naive_query,
 
19
  )
20
 
21
  from .utils import (
@@ -630,6 +631,25 @@ class LightRAG:
630
  embedding_func=None,
631
  ),
632
  )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
633
  else:
634
  raise ValueError(f"Unknown mode {param.mode}")
635
  await self._query_done()
 
16
  # local_query,global_query,hybrid_query,
17
  kg_query,
18
  naive_query,
19
+ mix_kg_vector_query,
20
  )
21
 
22
  from .utils import (
 
631
  embedding_func=None,
632
  ),
633
  )
634
+ elif param.mode == "mix":
635
+ response = await mix_kg_vector_query(
636
+ query,
637
+ self.chunk_entity_relation_graph,
638
+ self.entities_vdb,
639
+ self.relationships_vdb,
640
+ self.chunks_vdb,
641
+ self.text_chunks,
642
+ param,
643
+ asdict(self),
644
+ hashing_kv=self.llm_response_cache
645
+ if self.llm_response_cache
646
+ and hasattr(self.llm_response_cache, "global_config")
647
+ else self.key_string_value_json_storage_cls(
648
+ namespace="llm_response_cache",
649
+ global_config=asdict(self),
650
+ embedding_func=None,
651
+ ),
652
+ )
653
  else:
654
  raise ValueError(f"Unknown mode {param.mode}")
655
  await self._query_done()
lightrag/operate.py CHANGED
@@ -1147,3 +1147,195 @@ async def naive_query(
1147
  )
1148
 
1149
  return response
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1147
  )
1148
 
1149
  return response
1150
+
1151
+
1152
+ async def mix_kg_vector_query(
1153
+ query,
1154
+ knowledge_graph_inst: BaseGraphStorage,
1155
+ entities_vdb: BaseVectorStorage,
1156
+ relationships_vdb: BaseVectorStorage,
1157
+ chunks_vdb: BaseVectorStorage,
1158
+ text_chunks_db: BaseKVStorage[TextChunkSchema],
1159
+ query_param: QueryParam,
1160
+ global_config: dict,
1161
+ hashing_kv: BaseKVStorage = None,
1162
+ ) -> str:
1163
+ """
1164
+ Hybrid retrieval implementation combining knowledge graph and vector search.
1165
+
1166
+ This function performs a hybrid search by:
1167
+ 1. Extracting semantic information from knowledge graph
1168
+ 2. Retrieving relevant text chunks through vector similarity
1169
+ 3. Combining both results for comprehensive answer generation
1170
+ """
1171
+ # 1. Cache handling
1172
+ use_model_func = global_config["llm_model_func"]
1173
+ args_hash = compute_args_hash("mix", query)
1174
+ cached_response, quantized, min_val, max_val = await handle_cache(
1175
+ hashing_kv, args_hash, query, "mix"
1176
+ )
1177
+ if cached_response is not None:
1178
+ return cached_response
1179
+
1180
+ # 2. Execute knowledge graph and vector searches in parallel
1181
+ async def get_kg_context():
1182
+ try:
1183
+ # Reuse keyword extraction logic from kg_query
1184
+ example_number = global_config["addon_params"].get("example_number", None)
1185
+ if example_number and example_number < len(
1186
+ PROMPTS["keywords_extraction_examples"]
1187
+ ):
1188
+ examples = "\n".join(
1189
+ PROMPTS["keywords_extraction_examples"][: int(example_number)]
1190
+ )
1191
+ else:
1192
+ examples = "\n".join(PROMPTS["keywords_extraction_examples"])
1193
+
1194
+ language = global_config["addon_params"].get(
1195
+ "language", PROMPTS["DEFAULT_LANGUAGE"]
1196
+ )
1197
+
1198
+ # Extract keywords using LLM
1199
+ kw_prompt = PROMPTS["keywords_extraction"].format(
1200
+ query=query, examples=examples, language=language
1201
+ )
1202
+ result = await use_model_func(kw_prompt, keyword_extraction=True)
1203
+
1204
+ match = re.search(r"\{.*\}", result, re.DOTALL)
1205
+ if not match:
1206
+ logger.warning(
1207
+ "No JSON-like structure found in keywords extraction result"
1208
+ )
1209
+ return None
1210
+
1211
+ result = match.group(0)
1212
+ keywords_data = json.loads(result)
1213
+ hl_keywords = keywords_data.get("high_level_keywords", [])
1214
+ ll_keywords = keywords_data.get("low_level_keywords", [])
1215
+
1216
+ if not hl_keywords and not ll_keywords:
1217
+ logger.warning("Both high-level and low-level keywords are empty")
1218
+ return None
1219
+
1220
+ # Convert keyword lists to strings
1221
+ ll_keywords_str = ", ".join(ll_keywords) if ll_keywords else ""
1222
+ hl_keywords_str = ", ".join(hl_keywords) if hl_keywords else ""
1223
+
1224
+ # Set query mode based on available keywords
1225
+ if not ll_keywords_str and not hl_keywords_str:
1226
+ return None
1227
+ elif not ll_keywords_str:
1228
+ query_param.mode = "global"
1229
+ elif not hl_keywords_str:
1230
+ query_param.mode = "local"
1231
+ else:
1232
+ query_param.mode = "hybrid"
1233
+
1234
+ # Build knowledge graph context
1235
+ context = await _build_query_context(
1236
+ [ll_keywords_str, hl_keywords_str],
1237
+ knowledge_graph_inst,
1238
+ entities_vdb,
1239
+ relationships_vdb,
1240
+ text_chunks_db,
1241
+ query_param,
1242
+ )
1243
+
1244
+ return context
1245
+
1246
+ except Exception as e:
1247
+ logger.error(f"Error in get_kg_context: {str(e)}")
1248
+ return None
1249
+
1250
+ async def get_vector_context():
1251
+ # Reuse vector search logic from naive_query
1252
+ try:
1253
+ # Reduce top_k for vector search in hybrid mode since we have structured information from KG
1254
+ mix_topk = min(10, query_param.top_k)
1255
+ results = await chunks_vdb.query(query, top_k=mix_topk)
1256
+ if not results:
1257
+ return None
1258
+
1259
+ chunks_ids = [r["id"] for r in results]
1260
+ chunks = await text_chunks_db.get_by_ids(chunks_ids)
1261
+
1262
+ valid_chunks = [
1263
+ chunk for chunk in chunks if chunk is not None and "content" in chunk
1264
+ ]
1265
+
1266
+ if not valid_chunks:
1267
+ return None
1268
+
1269
+ maybe_trun_chunks = truncate_list_by_token_size(
1270
+ valid_chunks,
1271
+ key=lambda x: x["content"],
1272
+ max_token_size=query_param.max_token_for_text_unit,
1273
+ )
1274
+
1275
+ if not maybe_trun_chunks:
1276
+ return None
1277
+
1278
+ return "\n--New Chunk--\n".join([c["content"] for c in maybe_trun_chunks])
1279
+ except Exception as e:
1280
+ logger.error(f"Error in get_vector_context: {e}")
1281
+ return None
1282
+
1283
+ # 3. Execute both retrievals in parallel
1284
+ kg_context, vector_context = await asyncio.gather(
1285
+ get_kg_context(), get_vector_context()
1286
+ )
1287
+
1288
+ # 4. Merge contexts
1289
+ if kg_context is None and vector_context is None:
1290
+ return PROMPTS["fail_response"]
1291
+
1292
+ if query_param.only_need_context:
1293
+ return {"kg_context": kg_context, "vector_context": vector_context}
1294
+
1295
+ # 5. Construct hybrid prompt
1296
+ sys_prompt = PROMPTS["mix_rag_response"].format(
1297
+ kg_context=kg_context
1298
+ if kg_context
1299
+ else "No relevant knowledge graph information found",
1300
+ vector_context=vector_context
1301
+ if vector_context
1302
+ else "No relevant text information found",
1303
+ response_type=query_param.response_type,
1304
+ )
1305
+
1306
+ if query_param.only_need_prompt:
1307
+ return sys_prompt
1308
+
1309
+ # 6. Generate response
1310
+ response = await use_model_func(
1311
+ query,
1312
+ system_prompt=sys_prompt,
1313
+ stream=query_param.stream,
1314
+ )
1315
+
1316
+ if isinstance(response, str) and len(response) > len(sys_prompt):
1317
+ response = (
1318
+ response.replace(sys_prompt, "")
1319
+ .replace("user", "")
1320
+ .replace("model", "")
1321
+ .replace(query, "")
1322
+ .replace("<system>", "")
1323
+ .replace("</system>", "")
1324
+ .strip()
1325
+ )
1326
+
1327
+ # 7. Save cache
1328
+ await save_to_cache(
1329
+ hashing_kv,
1330
+ CacheData(
1331
+ args_hash=args_hash,
1332
+ content=response,
1333
+ prompt=query,
1334
+ quantized=quantized,
1335
+ min_val=min_val,
1336
+ max_val=max_val,
1337
+ mode="mix",
1338
+ ),
1339
+ )
1340
+
1341
+ return response
lightrag/prompt.py CHANGED
@@ -284,3 +284,81 @@ Similarity score criteria:
284
  0.5: Partially related and answer needs modification to be used
285
  Return only a number between 0-1, without any additional content.
286
  """
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
284
  0.5: Partially related and answer needs modification to be used
285
  Return only a number between 0-1, without any additional content.
286
  """
287
+
288
+ PROMPTS["mix_rag_response"] = """---Role Definition---
289
+ You are a professional knowledge integration assistant, responsible for answering questions strictly based on provided knowledge graph and text information. You must follow these rules:
290
+ 1. Only use provided knowledge graph and text information
291
+ 2. Do not use your own knowledge or experience
292
+ 3. Do not make any assumptions or speculations
293
+ 4. Analyze the language used in the user message and respond in the same language
294
+ 5. Include relevant images from the source information using HTML img tags
295
+
296
+ ---Objective---
297
+ Generate comprehensive and accurate answers based on knowledge graph and vector search information.
298
+ First analyze the language of the user's question (Chinese/English/Others), then respond in the same language.
299
+ In the following cases, respond politely with "I apologize, but I am unable to provide a complete answer to this question" in the user's language:
300
+ 1. No relevant information found in provided sources
301
+ 2. Question is beyond the scope of provided information
302
+ 3. Requires knowledge beyond provided information
303
+ 4. Requires speculation or assumptions
304
+
305
+ ---Information Sources---
306
+ 1. Knowledge Graph Analysis Results (Structured Information):
307
+ {kg_context}
308
+
309
+ 2. Vector Search Results (Original Text):
310
+ {vector_context}
311
+
312
+ ---Response Format---
313
+ Target response format and length requirements: {response_type}
314
+ Response language: Analyze user message language and respond in the same language
315
+ Image inclusion: If source information contains relevant images in HTML img tags, include them in the response
316
+
317
+ ---Guidelines---
318
+ 1. Language Recognition and Usage:
319
+ - Carefully analyze the language used in user message
320
+ - If question is in Chinese, respond in Chinese (e.g., "非常抱歉,基于现有信息我无法完整回答这个问题")
321
+ - If question is in English, respond in English
322
+ - If question is in other languages, respond in the same language
323
+
324
+ 2. Information Usage Rules:
325
+ - Must reference both knowledge graph and vector search results
326
+ - Each statement must clearly indicate its source
327
+ - Forbidden to use information outside provided sources
328
+ - If information is insufficient, politely state inability to answer in user's language
329
+ - When relevant images are found in source information, include them using HTML img tags
330
+
331
+ 3. Response Standards:
332
+ - Strictly follow specified format and length requirements
333
+ - Use markdown format for organization
334
+ - Use quotation marks for direct quotes
335
+ - Clearly distinguish between factual statements and sources
336
+ - No speculation or assumptions allowed
337
+ - Preserve and include HTML img tags for relevant images
338
+ - Place images appropriately within the context of the answer
339
+
340
+ 4. Information Integration Requirements:
341
+ - Only integrate directly relevant information
342
+ - No excessive interpretation or reasoning
343
+ - Maintain objectivity, no personal views
344
+ - If information conflicts, note it and prioritize knowledge graph
345
+ - When information is incomplete, clearly state the gaps
346
+ - Include relevant images that support or illustrate the answer
347
+
348
+ 5. Quality Control:
349
+ - Every answer must be traceable to provided sources
350
+ - No vague or uncertain expressions
351
+ - No subjective judgments
352
+ - No filling in information gaps
353
+ - No supplementing with common sense or background knowledge
354
+ - Only include images that are directly relevant to the question
355
+ - Maintain original img tags without modification
356
+
357
+ Processing Flow:
358
+ 1. First identify the language of user message
359
+ 2. Analyze provided knowledge graph and vector search information
360
+ 3. Identify relevant images in HTML img tags from the sources
361
+ 4. Organize and generate response in the same language as user, incorporating relevant images
362
+ 5. If unable to answer, express this politely in user's language with an explanation
363
+
364
+ Remember: It's better to say "I apologize, but I am unable to provide a complete answer to this question" (in the user's language, maintaining politeness) than to use information outside provided sources or make speculations. When including images, only use those that are directly relevant and helpful to the answer."""