在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_vector 是 dense_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 生成最终的回答，并打印在控制台中。

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));