基于LlamaIndex实现Rag向量数据存储与召回

470 阅读2分钟

基于LlamaIndex实现Rag向量数据存储与召回

1.什么是LlamaIndex?

LlamaIndex一个基于LLM(大语言模型)的AI应用开发框架,多适用于RAG内容检索场景。 提供了灵活的数据的抽象层与工具,以便更容易地摄取、结构化的处理与访问外部数据,从而更高效地将这些数据注入LLM 中,以实现更准确的文本生成。

2.数据存储&召回流程

AI智能体平台开发文档.png

3.自定义LlamaIndex向量化模型

export class CustomEmbeddings extends BaseEmbedding {
  constructor() {
    super();
  }
  async getTextEmbedding(text) {
    // 请求Embedding模型,并返回向量数据.....
    return [];
  }
  async getQueryEmbedding(text) {
    // 如果查询嵌入与文档嵌入相同,可以直接调用 getTextEmbedding
    return this.getTextEmbedding(text.text);
  }
}

4.文本数据处理方案

4.1数据存储过程
import { Document, SentenceSplitter } from 'llamaindex';

// 创建文档
const document = new Document({
  text: documentString,
});

// 创建文本分割器
const textSplitter = new SentenceSplitter({
  chunkSize: 1000,
  chunkOverlap: 200,
});

// 分隔的文档片段
const splitDocs = await textSplitter.splitText(document.text);

// 文档片段向量化
const docEmbeddings = await Promise.all(
  splitDocs.map(async docString => {
    // 自定义Embedding方法
    const embedding = await getTextEmbedding(docString);
    return {
      metaData: docString,
      vectorData: embedding,
    };
  })
);

// 向量化数据存储
await Promise.all(
  docEmbeddings.map(async ({ metaData, vectorData }) => {
    return VectorDB.saveDocumentData({
      metaData,
      vectorData,
    });
  })
);

4.2数据召回过程
import {
  Document,
  VectorStoreIndex,
  serviceContextFromDefaults,
  BaseEmbedding,
} from 'llamaindex';
import { CustomEmbeddings } from "./CustomEmbeddings"

// 定义问题
const question = '这是一个问题';
// 自定义Embedding
const customEmbeddings = new CustomEmbeddings();
// 问题向量化
const embedding = await customEmbeddings.getTextEmbedding(question);
// 向量化检索,得到最相似的5条内容(阿里云 PostgreSQL)
const queryResultList = await DocumentsVectorTableModel.query(
  `
      SELECT
        id,
        meta_data,
        cosine_similarity(vector_data, ARRAY[${embedding.join(
          ','
        )}]) AS similarity
      FROM
        "documents_vector_table"
      ORDER BY
        similarity DESC
      LIMIT 5 
    `
);
// 定制Embedding服务
const serviceContext = serviceContextFromDefaults({
  embedModel: customEmbeddings,
});
// 基于文档构建查询索引
const documents = queryResultList.map(
  queryString =>
    new Document({
      text: queryString,
    })
);
const index = await VectorStoreIndex.fromDocuments(documents, {
  serviceContext,
});
// 查询结果
const data = await index.asQueryEngine().retrieve(question);

5.文件类型处理方案

import * as path from 'path';
import {
  VectorStoreIndex,
  serviceContextFromDefaults,
  SimpleDirectoryReader,
} from 'llamaindex';

const reader = new SimpleDirectoryReader();

// 指定文件的路径,加载该文件下的所有文档
const filePath = path.resolve(
  'filePath'
);

// 使用 MarkdownReader 加载数据
const documents = await reader.loadData(filePath);

// 定制Embedding服务
const customEmbeddings = new CustomEmbeddings();
const serviceContext = serviceContextFromDefaults({
  embedModel: customEmbeddings,
});

// 创建LlamaIndex查询索引
const index = await VectorStoreIndex.fromDocuments(documents, {
  serviceContext,
});

// 查询结果
const data = await index.asQueryEngine().retrieve('这是一个问题');