diff --git a/lightrag/kg/postgres_impl.py b/lightrag/kg/postgres_impl.py index 71be88b9e7..eafffb45ab 100644 --- a/lightrag/kg/postgres_impl.py +++ b/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: diff --git a/lightrag/lightrag.py b/lightrag/lightrag.py index e9bea3f0bb..f5ff399b07 100644 --- a/lightrag/lightrag.py +++ b/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"), } ) diff --git a/lightrag/operate.py b/lightrag/operate.py index 207f061450..0a54c0d9cb 100644 --- a/lightrag/operate.py +++ b/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, diff --git a/lightrag/utils_graph.py b/lightrag/utils_graph.py index cedad57581..57ce8ff817 100644 --- a/lightrag/utils_graph.py +++ b/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