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