在nodejs中使用ElasticSearch(三)通过ES语义检索,实现RAG

RAG(Retrieval-Augmented Generation)是一种结合了信息检索和生成模型的技术,旨在提高生成模型的知识获取和生成能力。它通过在生成的过程中引入外部知识库或文档(如数据库、搜索引擎或文档存储),帮助生成更为准确和丰富的答案。 RAG 在自然语言处理(NLP)领域,特别是在对话生成、问答系统和文本摘要等任务中,具有非常重要的应用。它的核心思想是,生成模型不仅依赖于模型内部的知识,还可以在生成过程中动态地检索外部信息,从而增强生成内容的准确性和信息丰富度。

示例

下面这段代码实现了一个完整的工作流,涉及到以下几个关键步骤:

1. 创建 Elasticsearch 索引

  • 首先,代码检查是否已经存在名为 my_index4 的 Elasticsearch 索引。如果存在,则先删除该索引,确保创建的是新的索引。
  • 接着,创建一个新的索引 my_index4,并定义其映射(mappings):
    • id 字段是文本类型。
    • content 字段是 text 类型,并指定了 ik_max_word 分词器,用于中文分词。
    • content_vectordense_vector 类型,表示用于存储文档的嵌入向量,并使用余弦相似度进行向量检索。

2. 加载和处理 PDF 文件

  • 使用 pdfjsLib 库加载本地的 PDF 文件(t.pdf),并从中提取文本内容。文本按页面分批读取,直到所有页面的内容都被提取出来。
  • 提取的文本内容会经过 sentence-splitter 库进行分句处理,将文档拆分成句子节点,便于后续的处理。

3. 嵌入向量化

  • 为了让文本内容可以进行语义检索,代码使用了 OpenAI 的 Embeddings API(调用 openai.embeddings.create)来生成文本的嵌入向量(embedding)。
  • 该嵌入向量将被存储在 Elasticsearch 中,以便后续进行基于向量的检索。
  • 将每个句子的文本与对应的嵌入向量一同批量插入 Elasticsearch 索引(my_index4)中。

4. 批量插入数据到 Elasticsearch

  • 为每一个分句生成一个 id(使用 UUID)以及 content(文本内容)和 content_vector(嵌入向量)。这些数据通过 es.bulk 方法批量插入到 Elasticsearch 中。
  • bulk 方法使用批量操作提高插入效率。

5. 查询和检索

  • 用户的查询被定义为 query(例子中是 'how many parameters does llama 2 have?')。
  • 使用混合检索方式进行查询,结合了 文本搜索match)和 向量检索knn):
    • 文本搜索:查找与查询文本最相似的文档。
    • 向量检索 :将查询文本转换为嵌入向量,并通过 Elasticsearch 的 knn 查询来查找与查询向量最相似的文档。
    • num_candidates: 100 指定了检索时考虑的候选文档数量,以提高准确性和性能平衡。
  • Elasticsearch 会返回匹配度最高的前 5 条文档(size: 5)。

6. 生成回答

  • 基于 Elasticsearch 检索到的文档内容,构建一个提示模板 (promptTemplate),并向 OpenAI 提交请求,生成回答:
    • 提示模板包括用户的查询和从 Elasticsearch 中检索到的相关文档信息。生成的回答应该是基于已知信息的。
    • 如果检索结果中没有足够的信息来回答问题,生成的回答将是"我无法回答您的问题"。

7. 输出回答

  • 通过 get_completion 函数,OpenAI 生成最终的回答,并打印在控制台中。
TypeScript 复制代码
import { Client } from 'npm:@elastic/elasticsearch';
import * as pdfjsLib from 'npm:pdfjs-dist';
import { split, TxtSentenceNodeChildren } from "npm:sentence-splitter";
import fs from "node:fs";
import process from "node:process";
import OpenAI from "npm:openai";
import { v4 } from "npm:uuid";
import 'npm:dotenv/config';


const es = new Client({ node: 'http://localhost:9200' });

const exists = await es.indices.exists({ index: 'my_index4' });

if (exists) {
  await es.indices.delete({ index: 'my_index4' });
}

await es.indices.create({
  index: 'my_index4',
  aliases: {
    'mi4': {},
  },
  mappings: {
    properties: {
      id: { type: 'text' },
      content: {
        type: 'text',
        analyzer: "ik_max_word",
        fields: { keyword: { type: 'keyword' } },
      },
      content_vector: {
        type: "dense_vector",
        dims: 1024,
        index: true,  // 启用对 content_vector 字段的索引功能
        similarity: "cosine", // 使用余弦相似度来评估不同向量之间的相似性
      },
    }
  },
  settings: {
    number_of_replicas: 1,
    number_of_shards: 1,
  },
});

// 加载PDF文件
const uint8Array = new Uint8Array(fs.readFileSync('t.pdf'));
const pdf = await pdfjsLib.getDocument(uint8Array).promise;

let article = '';
let finish = false;
let page = 1;
do {
  try {
    const res = await pdf.getPage(page++);
    const textContent = await res.getTextContent();
    const str = textContent.items.reduce((sum, item) => {
      if (!('str' in item)) {
        return sum;
      } else {
        return sum + item.str;
      }
    }, '');
    article += str;
  } catch {
    finish = true;
  }
} while (!finish);


let PDFContent = split(article);

const openai = new OpenAI(
  {
    // apiKey: process.env.DASHSCOPE_API_KEY,
    apiKey: process.env.OPENAI_API_KEY,
    baseURL: process.env.OPENAI_BASE_URL
  }
);

async function getEmbedding(text: string) {
  const response = await openai.embeddings.create({
    model: String(process.env.OPEN_MODEL_EMBEDDING),  // 使用嵌入模型
    input: [text],
  });
  return response.data;  // 返回向量
}

PDFContent = PDFContent.flatMap((doc) => {
  if ('children' in doc && Array.isArray(doc.children)) {
    const tmp = doc.children.filter((text: { type?: string }) => text?.type === 'Str');
    // console.log(tmp);
    // const embedding = await getEmbedding(doc.raw);
    return tmp;
  } else {
    return null;
  }
}).filter(item => item) as TxtSentenceNodeChildren[];

const operations = (await Promise.all(
  PDFContent.map(async (doc) => {
    try {
      const content = (doc as { value: string }).value;
      const content_vector = (await getEmbedding(content))[0].embedding;
      // console.log(content);
      return [
        { index: { _index: "my_index4" } },
        {
          id: v4(),
          content,
          content_vector
        },
      ];
    } catch {
      return false;
    }
  })
)).filter(item => item).flat(9);

// 批量添加
await es.bulk({
  refresh: true,
  operations,
});

async function get_completion(prompt: string, model: string = String(process.env.OPENAI_MODEL)) {
  const response = await openai.chat.completions.create({
    model,
    temperature: 0,
    messages: [{ "role": "user", "content": prompt }],
  });

  return response.choices[0].message.content;
}

const query = 'how many parameters does llama 2 have?';

// const context = await es.search({
//   index: 'my_index4',
//   body: {
//     knn: {
//       field: 'content_vector',  // 向量字段名
//       query_vector: (await getEmbedding(query))[0].embedding,
//       k: 5,                     // 返回最近邻数量
//       num_candidates: 100,      // 候选池大小(精度与性能的平衡)
//     },
//     _source: ['id', 'content'], // 返回的字段
//   }
// });


// 混合检索
const context = await es.search({
  index: 'my_index4',
  body: {
    query: {
      bool: {
        should: [
          // 语义搜索部分
          {
            knn: {
              field: 'content_vector',
              query_vector: (await getEmbedding(query))[0].embedding,
              k: 50,
              //  定义了候选文档的数量。num_candidates: 100,那么 Elasticsearch 会首先从 100 个候选文档中选择最相似的。
              // 这个数量的选择会影响 KNN 查询的性能与准确性。较大的候选数量可以增加精度,但也会增加计算量。
              num_candidates: 100, 
              boost: 0.5 // 权重
            }
          },
          // 文本搜索部分
          {
            match: {
              content: {
                query: query,
                boost: 0.5
              }
            }
          }
        ]
      }
    },
    _source: ['id', 'content'],
    size: 5
  }
});

console.log(context);
const promptTemplate = `
你是一个问答机器人。
你的任务是根据下述给定的已知信息回答用户问题。

已知信息:
${JSON.stringify(context)}

用户问:
${query}

如果已知信息不包含用户问题的答案,或者已知信息不足以回答用户的问题,请直接回复"我无法回答您的问题"。
请不要输出已知信息中不包含的信息或答案。
请用中文回答用户问题。
`;

console.log(await get_completion(promptTemplate));
相关推荐
Asthenia041228 分钟前
Spring扩展点与工具类获取容器Bean-基于ApplicationContextAware实现非IOC容器中调用IOC的Bean
后端
bobz9651 小时前
ovs patch port 对比 veth pair
后端
Asthenia04121 小时前
Java受检异常与非受检异常分析
后端
uhakadotcom1 小时前
快速开始使用 n8n
后端·面试·github
JavaGuide1 小时前
公司来的新人用字符串存储日期,被组长怒怼了...
后端·mysql
bobz9651 小时前
qemu 网络使用基础
后端
Asthenia04122 小时前
面试攻略:如何应对 Spring 启动流程的层层追问
后端
Asthenia04122 小时前
Spring 启动流程:比喻表达
后端
Asthenia04122 小时前
Spring 启动流程分析-含时序图
后端
ONE_Gua3 小时前
chromium魔改——CDP(Chrome DevTools Protocol)检测01
前端·后端·爬虫