gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #1046 from danielaskdd/fix-extraction-problem

Browse Source 
Improve Entity Extraction Robustness for Truncated LLM Responses
This commit is contained in:
zrguo
2025-03-11 15:21:38 +08:00
committed by GitHub
parent f5f21f6f9b 9d1dc2c9c3
commit 141de28015
2 changed files with 213 additions and 230 deletions
Download Patch File Download Diff File Expand all files Collapse all files

317
lightrag/kg/neo4j_impl.py
View File

@@ -176,23 +176,6 @@ class Neo4JStorage(BaseGraphStorage):
# Noe4J handles persistence automatically # Noe4J handles persistence automatically
pass pass
def _ensure_label(self, label: str) -> str:
"""Ensure a label is valid
Args:
label: The label to validate
Returns:
str: The cleaned label
Raises:
ValueError: If label is empty after cleaning
"""
clean_label = label.strip('"')
if not clean_label:
raise ValueError("Neo4j: Label cannot be empty")
return clean_label
async def has_node(self, node_id: str) -> bool: async def has_node(self, node_id: str) -> bool:
""" """
Check if a node with the given label exists in the database Check if a node with the given label exists in the database
@@ -207,20 +190,17 @@ class Neo4JStorage(BaseGraphStorage):
ValueError: If node_id is invalid ValueError: If node_id is invalid
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
entity_name_label = self._ensure_label(node_id)
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
query = f"MATCH (n:`{entity_name_label}`) RETURN count(n) > 0 AS node_exists" query = "MATCH (n:base {entity_id: $entity_id}) RETURN count(n) > 0 AS node_exists"
result = await session.run(query) result = await session.run(query, entity_id=node_id)
single_result = await result.single() single_result = await result.single()
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
return single_result["node_exists"] return single_result["node_exists"]
except Exception as e: except Exception as e:
logger.error( logger.error(f"Error checking node existence for {node_id}: {str(e)}")
f"Error checking node existence for {entity_name_label}: {str(e)}"
)
await result.consume() # Ensure results are consumed even on error await result.consume() # Ensure results are consumed even on error
raise raise
@@ -239,24 +219,25 @@ class Neo4JStorage(BaseGraphStorage):
ValueError: If either node_id is invalid ValueError: If either node_id is invalid
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
entity_name_label_source = self._ensure_label(source_node_id)
entity_name_label_target = self._ensure_label(target_node_id)
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
query = ( query = (
f"MATCH (a:`{entity_name_label_source}`)-[r]-(b:`{entity_name_label_target}`) " "MATCH (a:base {entity_id: $source_entity_id})-[r]-(b:base {entity_id: $target_entity_id}) "
"RETURN COUNT(r) > 0 AS edgeExists" "RETURN COUNT(r) > 0 AS edgeExists"
) )
result = await session.run(query) result = await session.run(
query,
source_entity_id=source_node_id,
target_entity_id=target_node_id,
)
single_result = await result.single() single_result = await result.single()
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
return single_result["edgeExists"] return single_result["edgeExists"]
except Exception as e: except Exception as e:
logger.error( logger.error(
f"Error checking edge existence between {entity_name_label_source} and {entity_name_label_target}: {str(e)}" f"Error checking edge existence between {source_node_id} and {target_node_id}: {str(e)}"
) )
await result.consume() # Ensure results are consumed even on error await result.consume() # Ensure results are consumed even on error
raise raise
@@ -275,13 +256,12 @@ class Neo4JStorage(BaseGraphStorage):
ValueError: If node_id is invalid ValueError: If node_id is invalid
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
entity_name_label = self._ensure_label(node_id)
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
query = f"MATCH (n:`{entity_name_label}` {{entity_id: $entity_id}}) RETURN n" query = "MATCH (n:base {entity_id: $entity_id}) RETURN n"
result = await session.run(query, entity_id=entity_name_label) result = await session.run(query, entity_id=node_id)
try: try:
records = await result.fetch( records = await result.fetch(
2 2
@@ -289,20 +269,25 @@ class Neo4JStorage(BaseGraphStorage):
if len(records) > 1: if len(records) > 1:
logger.warning( logger.warning(
f"Multiple nodes found with label '{entity_name_label}'. Using first node." f"Multiple nodes found with label '{node_id}'. Using first node."
) )
if records: if records:
node = records[0]["n"] node = records[0]["n"]
node_dict = dict(node) node_dict = dict(node)
logger.debug( # Remove base label from labels list if it exists
f"{inspect.currentframe().f_code.co_name}: query: {query}, result: {node_dict}" if "labels" in node_dict:
) node_dict["labels"] = [
label
for label in node_dict["labels"]
if label != "base"
]
logger.debug(f"Neo4j query node {query} return: {node_dict}")
return node_dict return node_dict
return None return None
finally: finally:
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
except Exception as e: except Exception as e:
logger.error(f"Error getting node for {entity_name_label}: {str(e)}") logger.error(f"Error getting node for {node_id}: {str(e)}")
raise raise
async def node_degree(self, node_id: str) -> int: async def node_degree(self, node_id: str) -> int:
@@ -320,43 +305,32 @@ class Neo4JStorage(BaseGraphStorage):
ValueError: If node_id is invalid ValueError: If node_id is invalid
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
entity_name_label = self._ensure_label(node_id)
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
query = f""" query = """
MATCH (n:`{entity_name_label}`) MATCH (n:base {entity_id: $entity_id})
OPTIONAL MATCH (n)-[r]-() OPTIONAL MATCH (n)-[r]-()
RETURN n, COUNT(r) AS degree RETURN COUNT(r) AS degree
""" """
result = await session.run(query) result = await session.run(query, entity_id=node_id)
try: try:
records = await result.fetch(100) record = await result.single()
if not records: if not record:
logger.warning( logger.warning(f"No node found with label '{node_id}'")
f"No node found with label '{entity_name_label}'"
)
return 0 return 0
if len(records) > 1: degree = record["degree"]
logger.warning(
f"Multiple nodes ({len(records)}) found with label '{entity_name_label}', using first node's degree"
)
degree = records[0]["degree"]
logger.debug( logger.debug(
f"{inspect.currentframe().f_code.co_name}:query:{query}:result:{degree}" "Neo4j query node degree for {node_id} return: {degree}"
) )
return degree return degree
finally: finally:
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
except Exception as e: except Exception as e:
logger.error( logger.error(f"Error getting node degree for {node_id}: {str(e)}")
f"Error getting node degree for {entity_name_label}: {str(e)}"
)
raise raise
async def edge_degree(self, src_id: str, tgt_id: str) -> int: async def edge_degree(self, src_id: str, tgt_id: str) -> int:
@@ -369,11 +343,8 @@ class Neo4JStorage(BaseGraphStorage):
Returns: Returns:
int: Sum of the degrees of both nodes int: Sum of the degrees of both nodes
""" """
entity_name_label_source = self._ensure_label(src_id) src_degree = await self.node_degree(src_id)
entity_name_label_target = self._ensure_label(tgt_id) trg_degree = await self.node_degree(tgt_id)
src_degree = await self.node_degree(entity_name_label_source)
trg_degree = await self.node_degree(entity_name_label_target)
# Convert None to 0 for addition # Convert None to 0 for addition
src_degree = 0 if src_degree is None else src_degree src_degree = 0 if src_degree is None else src_degree
@@ -399,24 +370,24 @@ class Neo4JStorage(BaseGraphStorage):
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
try: try:
entity_name_label_source = self._ensure_label(source_node_id)
entity_name_label_target = self._ensure_label(target_node_id)
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
query = f""" query = """
MATCH (start:`{entity_name_label_source}`)-[r]-(end:`{entity_name_label_target}`) MATCH (start:base {entity_id: $source_entity_id})-[r]-(end:base {entity_id: $target_entity_id})
RETURN properties(r) as edge_properties RETURN properties(r) as edge_properties
""" """
result = await session.run(
result = await session.run(query) query,
source_entity_id=source_node_id,
target_entity_id=target_node_id,
)
try: try:
records = await result.fetch(2) records = await result.fetch(2)
if len(records) > 1: if len(records) > 1:
logger.warning( logger.warning(
f"Multiple edges found between '{entity_name_label_source}' and '{entity_name_label_target}'. Using first edge." f"Multiple edges found between '{source_node_id}' and '{target_node_id}'. Using first edge."
) )
if records: if records:
try: try:
@@ -433,7 +404,7 @@ class Neo4JStorage(BaseGraphStorage):
if key not in edge_result: if key not in edge_result:
edge_result[key] = default_value edge_result[key] = default_value
logger.warning( logger.warning(
f"Edge between {entity_name_label_source} and {entity_name_label_target} " f"Edge between {source_node_id} and {target_node_id} "
f"missing {key}, using default: {default_value}" f"missing {key}, using default: {default_value}"
) )
@@ -443,8 +414,8 @@ class Neo4JStorage(BaseGraphStorage):
return edge_result return edge_result
except (KeyError, TypeError, ValueError) as e: except (KeyError, TypeError, ValueError) as e:
logger.error( logger.error(
f"Error processing edge properties between {entity_name_label_source} " f"Error processing edge properties between {source_node_id} "
f"and {entity_name_label_target}: {str(e)}" f"and {target_node_id}: {str(e)}"
) )
# Return default edge properties on error # Return default edge properties on error
return { return {
@@ -455,7 +426,7 @@ class Neo4JStorage(BaseGraphStorage):
} }
logger.debug( logger.debug(
f"{inspect.currentframe().f_code.co_name}: No edge found between {entity_name_label_source} and {entity_name_label_target}" f"{inspect.currentframe().f_code.co_name}: No edge found between {source_node_id} and {target_node_id}"
) )
# Return default edge properties when no edge found # Return default edge properties when no edge found
return { return {
@@ -488,29 +459,33 @@ class Neo4JStorage(BaseGraphStorage):
Exception: If there is an error executing the query Exception: If there is an error executing the query
""" """
try: try:
node_label = self._ensure_label(source_node_id)
query = f"""MATCH (n:`{node_label}`)
OPTIONAL MATCH (n)-[r]-(connected)
RETURN n, r, connected"""
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
results = await session.run(query) query = """MATCH (n:base {entity_id: $entity_id})
edges = [] OPTIONAL MATCH (n)-[r]-(connected:base)
WHERE connected.entity_id IS NOT NULL
RETURN n, r, connected"""
results = await session.run(query, entity_id=source_node_id)
edges = []
async for record in results: async for record in results:
source_node = record["n"] source_node = record["n"]
connected_node = record["connected"] connected_node = record["connected"]
# Skip if either node is None
if not source_node or not connected_node:
continue
source_label = ( source_label = (
list(source_node.labels)[0] if source_node.labels else None source_node.get("entity_id")
if source_node.get("entity_id")
else None
) )
target_label = ( target_label = (
list(connected_node.labels)[0] connected_node.get("entity_id")
if connected_node and connected_node.labels if connected_node.get("entity_id")
else None else None
) )
@@ -520,7 +495,9 @@ class Neo4JStorage(BaseGraphStorage):
await results.consume() # Ensure results are consumed await results.consume() # Ensure results are consumed
return edges return edges
except Exception as e: except Exception as e:
logger.error(f"Error getting edges for node {node_label}: {str(e)}") logger.error(
f"Error getting edges for node {source_node_id}: {str(e)}"
)
await results.consume() # Ensure results are consumed even on error await results.consume() # Ensure results are consumed even on error
raise raise
except Exception as e: except Exception as e:
@@ -547,8 +524,9 @@ class Neo4JStorage(BaseGraphStorage):
node_id: The unique identifier for the node (used as label) node_id: The unique identifier for the node (used as label)
node_data: Dictionary of node properties node_data: Dictionary of node properties
""" """
label = self._ensure_label(node_id)
properties = node_data properties = node_data
entity_type = properties["entity_type"]
entity_id = properties["entity_id"]
if "entity_id" not in properties: if "entity_id" not in properties:
raise ValueError("Neo4j: node properties must contain an 'entity_id' field") raise ValueError("Neo4j: node properties must contain an 'entity_id' field")
@@ -556,13 +534,17 @@ class Neo4JStorage(BaseGraphStorage):
async with self._driver.session(database=self._DATABASE) as session: async with self._driver.session(database=self._DATABASE) as session:
async def execute_upsert(tx: AsyncManagedTransaction): async def execute_upsert(tx: AsyncManagedTransaction):
query = f""" query = (
MERGE (n:`{label}` {{entity_id: $properties.entity_id}})
SET n += $properties
""" """
MERGE (n:base {entity_id: $properties.entity_id})
SET n += $properties
SET n:`%s`
"""
% entity_type
)
result = await tx.run(query, properties=properties) result = await tx.run(query, properties=properties)
logger.debug( logger.debug(
f"Upserted node with label '{label}' and properties: {properties}" f"Upserted node with entity_id '{entity_id}' and properties: {properties}"
) )
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
@@ -583,52 +565,6 @@ class Neo4JStorage(BaseGraphStorage):
) )
), ),
) )
async def _get_unique_node_entity_id(self, node_label: str) -> str:
"""
Get the entity_id of a node with the given label, ensuring the node is unique.
Args:
node_label (str): Label of the node to check
Returns:
str: The entity_id of the unique node
Raises:
ValueError: If no node with the given label exists or if multiple nodes have the same label
"""
async with self._driver.session(
database=self._DATABASE, default_access_mode="READ"
) as session:
query = f"""
MATCH (n:`{node_label}`)
RETURN n, count(n) as node_count
"""
result = await session.run(query)
try:
records = await result.fetch(
2
) # We only need to know if there are 0, 1, or >1 nodes
if not records or records[0]["node_count"] == 0:
raise ValueError(
f"Neo4j: node with label '{node_label}' does not exist"
)
if records[0]["node_count"] > 1:
raise ValueError(
f"Neo4j: multiple nodes found with label '{node_label}', cannot determine unique node"
)
node = records[0]["n"]
if "entity_id" not in node:
raise ValueError(
f"Neo4j: node with label '{node_label}' does not have an entity_id property"
)
return node["entity_id"]
finally:
await result.consume() # Ensure result is fully consumed
@retry( @retry(
stop=stop_after_attempt(3), stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=4, max=10), wait=wait_exponential(multiplier=1, min=4, max=10),
@@ -657,38 +593,30 @@ class Neo4JStorage(BaseGraphStorage):
Raises: Raises:
ValueError: If either source or target node does not exist or is not unique ValueError: If either source or target node does not exist or is not unique
""" """
source_label = self._ensure_label(source_node_id)
target_label = self._ensure_label(target_node_id)
edge_properties = edge_data
# Get entity_ids for source and target nodes, ensuring they are unique
source_entity_id = await self._get_unique_node_entity_id(source_label)
target_entity_id = await self._get_unique_node_entity_id(target_label)
try: try:
edge_properties = edge_data
async with self._driver.session(database=self._DATABASE) as session: async with self._driver.session(database=self._DATABASE) as session:
async def execute_upsert(tx: AsyncManagedTransaction): async def execute_upsert(tx: AsyncManagedTransaction):
query = f""" query = """
MATCH (source:`{source_label}` {{entity_id: $source_entity_id}}) MATCH (source:base {entity_id: $source_entity_id})
WITH source WITH source
MATCH (target:`{target_label}` {{entity_id: $target_entity_id}}) MATCH (target:base {entity_id: $target_entity_id})
MERGE (source)-[r:DIRECTED]-(target) MERGE (source)-[r:DIRECTED]-(target)
SET r += $properties SET r += $properties
RETURN r, source, target RETURN r, source, target
""" """
result = await tx.run( result = await tx.run(
query, query,
source_entity_id=source_entity_id, source_entity_id=source_node_id,
target_entity_id=target_entity_id, target_entity_id=target_node_id,
properties=edge_properties, properties=edge_properties,
) )
try: try:
records = await result.fetch(100) records = await result.fetch(2)
if records: if records:
logger.debug( logger.debug(
f"Upserted edge from '{source_label}' (entity_id: {source_entity_id}) " f"Upserted edge from '{source_node_id}' to '{target_node_id}'"
f"to '{target_label}' (entity_id: {target_entity_id}) "
f"with properties: {edge_properties}" f"with properties: {edge_properties}"
) )
finally: finally:
@@ -726,7 +654,6 @@ class Neo4JStorage(BaseGraphStorage):
Returns: Returns:
KnowledgeGraph: Complete connected subgraph for specified node KnowledgeGraph: Complete connected subgraph for specified node
""" """
label = node_label.strip('"')
result = KnowledgeGraph() result = KnowledgeGraph()
seen_nodes = set() seen_nodes = set()
seen_edges = set() seen_edges = set()
@@ -735,7 +662,7 @@ class Neo4JStorage(BaseGraphStorage):
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
try: try:
if label == "*": if node_label == "*":
main_query = """ main_query = """
MATCH (n) MATCH (n)
OPTIONAL MATCH (n)-[r]-() OPTIONAL MATCH (n)-[r]-()
@@ -760,12 +687,11 @@ class Neo4JStorage(BaseGraphStorage):
# Main query uses partial matching # Main query uses partial matching
main_query = """ main_query = """
MATCH (start) MATCH (start)
WHERE any(label IN labels(start) WHERE WHERE
CASE CASE
WHEN $inclusive THEN label CONTAINS $label WHEN $inclusive THEN start.entity_id CONTAINS $entity_id
ELSE label = $label ELSE start.entity_id = $entity_id
END END
)
WITH start WITH start
CALL apoc.path.subgraphAll(start, { CALL apoc.path.subgraphAll(start, {
relationshipFilter: '', relationshipFilter: '',
@@ -799,7 +725,7 @@ class Neo4JStorage(BaseGraphStorage):
main_query, main_query,
{ {
"max_nodes": MAX_GRAPH_NODES, "max_nodes": MAX_GRAPH_NODES,
"label": label, "entity_id": node_label,
"inclusive": inclusive, "inclusive": inclusive,
"max_depth": max_depth, "max_depth": max_depth,
"min_degree": min_degree, "min_degree": min_degree,
@@ -818,7 +744,11 @@ class Neo4JStorage(BaseGraphStorage):
result.nodes.append( result.nodes.append(
KnowledgeGraphNode( KnowledgeGraphNode(
id=f"{node_id}", id=f"{node_id}",
labels=list(node.labels), labels=[
label
for label in node.labels
if label != "base"
],
properties=dict(node), properties=dict(node),
) )
) )
@@ -849,7 +779,7 @@ class Neo4JStorage(BaseGraphStorage):
except neo4jExceptions.ClientError as e: except neo4jExceptions.ClientError as e:
logger.warning(f"APOC plugin error: {str(e)}") logger.warning(f"APOC plugin error: {str(e)}")
if label != "*": if node_label != "*":
logger.warning( logger.warning(
"Neo4j: falling back to basic Cypher recursive search..." "Neo4j: falling back to basic Cypher recursive search..."
) )
@@ -857,12 +787,14 @@ class Neo4JStorage(BaseGraphStorage):
logger.warning( logger.warning(
"Neo4j: inclusive search mode is not supported in recursive query, using exact matching" "Neo4j: inclusive search mode is not supported in recursive query, using exact matching"
) )
return await self._robust_fallback(label, max_depth, min_degree) return await self._robust_fallback(
node_label, max_depth, min_degree
)
return result return result
async def _robust_fallback( async def _robust_fallback(
self, label: str, max_depth: int, min_degree: int = 0 self, node_label: str, max_depth: int, min_degree: int = 0
) -> KnowledgeGraph: ) -> KnowledgeGraph:
""" """
Fallback implementation when APOC plugin is not available or incompatible. Fallback implementation when APOC plugin is not available or incompatible.
@@ -895,12 +827,11 @@ class Neo4JStorage(BaseGraphStorage):
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
query = """ query = """
MATCH (a)-[r]-(b) MATCH (a:base {entity_id: $entity_id})-[r]-(b)
WHERE id(a) = toInteger($node_id)
WITH r, b, id(r) as edge_id, id(b) as target_id WITH r, b, id(r) as edge_id, id(b) as target_id
RETURN r, b, edge_id, target_id RETURN r, b, edge_id, target_id
""" """
results = await session.run(query, {"node_id": node.id}) results = await session.run(query, entity_id=node.id)
# Get all records and release database connection # Get all records and release database connection
records = await results.fetch( records = await results.fetch(
@@ -928,14 +859,16 @@ class Neo4JStorage(BaseGraphStorage):
edge_id = str(record["edge_id"]) edge_id = str(record["edge_id"])
if edge_id not in visited_edges: if edge_id not in visited_edges:
b_node = record["b"] b_node = record["b"]
target_id = str(record["target_id"]) target_id = b_node.get("entity_id")
if b_node.labels: # Only process if target node has labels if target_id: # Only process if target node has entity_id
# Create KnowledgeGraphNode for target # Create KnowledgeGraphNode for target
target_node = KnowledgeGraphNode( target_node = KnowledgeGraphNode(
id=f"{target_id}", id=f"{target_id}",
labels=list(b_node.labels), labels=[
properties=dict(b_node), label for label in b_node.labels if label != "base"
],
properties=dict(b_node.properties),
) )
# Create KnowledgeGraphEdge # Create KnowledgeGraphEdge
@@ -961,11 +894,11 @@ class Neo4JStorage(BaseGraphStorage):
async with self._driver.session( async with self._driver.session(
database=self._DATABASE, default_access_mode="READ" database=self._DATABASE, default_access_mode="READ"
) as session: ) as session:
query = f""" query = """
MATCH (n:`{label}`) MATCH (n:base {entity_id: $entity_id})
RETURN id(n) as node_id, n RETURN id(n) as node_id, n
""" """
node_result = await session.run(query) node_result = await session.run(query, entity_id=node_label)
try: try:
node_record = await node_result.single() node_record = await node_result.single()
if not node_record: if not node_record:
@@ -973,9 +906,11 @@ class Neo4JStorage(BaseGraphStorage):
# Create initial KnowledgeGraphNode # Create initial KnowledgeGraphNode
start_node = KnowledgeGraphNode( start_node = KnowledgeGraphNode(
id=f"{node_record['node_id']}", id=f"{node_record['n'].get('entity_id')}",
labels=list(node_record["n"].labels), labels=[
properties=dict(node_record["n"]), label for label in node_record["n"].labels if label != "base"
],
properties=dict(node_record["n"].properties),
) )
finally: finally:
await node_result.consume() # Ensure results are consumed await node_result.consume() # Ensure results are consumed
@@ -1000,9 +935,8 @@ class Neo4JStorage(BaseGraphStorage):
# Method 2: Query compatible with older versions # Method 2: Query compatible with older versions
query = """ query = """
MATCH (n) MATCH (n)
WITH DISTINCT labels(n) AS node_labels WHERE n.entity_id IS NOT NULL
UNWIND node_labels AS label RETURN DISTINCT n.entity_id AS label
RETURN DISTINCT label
ORDER BY label ORDER BY label
""" """
result = await session.run(query) result = await session.run(query)
@@ -1034,15 +968,14 @@ class Neo4JStorage(BaseGraphStorage):
Args: Args:
node_id: The label of the node to delete node_id: The label of the node to delete
""" """
label = self._ensure_label(node_id)
async def _do_delete(tx: AsyncManagedTransaction): async def _do_delete(tx: AsyncManagedTransaction):
query = f""" query = """
MATCH (n:`{label}`) MATCH (n:base {entity_id: $entity_id})
DETACH DELETE n DETACH DELETE n
""" """
result = await tx.run(query) result = await tx.run(query, entity_id=node_id)
logger.debug(f"Deleted node with label '{label}'") logger.debug(f"Deleted node with label '{node_id}'")
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
try: try:
@@ -1092,16 +1025,16 @@ class Neo4JStorage(BaseGraphStorage):
edges: List of edges to be deleted, each edge is a (source, target) tuple edges: List of edges to be deleted, each edge is a (source, target) tuple
""" """
for source, target in edges: for source, target in edges:
source_label = self._ensure_label(source)
target_label = self._ensure_label(target)
async def _do_delete_edge(tx: AsyncManagedTransaction): async def _do_delete_edge(tx: AsyncManagedTransaction):
query = f""" query = """
MATCH (source:`{source_label}`)-[r]-(target:`{target_label}`) MATCH (source:base {entity_id: $source_entity_id})-[r]-(target:base {entity_id: $target_entity_id})
DELETE r DELETE r
""" """
result = await tx.run(query) result = await tx.run(
logger.debug(f"Deleted edge from '{source_label}' to '{target_label}'") query, source_entity_id=source, target_entity_id=target
)
logger.debug(f"Deleted edge from '{source}' to '{target}'")
await result.consume() # Ensure result is fully consumed await result.consume() # Ensure result is fully consumed
try: try:

124
lightrag/operate.py
View File

@@ -141,18 +141,36 @@ async def _handle_single_entity_extraction(
): ):
if len(record_attributes) < 4 or record_attributes[0] != '"entity"': if len(record_attributes) < 4 or record_attributes[0] != '"entity"':
return None return None
# add this record as a node in the G
# Clean and validate entity name
entity_name = clean_str(record_attributes[1]).strip('"') entity_name = clean_str(record_attributes[1]).strip('"')
if not entity_name.strip(): if not entity_name.strip():
logger.warning(
f"Entity extraction error: empty entity name in: {record_attributes}"
)
return None return None
# Clean and validate entity type
entity_type = clean_str(record_attributes[2]).strip('"') entity_type = clean_str(record_attributes[2]).strip('"')
if not entity_type.strip() or entity_type.startswith('("'):
logger.warning(
f"Entity extraction error: invalid entity type in: {record_attributes}"
)
return None
# Clean and validate description
entity_description = clean_str(record_attributes[3]).strip('"') entity_description = clean_str(record_attributes[3]).strip('"')
entity_source_id = chunk_key if not entity_description.strip():
logger.warning(
f"Entity extraction error: empty description for entity '{entity_name}' of type '{entity_type}'"
)
return None
return dict( return dict(
entity_name=entity_name, entity_name=entity_name,
entity_type=entity_type, entity_type=entity_type,
description=entity_description, description=entity_description,
source_id=entity_source_id, source_id=chunk_key,
metadata={"created_at": time.time()}, metadata={"created_at": time.time()},
) )
@@ -438,47 +456,22 @@ async def extract_entities(
else: else:
return await use_llm_func(input_text) return await use_llm_func(input_text)
async def _process_single_content(chunk_key_dp: tuple[str, TextChunkSchema]): async def _process_extraction_result(result: str, chunk_key: str):
""" "Prpocess a single chunk """Process a single extraction result (either initial or gleaning)
Args: Args:
chunk_key_dp (tuple[str, TextChunkSchema]): result (str): The extraction result to process
("chunck-xxxxxx", {"tokens": int, "content": str, "full_doc_id": str, "chunk_order_index": int}) chunk_key (str): The chunk key for source tracking
Returns:
tuple: (nodes_dict, edges_dict) containing the extracted entities and relationships
""" """
nonlocal processed_chunks maybe_nodes = defaultdict(list)
chunk_key = chunk_key_dp[0] maybe_edges = defaultdict(list)
chunk_dp = chunk_key_dp[1]
content = chunk_dp["content"]
# hint_prompt = entity_extract_prompt.format(**context_base, input_text=content)
hint_prompt = entity_extract_prompt.format(
**context_base, input_text="{input_text}"
).format(**context_base, input_text=content)
final_result = await _user_llm_func_with_cache(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 _user_llm_func_with_cache(
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 _user_llm_func_with_cache(
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( records = split_string_by_multi_markers(
final_result, result,
[context_base["record_delimiter"], context_base["completion_delimiter"]], [context_base["record_delimiter"], context_base["completion_delimiter"]],
) )
maybe_nodes = defaultdict(list)
maybe_edges = defaultdict(list)
for record in records: for record in records:
record = re.search(r"\((.*)\)", record) record = re.search(r"\((.*)\)", record)
if record is None: if record is None:
@@ -487,6 +480,7 @@ async def extract_entities(
record_attributes = split_string_by_multi_markers( record_attributes = split_string_by_multi_markers(
record, [context_base["tuple_delimiter"]] record, [context_base["tuple_delimiter"]]
) )
if_entities = await _handle_single_entity_extraction( if_entities = await _handle_single_entity_extraction(
record_attributes, chunk_key record_attributes, chunk_key
) )
@@ -501,6 +495,62 @@ async def extract_entities(
maybe_edges[(if_relation["src_id"], if_relation["tgt_id"])].append( maybe_edges[(if_relation["src_id"], if_relation["tgt_id"])].append(
if_relation if_relation
) )
return maybe_nodes, maybe_edges
async def _process_single_content(chunk_key_dp: tuple[str, TextChunkSchema]):
"""Process a single chunk
Args:
chunk_key_dp (tuple[str, TextChunkSchema]):
("chunk-xxxxxx", {"tokens": int, "content": str, "full_doc_id": str, "chunk_order_index": int})
"""
nonlocal processed_chunks
chunk_key = chunk_key_dp[0]
chunk_dp = chunk_key_dp[1]
content = chunk_dp["content"]
# Get initial extraction
hint_prompt = entity_extract_prompt.format(
**context_base, input_text="{input_text}"
).format(**context_base, input_text=content)
final_result = await _user_llm_func_with_cache(hint_prompt)
history = pack_user_ass_to_openai_messages(hint_prompt, final_result)
# Process initial extraction
maybe_nodes, maybe_edges = await _process_extraction_result(
final_result, chunk_key
)
# Process additional gleaning results
for now_glean_index in range(entity_extract_max_gleaning):
glean_result = await _user_llm_func_with_cache(
continue_prompt, history_messages=history
)
history += pack_user_ass_to_openai_messages(continue_prompt, glean_result)
# Process gleaning result separately
glean_nodes, glean_edges = await _process_extraction_result(
glean_result, chunk_key
)
# Merge results
for entity_name, entities in glean_nodes.items():
maybe_nodes[entity_name].extend(entities)
for edge_key, edges in glean_edges.items():
maybe_edges[edge_key].extend(edges)
if now_glean_index == entity_extract_max_gleaning - 1:
break
if_loop_result: str = await _user_llm_func_with_cache(
if_loop_prompt, history_messages=history
)
if_loop_result = if_loop_result.strip().strip('"').strip("'").lower()
if if_loop_result != "yes":
break
processed_chunks += 1 processed_chunks += 1
entities_count = len(maybe_nodes) entities_count = len(maybe_nodes)
relations_count = len(maybe_edges) relations_count = len(maybe_edges)