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

use KnowledgeGraph typed dict for graph API response

Browse Source
This commit is contained in:
ArnoChen
2025-02-13 17:32:51 +08:00
parent e4562c761c
commit c674905a98
5 changed files with 54 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
lightrag/kg/neo4j_impl.py
View File

@@ -25,6 +25,7 @@ from tenacity import (
from ..utils import logger
from ..base import BaseGraphStorage
from ..types import KnowledgeGraph, KnowledgeGraphNode, KnowledgeGraphEdge
@dataclass
@@ -44,7 +45,8 @@ class Neo4JStorage(BaseGraphStorage):
URI = os.environ["NEO4J_URI"]
USERNAME = os.environ["NEO4J_USERNAME"]
PASSWORD = os.environ["NEO4J_PASSWORD"]
MAX_CONNECTION_POOL_SIZE = os.environ.get("NEO4J_MAX_CONNECTION_POOL_SIZE", 800)
MAX_CONNECTION_POOL_SIZE = os.environ.get(
"NEO4J_MAX_CONNECTION_POOL_SIZE", 800)
DATABASE = os.environ.get(
"NEO4J_DATABASE", re.sub(r"[^a-zA-Z0-9-]", "-", namespace)
)
@@ -74,19 +76,22 @@ class Neo4JStorage(BaseGraphStorage):
)
raise e
except neo4jExceptions.AuthError as e:
logger.error(f"Authentication failed for {database} at {URI}")
logger.error(
f"Authentication failed for {database} at {URI}")
raise e
except neo4jExceptions.ClientError as e:
if e.code == "Neo.ClientError.Database.DatabaseNotFound":
logger.info(
f"{database} at {URI} not found. Try to create specified database.".capitalize()
f"{database} at {URI} not found. Try to create specified database.".capitalize(
)
)
try:
with _sync_driver.session() as session:
session.run(
f"CREATE DATABASE `{database}` IF NOT EXISTS"
)
logger.info(f"{database} at {URI} created".capitalize())
logger.info(
f"{database} at {URI} created".capitalize())
connected = True
except (
neo4jExceptions.ClientError,
@@ -103,7 +108,8 @@ class Neo4JStorage(BaseGraphStorage):
"This Neo4j instance does not support creating databases. Try to use Neo4j Desktop/Enterprise version or DozerDB instead. Fallback to use the default database."
)
if database is None:
logger.error(f"Failed to create {database} at {URI}")
logger.error(
f"Failed to create {database} at {URI}")
raise e
if connected:
@@ -365,7 +371,7 @@ class Neo4JStorage(BaseGraphStorage):
async def get_knowledge_graph(
self, node_label: str, max_depth: int = 5
) -> Dict[str, List[Dict]]:
) -> KnowledgeGraph:
"""
Get complete connected subgraph for specified node (including the starting node itself)
@@ -376,7 +382,7 @@ class Neo4JStorage(BaseGraphStorage):
4. Add depth control
"""
label = node_label.strip('"')
result = {"nodes": [], "edges": []}
result = KnowledgeGraph()
seen_nodes = set()
seen_edges = set()
@@ -395,7 +401,8 @@ class Neo4JStorage(BaseGraphStorage):
validate_query = f"MATCH (n:`{label}`) RETURN n LIMIT 1"
validate_result = await session.run(validate_query)
if not await validate_result.single():
logger.warning(f"Starting node {label} does not exist!")
logger.warning(
f"Starting node {label} does not exist!")
return result
# Optimized query (including direction handling and self-loops)
@@ -420,11 +427,11 @@ class Neo4JStorage(BaseGraphStorage):
# Use node ID + label combination as unique identifier
node_id = node.id
if node_id not in seen_nodes:
node_data = {}
node_data["labels"] = list(node.labels) # Keep all labels
node_data["id"] = f"{node_id}"
node_data["properties"] = dict(node)
result["nodes"].append(node_data)
result.nodes.append(KnowledgeGraphNode(
id=f"{node_id}",
labels=list(node.labels),
properties=dict(node),
))
seen_nodes.add(node_id)
# Handle relationships (including direction information)
@@ -433,21 +440,17 @@ class Neo4JStorage(BaseGraphStorage):
if edge_id not in seen_edges:
start = rel.start_node
end = rel.end_node
edge_data = {}
edge_data.update(
{
"source": f"{start.id}",
"target": f"{end.id}",
"type": rel.type,
"id": f"{edge_id}",
"properties": dict(rel),
}
)
result["edges"].append(edge_data)
result.edges.append(KnowledgeGraphEdge(
id=f"{edge_id}",
type=rel.type,
source=f"{start.id}",
target=f"{end.id}",
properties=dict(rel),
))
seen_edges.add(edge_id)
logger.info(
f"Subgraph query successful | Node count: {len(result['nodes'])} | Edge count: {len(result['edges'])}"
f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}"
)
except neo4jExceptions.ClientError as e: