add `kùzu` as KuzuGraphStorage #1451

drahnreb · 2025-04-25T00:26:59Z

Description

Adds a new contender to the ring: KùzuDB! This PR introduces KuzuGraphStorage, a new graph storage backend option, leveraging the high-performance, embeddable Kùzu graph database.

Integrating Kùzu could act as a middle ground between the simplicity of the in-memory NetworkXStorage and the robustness (but higher operational overhead) of server-based solutions like Neo4JStorage or PGGraphStorage.

It provides:

High Performance: speed using columnar storage, vectorized execution, and modern join algorithms, it should accelerate graph traversals.
Embeddable Simplicity: Runs in-process with the LightRAG application, requiring only a library dependency and a database file path (or using :memory:). This reduces setup complexity and infrastructure management compared to client-server databases.
Persistence: Offers reliable on-disk persistence, unlike the default in-memory storage.
Python & AI Ecosystem Integration: Provides a smooth Python API and is designed with AI/ML integrations in mind (see notes below).
Permissive Licensing: Uses the MIT license, ensuring compatibility with LightRAG.
Cypher Support: Leverages the popular Cypher query language.

This addition enables more complex knowledge graph structures and sophisticated graph-based retrieval strategies without significant operational hurdles.

Related Issues

None

Changes Made

Added kuzu_impl.py containing the KuzuGraphStorage class implementation, adhering to the BaseGraphStorage interface.
Included KuzuClientManager for managing singleton database and connection instances.
Added KùzuDB dependency (kuzu>=0.9.0) check and installation prompt via pipmaster.
default configuration (config.ini.example) for optional kùzu option.

Checklist

Changes tested locally (Please confirm specific tests run)
Code reviewed
Documentation updated
successfully passed all tests for KuzuGraphStorage (see picture attached)

Additional Notes

Kùzu Strengths for Advanced LightRAG:

As discussed, KùzuDB's strengths align well with potential future enhancements for graph-based retrieval in LightRAG:

PyTorch Geometric (PyG) Integration: Kùzu offers excellent, direct integration with PyG via get_torch_geometric_remote_backend(). This facilitates using graph embeddings generated by GNNs directly within Kùzu, potentially enabling more semantically rich graph traversals for retrieval (e.g., in mix mode).
Advanced Graph Algorithms: While Kùzu focuses on core storage and query performance, it provides strong interoperability. Complex graph algorithms (like centrality metrics, community detection) can be efficiently computed and these metrics could then inform retrieval strategies or KG analysis.

python test_graph_storage.py

    ╔══════════════════════════════════════════════════════════════╗
    ║                  通用图存储测试程序                          ║
    ╚══════════════════════════════════════════════════════════════╝


当前配置的图存储类型: KuzuGraphStorage
支持的图存储类型: NetworkXStorage, Neo4JStorage, PGGraphStorage, KuzuGraphStorage
config.ini not found at /Users/drahnreb/GitHub/LightRAG/lightrag/kg/config.ini, using environment variables or defaults.

请选择测试类型:
1. 基本测试 (节点和边的插入、读取)
2. 高级测试 (度数、标签、知识图谱、删除操作等)
3. 批量操作测试 (批量获取节点、边属性和度数等)
4. 无向图特性测试 (验证存储的无向图特性)
5. 特殊字符测试 (验证单引号、双引号和反斜杠等特殊字符)
6. 全部测试

请输入选项 (1/2/3/4/5/6): 6

执行测试前清理数据...
Dropping all data from Kuzu graph (namespace: test_graph)...
数据清理完成


=== 开始基本测试 ===
插入节点1: 人工智能
插入节点2: 机器学习
插入边: 人工智能 -> 机器学习
读取节点属性: 人工智能
成功读取节点属性: 人工智能
节点描述: 人工智能是计算机科学的一个分支，它企图了解智能的实质，并生产出一种新的能以人类智能相似的方式做出反应的智能机器。
节点类型: 技术领域
节点关键词: AI,机器学习,深度学习
读取边属性: 人工智能 -> 机器学习
成功读取边属性: 人工智能 -> 机器学习
边关系: 包含
边描述: 人工智能领域包含机器学习这个子领域
边权重: 1.0
读取反向边属性: 机器学习 -> 人工智能
成功读取反向边属性: 机器学习 -> 人工智能
反向边关系: 包含
反向边描述: 人工智能领域包含机器学习这个子领域
反向边权重: 1.0
无向图特性验证成功：正向和反向边属性一致
基本测试完成，数据已保留在数据库中

=== 开始高级测试 ===
插入节点1: 人工智能
插入节点2: 机器学习
插入节点3: 深度学习
插入边1: 人工智能 -> 机器学习
插入边2: 机器学习 -> 深度学习
== 测试 node_degree: 人工智能
节点 人工智能 的度数: 1
== 测试所有节点的度数
节点 机器学习 的度数: 2
节点 深度学习 的度数: 1
== 测试 edge_degree: 人工智能 -> 机器学习
边 人工智能 -> 机器学习 的度数: 3
== 测试反向边的度数: 机器学习 -> 人工智能
反向边 机器学习 -> 人工智能 的度数: 3
无向图特性验证成功：正向边和反向边的度数一致
== 测试 get_node_edges: 机器学习
节点 机器学习 的所有边: [('机器学习', '人工智能'), ('机器学习', '深度学习')]
== 验证节点边的无向图特性
无向图特性验证成功：节点 机器学习 的边列表包含所有相关的边
== 测试 get_all_labels
所有标签: ['人工智能', '机器学习', '深度学习']
== 测试 get_knowledge_graph
知识图谱节点数: 3
知识图谱边数: 2
== 测试 delete_node: 深度学习
删除后查询节点属性 深度学习: None
== 测试 remove_edges: 机器学习 -> 深度学习
删除后查询边属性 机器学习 -> 深度学习: None
== 验证删除边的无向图特性: 深度学习 -> 机器学习
删除后查询反向边属性 深度学习 -> 机器学习: None
无向图特性验证成功：删除一个方向的边后，反向边也被删除
== 测试 remove_nodes: [机器学习, 深度学习]
删除后查询节点属性 机器学习: None
删除后查询节点属性 深度学习: None

高级测试完成

=== 开始批量操作测试 ===
插入节点1: 人工智能
插入节点2: 机器学习
插入节点3: 深度学习
插入节点4: 自然语言处理
插入节点5: 计算机视觉
插入边1: 人工智能 -> 机器学习
插入边2: 机器学习 -> 深度学习
插入边3: 人工智能 -> 自然语言处理
插入边4: 人工智能 -> 计算机视觉
插入边5: 深度学习 -> 自然语言处理
插入边6: 深度学习 -> 计算机视觉
== 测试 get_nodes_batch
批量获取节点属性结果: dict_keys(['人工智能', '机器学习', '深度学习'])
== 测试 node_degrees_batch
批量获取节点度数结果: {'人工智能': 3, '机器学习': 2, '深度学习': 3}
== 测试 edge_degrees_batch
批量获取边度数结果: {('人工智能', '机器学习'): 5, ('机器学习', '深度学习'): 5, ('深度学习', '自然语言处理'): 5}
== 测试 get_edges_batch
批量获取边属性结果: dict_keys([('人工智能', '机器学习'), ('机器学习', '深度学习'), ('深度学习', '自然语言处理')])
== 测试反向边的批量获取
批量获取反向边属性结果: dict_keys([('机器学习', '人工智能'), ('深度学习', '机器学习'), ('自然语言处理', '深度学习')])
无向图特性验证成功：批量获取的正向和反向边属性一致
== 测试 get_nodes_edges_batch
批量获取节点边结果: dict_keys(['人工智能', '深度学习'])
== 验证批量获取节点边的无向图特性
节点 人工智能 的出边: [('人工智能', '计算机视觉'), ('人工智能', '自然语言处理'), ('人工智能', '机器学习')]
节点 人工智能 的入边: []
节点 深度学习 的出边: [('深度学习', '计算机视觉'), ('深度学习', '自然语言处理'), ('深度学习', '机器学习')]
节点 深度学习 的入边: []
无向图特性验证成功：批量获取的节点边包含所有相关的边（无论方向）

批量操作测试完成

=== 开始无向图特性测试 ===
插入节点1: 计算机科学
插入节点2: 数据结构
插入节点3: 算法

== 测试插入边后的无向图特性
插入边1: 计算机科学 -> 数据结构
正向边属性: {'relationship': '包含', 'weight': 1.0, 'description': '计算机科学包含数据结构这个概念'}
反向边属性: {'relationship': '包含', 'weight': 1.0, 'description': '计算机科学包含数据结构这个概念'}
无向图特性验证成功：正向和反向边属性一致

== 测试边的度数的无向图特性
插入边2: 计算机科学 -> 算法
正向边 计算机科学 -> 数据结构 的度数: 3
反向边 数据结构 -> 计算机科学 的度数: 3
无向图特性验证成功：正向和反向边的度数一致

== 测试删除边的无向图特性
删除边: 计算机科学 -> 数据结构
删除后查询正向边属性 计算机科学 -> 数据结构: None
删除后查询反向边属性 数据结构 -> 计算机科学: None
无向图特性验证成功：删除一个方向的边后，反向边也被删除

== 测试批量操作中的无向图特性
批量获取正向边属性结果: dict_keys([('计算机科学', '数据结构'), ('计算机科学', '算法')])
批量获取反向边属性结果: dict_keys([('数据结构', '计算机科学'), ('算法', '计算机科学')])
无向图特性验证成功：批量获取的正向和反向边属性一致

== 测试批量获取节点边的无向图特性
批量获取节点边结果: dict_keys(['计算机科学', '数据结构'])
无向图特性验证成功：批量获取的节点边包含所有相关的边（无论方向）

无向图特性测试完成

=== 开始特殊字符测试 ===
插入包含特殊字符的节点1: 包含'单引号'的节点
插入包含特殊字符的节点2: 包含"双引号"的节点
插入包含特殊字符的节点3: 包含\反斜杠\的节点
插入包含特殊字符的边: 包含'单引号'的节点 -> 包含"双引号"的节点
插入包含复杂特殊字符的边: 包含"双引号"的节点 -> 包含\反斜杠\的节点

== 验证节点特殊字符
成功读取节点: 包含'单引号'的节点
节点描述: 这个描述包含'单引号'、"双引号"和\反斜杠
节点 包含'单引号'的节点 特殊字符验证成功
成功读取节点: 包含"双引号"的节点
节点描述: 这个描述同时包含'单引号'和"双引号"以及\反斜杠\路径
节点 包含"双引号"的节点 特殊字符验证成功
成功读取节点: 包含\反斜杠\的节点
节点描述: 这个描述包含Windows路径C:\Program Files\和转义字符\n\t
节点 包含\反斜杠\的节点 特殊字符验证成功

== 验证边特殊字符
成功读取边: 包含'单引号'的节点 -> 包含"双引号"的节点
边关系: 特殊'关系'
边描述: 这个边描述包含'单引号'、"双引号"和\反斜杠
边 包含'单引号'的节点 -> 包含"双引号"的节点 特殊字符验证成功
成功读取边: 包含"双引号"的节点 -> 包含\反斜杠\的节点
边关系: 复杂"关系"\类型
边描述: 包含SQL注入尝试: SELECT * FROM users WHERE name='admin'--
边 包含"双引号"的节点 -> 包含\反斜杠\的节点 特殊字符验证成功

特殊字符测试完成，数据已保留在数据库中

存储连接已关闭

drahnreb added 6 commits April 25, 2025 00:55

first kuzu implementation

8a3d475

test 1 and 5 working. 2,3,4 seg faults. needs fix

bb75578

test 1,2,4,5 working. 2 needs fix

a866920

all tests working.

52d44c4

cleanup. add config and docs.

3b469ce

Merge branch 'main' into drahnreb/add-kuzu-graph-db

b473c67

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

add `kùzu` as KuzuGraphStorage #1451

add `kùzu` as KuzuGraphStorage #1451

drahnreb commented Apr 25, 2025 •

edited

Loading

add kùzu as KuzuGraphStorage #1451

Are you sure you want to change the base?

add kùzu as KuzuGraphStorage #1451

Conversation

drahnreb commented Apr 25, 2025 • edited Loading

Description

Related Issues

Changes Made

Checklist

Additional Notes

add `kùzu` as KuzuGraphStorage #1451

add `kùzu` as KuzuGraphStorage #1451

drahnreb commented Apr 25, 2025 •

edited

Loading