"""Entity and relationship extraction. This module provides LLM-based entity and relationship extraction from text chunks for knowledge graph construction. """ import asyncio import re import time from collections import Counter, defaultdict from typing import Dict, List, Union, Any from tqdm.asyncio import tqdm as tqdm_async from ..utils import ( logger, clean_str, compute_mdhash_id, encode_string_by_tiktoken, decode_tokens_by_tiktoken, is_float_regex, pack_user_ass_to_openai_messages, split_string_by_multi_markers, ) from ..base import BaseGraphStorage, BaseVectorStorage, TextChunkSchema from ..prompts import GRAPH_FIELD_SEP, PROMPTS # ============================================================================= # Helper Functions for Entity/Relationship Processing # ============================================================================= async def _handle_entity_relation_summary( entity_or_relation_name: str, description: str, global_config: dict, ) -> str: """Summarize entity or relationship descriptions using LLM. If the description exceeds token limits, uses LLM to create a summary. Args: entity_or_relation_name: Name of the entity or relationship description: Full description text global_config: Configuration with LLM settings Returns: Original description if under limit, otherwise summarized text """ use_llm_func: callable = global_config["llm_model_func"] llm_max_tokens = global_config["llm_model_max_token_size"] tiktoken_model_name = global_config["tiktoken_model_name"] summary_max_tokens = global_config["entity_summary_to_max_tokens"] language = global_config["addon_params"].get( "language", PROMPTS["DEFAULT_LANGUAGE"] ) tokens = encode_string_by_tiktoken(description, model_name=tiktoken_model_name) if len(tokens) < summary_max_tokens: return description prompt_template = PROMPTS["summarize_entity_descriptions"] use_description = decode_tokens_by_tiktoken( tokens[:llm_max_tokens], model_name=tiktoken_model_name ) context_base = dict( entity_name=entity_or_relation_name, description_list=use_description.split(GRAPH_FIELD_SEP), language=language, ) use_prompt = prompt_template.format(**context_base) logger.debug(f"Trigger summary: {entity_or_relation_name}") summary = await use_llm_func(use_prompt, max_tokens=summary_max_tokens) return summary async def _handle_single_entity_extraction( record_attributes: list[str], chunk_key: str, ) -> Union[dict, None]: """Parse a single entity record from LLM output. Args: record_attributes: List of record attributes from LLM chunk_key: Source chunk identifier Returns: Entity dictionary or None if invalid """ if len(record_attributes) < 4 or record_attributes[0] != '"entity"': return None entity_name = clean_str(record_attributes[1].lower()) if not entity_name.strip(): return None entity_type = clean_str(record_attributes[2].lower()) entity_description = clean_str(record_attributes[3]) entity_source_id = chunk_key return dict( entity_name=entity_name, entity_type=entity_type, description=entity_description, source_id=entity_source_id, ) async def _handle_single_relationship_extraction( record_attributes: list[str], chunk_key: str, ) -> Union[dict, None]: """Parse a single relationship record from LLM output. Args: record_attributes: List of record attributes from LLM chunk_key: Source chunk identifier Returns: Relationship dictionary or None if invalid """ if len(record_attributes) < 5 or record_attributes[0] != '"relationship"': return None source = clean_str(record_attributes[1].lower()) target = clean_str(record_attributes[2].lower()) edge_description = clean_str(record_attributes[3]) edge_keywords = clean_str(record_attributes[4]) edge_source_id = chunk_key weight = ( float(record_attributes[-1]) if is_float_regex(record_attributes[-1]) else 1.0 ) return dict( src_id=source, tgt_id=target, weight=weight, description=edge_description, keywords=edge_keywords, source_id=edge_source_id, ) async def _merge_nodes_then_upsert( entity_name: str, nodes_data: list[dict], knowledge_graph_inst: BaseGraphStorage, global_config: dict, ) -> dict: """Merge duplicate entity nodes and upsert to knowledge graph. Args: entity_name: Name of the entity nodes_data: List of node data dictionaries to merge knowledge_graph_inst: Knowledge graph storage global_config: Configuration dictionary Returns: Merged node data dictionary """ already_entity_types = [] already_source_ids = [] already_description = [] already_node = await knowledge_graph_inst.get_node(entity_name) if already_node is not None: already_entity_types.append(already_node["entity_type"]) already_source_ids.extend( split_string_by_multi_markers(already_node["source_id"], [GRAPH_FIELD_SEP]) ) already_description.append(already_node["description"]) entity_type = sorted( Counter( [dp["entity_type"] for dp in nodes_data] + already_entity_types ).items(), key=lambda x: x[1], reverse=True, )[0][0] description = GRAPH_FIELD_SEP.join( sorted(set([dp["description"] for dp in nodes_data] + already_description)) ) source_id = GRAPH_FIELD_SEP.join( set([dp["source_id"] for dp in nodes_data] + already_source_ids) ) description = await _handle_entity_relation_summary( entity_name, description, global_config ) node_data = dict( entity_type=entity_type, description=description, source_id=source_id, ) await knowledge_graph_inst.upsert_node( entity_name, node_data=node_data, ) node_data["entity_name"] = entity_name return node_data async def _merge_edges_then_upsert( src_id: str, tgt_id: str, edges_data: list[dict], knowledge_graph_inst: BaseGraphStorage, global_config: dict, ) -> dict: """Merge duplicate relationship edges and upsert to knowledge graph. Args: src_id: Source entity name tgt_id: Target entity name edges_data: List of edge data dictionaries to merge knowledge_graph_inst: Knowledge graph storage global_config: Configuration dictionary Returns: Merged edge data dictionary """ already_weights = [] already_source_ids = [] already_description = [] already_keywords = [] if await knowledge_graph_inst.has_edge(src_id, tgt_id): already_edge = await knowledge_graph_inst.get_edge(src_id, tgt_id) already_weights.append(already_edge["weight"]) already_source_ids.extend( split_string_by_multi_markers(already_edge["source_id"], [GRAPH_FIELD_SEP]) ) already_description.append(already_edge["description"]) already_keywords.extend( split_string_by_multi_markers(already_edge["keywords"], [GRAPH_FIELD_SEP]) ) weight = sum([dp["weight"] for dp in edges_data] + already_weights) description = GRAPH_FIELD_SEP.join( sorted(set([dp["description"] for dp in edges_data] + already_description)) ) keywords = GRAPH_FIELD_SEP.join( sorted(set([dp["keywords"] for dp in edges_data] + already_keywords)) ) source_id = GRAPH_FIELD_SEP.join( set([dp["source_id"] for dp in edges_data] + already_source_ids) ) for need_insert_id in [src_id, tgt_id]: if not (await knowledge_graph_inst.has_node(need_insert_id)): await knowledge_graph_inst.upsert_node( need_insert_id, node_data={ "source_id": source_id, "description": description, "entity_type": '"UNKNOWN"', }, ) description = await _handle_entity_relation_summary( f"({src_id}, {tgt_id})", description, global_config ) await knowledge_graph_inst.upsert_edge( src_id, tgt_id, edge_data=dict( weight=weight, description=description, keywords=keywords, source_id=source_id, ), ) edge_data = dict( src_id=src_id, tgt_id=tgt_id, description=description, keywords=keywords, ) return edge_data # ============================================================================= # Main Extraction Function # ============================================================================= async def extract_entities( chunks: dict[str, TextChunkSchema], knowledge_graph_inst: BaseGraphStorage, entity_vdb: BaseVectorStorage, relationships_vdb: BaseVectorStorage, global_config: dict, ) -> Union[BaseGraphStorage, None]: """Extract entities and relationships from text chunks using LLM. This function processes text chunks in parallel, using an LLM to extract entities and relationships, then merges duplicates and stores them in the knowledge graph and vector databases. Args: chunks: Dictionary of chunk_id -> chunk data knowledge_graph_inst: Knowledge graph storage entity_vdb: Vector database for entity embeddings relationships_vdb: Vector database for relationship embeddings global_config: Configuration dictionary with LLM settings Returns: Updated knowledge graph instance, or None if extraction failed """ use_llm_func: callable = global_config["llm_model_func"] entity_extract_max_gleaning = global_config["entity_extract_max_gleaning"] ordered_chunks = list(chunks.items()) language = global_config["addon_params"].get( "language", PROMPTS["DEFAULT_LANGUAGE"] ) entity_types = global_config["addon_params"].get( "entity_types", PROMPTS["DEFAULT_ENTITY_TYPES"] ) # Check if we should use database schema entity extraction use_database_schema_prompt = global_config["addon_params"].get("use_database_schema_prompt", False) # Unified approach: use database_schema_entity_extraction for all database-related content if use_database_schema_prompt: entity_extract_prompt = PROMPTS["database_schema_entity_extraction"] examples_key = "database_schema_entity_extraction_examples" else: entity_extract_prompt = PROMPTS["entity_extraction"] examples_key = "entity_extraction_examples" example_number = global_config["addon_params"].get("example_number", None) example_index = global_config["addon_params"].get("example_index", None) if example_index is not None and example_index < len(PROMPTS[examples_key]): examples = PROMPTS[examples_key][example_index] elif example_number and example_number < len(PROMPTS[examples_key]): examples = "\n".join( PROMPTS[examples_key][: int(example_number)] ) else: examples = "\n".join(PROMPTS[examples_key]) example_context_base = dict( tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"], record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"], completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"], entity_types=",".join(entity_types), language=language, metadata="metadata", ) examples = examples.format(**example_context_base) context_base = dict( tuple_delimiter=PROMPTS["DEFAULT_TUPLE_DELIMITER"], record_delimiter=PROMPTS["DEFAULT_RECORD_DELIMITER"], completion_delimiter=PROMPTS["DEFAULT_COMPLETION_DELIMITER"], entity_types=",".join(entity_types), examples=examples, language=language, metadata="metadata", ) continue_prompt = PROMPTS["entiti_continue_extraction"] if_loop_prompt = PROMPTS["entiti_if_loop_extraction"] already_processed = 0 already_entities = 0 already_relations = 0 async def _process_single_content(chunk_key_dp: tuple[str, TextChunkSchema]): nonlocal already_processed, already_entities, already_relations chunk_key = chunk_key_dp[0] chunk_dp = chunk_key_dp[1] content = chunk_dp["content"] # Use a safer approach: first format the prompt template with a placeholder for input_text context_with_placeholder = context_base.copy() context_with_placeholder["input_text"] = "{input_text}" formatted_prompt = entity_extract_prompt.format(**context_with_placeholder) hint_prompt = formatted_prompt.replace("{input_text}", content) final_result = await use_llm_func(hint_prompt) history = pack_user_ass_to_openai_messages(hint_prompt, final_result) for now_glean_index in range(entity_extract_max_gleaning): glean_result = await use_llm_func(continue_prompt, history_messages=history) history += pack_user_ass_to_openai_messages(continue_prompt, glean_result) final_result += glean_result if now_glean_index == entity_extract_max_gleaning - 1: break if_loop_result: str = await use_llm_func( if_loop_prompt, history_messages=history ) if_loop_result = if_loop_result.strip().strip('"').strip("'").lower() if if_loop_result != "yes": break records = split_string_by_multi_markers( final_result, [context_base["record_delimiter"], context_base["completion_delimiter"]], ) maybe_nodes = defaultdict(list) maybe_edges = defaultdict(list) for record in records: record = re.search(r"\((.*)\)", record) if record is None: continue record = record.group(1) record_attributes = split_string_by_multi_markers( record, [context_base["tuple_delimiter"]] ) if_entities = await _handle_single_entity_extraction( record_attributes, chunk_key ) if if_entities is not None: maybe_nodes[if_entities["entity_name"]].append(if_entities) continue if_relation = await _handle_single_relationship_extraction( record_attributes, chunk_key ) if if_relation is not None: maybe_edges[(if_relation["src_id"], if_relation["tgt_id"])].append( if_relation ) already_processed += 1 already_entities += len(maybe_nodes) already_relations += len(maybe_edges) now_ticks = PROMPTS["process_tickers"][ already_processed % len(PROMPTS["process_tickers"]) ] print( f"{now_ticks} Processed {already_processed} chunks, {already_entities} entities(duplicated), {already_relations} relations(duplicated)\r", end="", flush=True, ) return dict(maybe_nodes), dict(maybe_edges) results = [] for result in tqdm_async( asyncio.as_completed([_process_single_content(c) for c in ordered_chunks]), total=len(ordered_chunks), desc="Extracting entities from chunks", unit="chunk", ): results.append(await result) maybe_nodes = defaultdict(list) maybe_edges = defaultdict(list) for m_nodes, m_edges in results: for k, v in m_nodes.items(): maybe_nodes[k].extend(v) for k, v in m_edges.items(): maybe_edges[k].extend(v) logger.info("Inserting entities into storage...") all_entities_data = [] for result in tqdm_async( asyncio.as_completed( [ _merge_nodes_then_upsert(k, v, knowledge_graph_inst, global_config) for k, v in maybe_nodes.items() ] ), total=len(maybe_nodes), desc="Inserting entities", unit="entity", ): all_entities_data.append(await result) logger.info("Inserting relationships into storage...") all_relationships_data = [] for result in tqdm_async( asyncio.as_completed( [ _merge_edges_then_upsert( k[0], k[1], v, knowledge_graph_inst, global_config ) for k, v in maybe_edges.items() ] ), total=len(maybe_edges), desc="Inserting relationships", unit="relationship", ): all_relationships_data.append(await result) if not len(all_entities_data) and not len(all_relationships_data): logger.warning( "Didn't extract any entities and relationships, maybe your LLM is not working" ) return None if not len(all_entities_data): logger.warning("Didn't extract any entities") if not len(all_relationships_data): logger.warning("Didn't extract any relationships") if entity_vdb is not None: data_for_vdb = { compute_mdhash_id(dp["entity_name"], prefix="ent-"): { "content": dp["entity_name"] + dp["description"], "entity_name": dp["entity_name"], } for dp in all_entities_data } await entity_vdb.upsert(data_for_vdb) if relationships_vdb is not None: data_for_vdb = { compute_mdhash_id(dp["src_id"] + dp["tgt_id"], prefix="rel-"): { "src_id": dp["src_id"], "tgt_id": dp["tgt_id"], "content": dp["keywords"] + dp["src_id"] + dp["tgt_id"] + dp["description"], } for dp in all_relationships_data } await relationships_vdb.upsert(data_for_vdb) return knowledge_graph_inst