写点什么

使用 TiDB Vector 构建 LightRAG 知识库

  • 2024-12-27
    北京
  • 本文字数:3848 字

    阅读完需:约 13 分钟

作者: 魔人逗逗原文来源:https://tidb.net/blog/a3202a6b

引言

GraphRAG 是大家都耳熟能详的一种 RAG 算法,LightRAG 可以视为是 GraphRAG 青春版或轻量版,重点在于更轻量、提升效率、持久化等,例如:


  • LightRAG 去掉了 GraphRAG 中 Community 的概念,只聚焦于 EntityRelationGraph 本身,大大简化了整个知识库构建的流程

  • LightRAG 提供了一些持久化机制,将存储分为了 kv 存储vector 存储graph 存储,并支持了 oraclechromaneo4j 等不同类型的存储形式

  • LightRAG 还简化了查询效率,查询的时候会做向量检索;然后根据图检索相关的Entity 或者Relation ,以此来构建整个上下文;最后只需要调用一次大模型生成回复即可


本篇文章主要是想介绍 TiDB Vector 是怎么和 LightRAG 结合的。

LightRAG 流程

我们先回顾一下 LightRAG 的整体流程,主要是从索引构建查询两个角度来讲。 想看完整的论文可以看这里,完整的源码在 HKUDS/LightRAG



如果你对 GraphRAG 不太理解,我们先来简单解释一下 ChunkEntityRelation


  • Chunk 其实就是对完整文档的切片,LightRAG 主要是支持了 token 切片,当然这远远无法满足所有文档,很多 markdown 文档甚至也可以做一些特殊的结构化的切片

  • Entity 可以理解为 Chunk 中的关键字 keywords 或者一些关键的信息、高频词汇等等,Entity 也会分为不同的类型。当然这里的 Entity 实例是让大模型帮助总结的,具体的提示词可以在 lightrag/prompt.py 代码中找到

  • 最后就是 Relation ,前面说到了 Entity ,那么构建基于 Entity 的图就必须要各个节点直接的关系,这就是 Relation。同理,Relation 也是让大模型总结提取的。


知道了这些概念之后让我们看看具体的流程。


索引构建的流程大概可以分为以下几步:


  1. 解析文档

  2. 对文档进行 token 切片 (Chunk) 并入库 vector 存储实例

  3. 解析 Chunk,提取其中的 EntityRelation

  4. 合并 EntityRelation 并存入 graph 存储实例

  5. EntityRelation 存入 vector 存储实例

  6. 将原始文件和切片存入 kv 存储引擎


检索的方式主要分为了两类:


  • 简单检索 (naive):和原始的 RAG 检索类似,直接对切片 (chunk) 进行向量检索,并调用大模型生成回复

  • 图检索:LightRAG 在这里将图检索分为了三类:

  • local:侧重于对 Entity 进行语义检索,其次通过 Graph 查询并合并和与检索结果相关的 EntityRelation,然后根据等级 + 权重进行排序,最后将EntityRelationChunk 拼接为一个完整的上下文并请求大模型生成结果

  • global:侧重于对 Relation 进行语义检索,其次通过 Graph 查询并合并和与检索结果相关的 EntityRelation ,然后根据等级 + 权重进行排序,最后将EntityRelationChunk 拼接为一个完整的上下文并请求大模型生成结果

  • hybrid:混合检索则是 local + global 的方式,分别对 EntityRelation 进行语义检索并合并,最后走生成一个合并后的上下文

表结构

我这里主要分为了 4 张表,结合上面提到的 kv 存储vector 存储graph 存储,下面依次介绍:


LIGHTRAG_DOC_FULL 表主要是用来存原始文档的。在 LightRAG 中 kv 存储 中被使用。


CREATE TABLE LIGHTRAG_DOC_FULL (    `id` BIGINT PRIMARY KEY AUTO_RANDOM,    `doc_id` VARCHAR(256) NOT NULL,    `workspace` varchar(1024),    `content` LONGTEXT,    `meta` JSON,    `createtime` TIMESTAMP DEFAULT CURRENT_TIMESTAMP,    `updatetime` TIMESTAMP DEFAULT NULL,    UNIQUE KEY (`doc_id`));
复制代码


LIGHTRAG_DOC_CHUNKS 表主要是用来存储原始文档的切片,即 Chunk


在 LightRAG 中既会被 kv 存储来存切片,也会被 vector 存储来存切片的向量数据。


CREATE TABLE LIGHTRAG_DOC_CHUNKS (    `id` BIGINT PRIMARY KEY AUTO_RANDOM,    `chunk_id` VARCHAR(256) NOT NULL,    `full_doc_id` VARCHAR(256) NOT NULL,    `workspace` varchar(1024),    `chunk_order_index` INT,    `tokens` INT,    `content` LONGTEXT,    `content_vector` VECTOR,    `createtime` DATETIME DEFAULT CURRENT_TIMESTAMP,    `updatetime` DATETIME DEFAULT NULL,    UNIQUE KEY (`chunk_id`));
复制代码


LIGHTRAG_GRAPH_NODES 表主要是用来存储解析后的 Entity


在 LightRAG 中既会被 vector 存储来存 Entity 中的内容和向量,也会被 graph 存储来存Entity 在 Graph 中的信息如 entity_typedescription 额外字段,作为 Entity 在整个图中的属性。


还有一点是,LightRAG 在存储 Entity 前已经进行了去重和合并,以保证构建图的时候不会有完全一样的节点。


CREATE TABLE LIGHTRAG_GRAPH_NODES (    `id` BIGINT PRIMARY KEY AUTO_RANDOM,    `entity_id` VARCHAR(256),    `workspace` varchar(1024),    `name` VARCHAR(2048),    `entity_type` VARCHAR(1024),    `description` LONGTEXT,    `source_chunk_id` VARCHAR(256),    `content` LONGTEXT,    `content_vector` VECTOR,    `createtime` DATETIME DEFAULT CURRENT_TIMESTAMP,    `updatetime` DATETIME DEFAULT NULL,    KEY (`entity_id`));
复制代码


LIGHTRAG_GRAPH_EDGES 表主要是用来存储 Relations


同样的,在 LightRAG 中既会被 vector 存储来存 Relations 内容及其向量,也会被 graph 存储来存 Relations 在 Graph 中的信息如 weightkeywordsdescription 等。


这里同样需要提醒一点,GraphRAGLightRAG 最终构建的都是一个有权重且无方向的图。所以对下面这张表来说 source_nametarget_name 的组合对只会出现一次,这部分的去重和合并也是 LightRAG 在代码里做的。


至于为什么选择有权重且无方向的图?个人猜想 RAG 检索主要是关注的是图中节点之间的连接关系(语义之间的相关性)和权重(关联的强弱),而不太在意节点在图中的前后顺序。这样一来,只需要关注 relation 关系本身和 weight 权重即可,不需要关心谁是 source 谁是 target;当然,无向也降低了整个结构的复杂性。


CREATE TABLE LIGHTRAG_GRAPH_EDGES (    `id` BIGINT PRIMARY KEY AUTO_RANDOM,    `relation_id` VARCHAR(256),    `workspace` varchar(1024),    `source_name` VARCHAR(2048),    `target_name` VARCHAR(2048),    `weight` DECIMAL,    `keywords` TEXT,    `description` LONGTEXT,    `source_chunk_id` varchar(256),    `content` LONGTEXT,    `content_vector` VECTOR,    `createtime` DATETIME DEFAULT CURRENT_TIMESTAMP,    `updatetime` DATETIME DEFAULT NULL,    KEY (`relation_id`));
复制代码

SQL 和使用

前面说到只要分为了索引构建检索两个部分。索引构建部分比较简单,就是 CRUD 没什么特殊的,这里我们主要说一下检索相关的 SQL。


查询阶段会先进行向量检索,然后根据 Graph 图去查询相关的 EntityRelation


我们先看看向量检索的部分,以下只是 SQL 示例,可以根据具体场景进行优化改写


# Search for Entity/Node VectorSELECT n.name as entity_nameFROM (        SELECT entity_id as id,            name,            VEC_COSINE_DISTANCE(content_vector, :embedding_string) as distance        FROM LIGHTRAG_GRAPH_NODES        WHERE workspace = :workspace    ) nWHERE n.distance > :better_than_thresholdORDER BY n.distance DESCLIMIT :top_k;
# Search for Relation/Edge VectorSELECT e.source_name as src_id, e.target_name as tgt_idFROM ( SELECT source_name, target_name, VEC_COSINE_DISTANCE(content_vector, :embedding_string) as distance FROM LIGHTRAG_GRAPH_EDGES WHERE workspace = :workspace ) eWHERE e.distance > :better_than_thresholdORDER BY e.distance DESCLIMIT :top_k
复制代码


其次是查询向量结果的明细信息,并获取 EntityRelation 的度数 degree


  • Entitydegree 指的是与这个节点相连的边的数量

  • Relationdegree 指的是构成 Relation 的两个 Entitydegree 之和


这里只需要一个 SQL 语句,需要注意这里只适用于无向图


SELECT COUNT(id) AS cntFROM LIGHTRAG_GRAPH_EDGESWHERE workspace = :workspace    AND :name IN (source_name, target_name)
复制代码


到这里,通过向量检索出来的 Entity 或者 Relation 的内容和度数已经拿出来了,最后就是根据这写结果去查询相关的 Relation 或者 Entity


  • 根据向量检索后的 Entity 查询这些节点相关的边,即 Relation

  • 根据向量检索后的 Relation 找出对应的 src_idtgt_id ,查询这些节点,即 Entity


还有最后一步就是查询这些 EntityRelation 相关的 Chunk ,这里都是一些比较简单的 SQL 这里就不再赘述,感兴趣的可以去下载源码运行。

总结

显而易见,从 LightRAG 的角度来看 TiDB 相比于很多只做关系型或者只做向量的存储引擎确实优势很大。


对于 Graph 图这部分搜到了 TiGraph 和相关的运算符,但是具体的使用示例比较少就暂时没有深入了解。 ——TiGraph: 8,700x Computing Performance Achieved by Combining Graphs + the RDBMS Syntax


从目前我对 TiDB 和 TiDB Vector 的了解,在无向图的情况下,基本一个 TiDB 便足矣来完成 GraphRAG 及其衍生的大部分场景的持久化。除此之外,能够与之一战的,大概也就只有商业化完善的 Oracle 或是社区生态丰富的 PostgreSQL 了。


最后推荐一下 TiDB 发布的基于 GraphRAG 的开源代码 pingcap/autoflow ,最近太忙了,等之后有时间了深入了解了解 x


发布于: 刚刚阅读数: 4
用户头像

TiDB 社区官网:https://tidb.net/ 2021-12-15 加入

TiDB 社区干货传送门是由 TiDB 社区中布道师组委会自发组织的 TiDB 社区优质内容对外宣布的栏目,旨在加深 TiDBer 之间的交流和学习。一起构建有爱、互助、共创共建的 TiDB 社区 https://tidb.net/

评论

发布
暂无评论
使用 TiDB Vector 构建 LightRAG 知识库_实践案例_TiDB 社区干货传送门_InfoQ写作社区