Enhance entity extraction and keyword extraction prompts

lightrag/kg/postgres_impl.py

@@ -4796,15 +4796,19 @@ async def get_nodes_edges_batch(
 
             for result in outgoing_results:
                 if result["node_id"] and result["connected_id"]:
-                    edges_norm[result["node_id"]].append(
-                        (result["node_id"], result["connected_id"])
-                    )
+                    normalized_node_id = self._normalize_node_id(result["node_id"])
+                    if normalized_node_id in edges_norm:
+                        edges_norm[normalized_node_id].append(
+                            (result["node_id"], result["connected_id"])
+                        )
 
             for result in incoming_results:
                 if result["node_id"] and result["connected_id"]:
-                    edges_norm[result["node_id"]].append(
-                        (result["connected_id"], result["node_id"])
-                    )
+                    normalized_node_id = self._normalize_node_id(result["node_id"])
+                    if normalized_node_id in edges_norm:
+                        edges_norm[normalized_node_id].append(
+                            (result["connected_id"], result["node_id"])
+                        )
 
         out: dict[str, list[tuple[str, str]]] = {}
         for orig in node_ids:

lightrag/lightrag.py

@@ -422,6 +422,21 @@ class LightRAG:
                 "SUMMARY_LANGUAGE", DEFAULT_SUMMARY_LANGUAGE, str
             ),
             "entity_types": get_env_value("ENTITY_TYPES", DEFAULT_ENTITY_TYPES, list),
+            "entity_extraction_system_prompt": PROMPTS.get(
+                "entity_extraction_system_prompt"
+            ),
+            "entity_extraction_user_prompt": PROMPTS.get(
+                "entity_extraction_user_prompt"
+            ),
+            "entity_continue_extraction_user_prompt": PROMPTS.get(
+                "entity_continue_extraction_user_prompt"
+            ),
+            "entity_extraction_examples": PROMPTS.get("entity_extraction_examples"),
+            "summarize_entity_descriptions": PROMPTS.get(
+                "summarize_entity_descriptions"
+            ),
+            "keywords_extraction": PROMPTS.get("keywords_extraction"),
+            "keywords_extraction_examples": PROMPTS.get("keywords_extraction_examples"),
         }
     )
 

lightrag/operate.py

@@ -320,7 +320,7 @@ async def _summarize_descriptions(
 
     summary_length_recommended = global_config["summary_length_recommended"]
 
-    prompt_template = PROMPTS["summarize_entity_descriptions"]
+    prompt_template = global_config["addon_params"]["summarize_entity_descriptions"]
 
     # Convert descriptions to JSONL format and apply token-based truncation
     tokenizer = global_config["tokenizer"]
@@ -2791,7 +2791,7 @@ async def extract_entities(
         "entity_types", DEFAULT_ENTITY_TYPES
     )
 
-    examples = "\n".join(PROMPTS["entity_extraction_examples"])
+    examples = "\n".join(global_config["addon_params"]["entity_extraction_examples"])
 
     example_context_base = dict(
         tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"],
@@ -2833,14 +2833,14 @@ async def _process_single_content(chunk_key_dp: tuple[str, TextChunkSchema]):
 
         # Get initial extraction
         # Format system prompt without input_text for each chunk (enables OpenAI prompt caching across chunks)
-        entity_extraction_system_prompt = PROMPTS[
+        entity_extraction_system_prompt = global_config["addon_params"][
             "entity_extraction_system_prompt"
         ].format(**context_base)
         # Format user prompts with input_text for each chunk
-        entity_extraction_user_prompt = PROMPTS["entity_extraction_user_prompt"].format(
-            **{**context_base, "input_text": content}
-        )
-        entity_continue_extraction_user_prompt = PROMPTS[
+        entity_extraction_user_prompt = global_config["addon_params"][
+            "entity_extraction_user_prompt"
+        ].format(**{**context_base, "input_text": content})
+        entity_continue_extraction_user_prompt = global_config["addon_params"][
             "entity_continue_extraction_user_prompt"
         ].format(**{**context_base, "input_text": content})
 
@@ -3267,7 +3267,7 @@ async def extract_keywords_only(
     """
 
     # 1. Build the examples
-    examples = "\n".join(PROMPTS["keywords_extraction_examples"])
+    examples = "\n".join(global_config["addon_params"]["keywords_extraction_examples"])
 
     language = global_config["addon_params"].get("language", DEFAULT_SUMMARY_LANGUAGE)
 
@@ -3293,7 +3293,7 @@ async def extract_keywords_only(
             )
 
     # 3. Build the keyword-extraction prompt
-    kw_prompt = PROMPTS["keywords_extraction"].format(
+    kw_prompt = global_config["addon_params"]["keywords_extraction"].format(
         query=text,
         examples=examples,
         language=language,

lightrag/utils_graph.py

@@ -1258,6 +1258,9 @@ async def _merge_entities_impl(
                 if src == entity_name:
                     edge_data = await chunk_entity_relation_graph.get_edge(src, tgt)
                     all_relations.append((src, tgt, edge_data))
+                elif tgt == entity_name:
+                    edge_data = await chunk_entity_relation_graph.get_edge(src, tgt)
+                    all_relations.append((src, tgt, edge_data))
 
     # 5. Create or update the target entity
     merged_entity_data["entity_id"] = target_entity