RAG知识库-LlamaIndex文档切片

一、为什么需要文档切片

在处理文本数据时，我们需要将这些信息转变成向量格式，比如说把1000个文字变成1024维的向量。为了让向量化流程更加精准做，避免内容丢失，我们需要先对文档进行切片，并进行的数据标准化处理。

1. 提高向量化质量：通过适当的切片，可以确保每个切片的内容具有一致的格式和上下文。这有助于在转换为向量时保留更多的语义信息，提高向量化表示的质量和准确性。

2. 数据处理更高效：处理非常大的文档或长篇文章时，如果不进行切片，处理这些数据将非常耗时且资源密集。文档切片可以将数据分成较小的部分，提高处理速度，并减少内存消耗。

3. 避免信息丢失：长文档在进行向量化时，如果不进行切片，可能会因信息过多而导致信息丢失或重要内容被忽略。切片可以确保每个片段的独立性和信息完备性。

4. 提高检索准确性：对于文本搜索和信息检索，经过切片处理的文档能够提供更精细的细节和上下文信息，从而提高搜索结果的相关性和精确度。

二、切片工具种类

○ 语义段落分割
    ■ 特点：对纯文本、对话、书籍等自然语言文档分割更加出色，更加关注分割语句的完整性。
    ■ 原理：语句边界匹配
○  Token计数分割
    ■ 特点：基于Token的分割工具，关注分割的文本长度或固定Tokens长度的分割，更适合需精确控制 Token 数的场景。
    ■ 原理：Tokens计数分割
○ 代码语法分割
    ■ 特点：
        ● 针对编程语言语法分割（如 Python、JavaScript）。
        ● 按函数、类或代码块切分，保留逻辑完整性。
    ■ 原理：代码语法结构分析（AST）
○ Markdown语法分割
    ■ 特点：基于Markdown语法处理md格式内容，可以按文档结构进行分割处理。
    ■ 原理：基于Markdown语法进行字符匹配与分割

三、各框架切片工具种类横向对比

四、数据入库流程

五、具体代码实现

1. 文本处理类

import { Provide, Inject, Config, Logger } from '@midwayjs/core';
import { ILogger } from '@midwayjs/logger';
import {
  SentenceSplitter,
  TokenTextSplitter,
  CodeSplitter,
  Document,
} from 'llamaindex';
import * as Parser from 'tree-sitter';
import { RagMdParseService } from './RagMdParseService';
import axios from 'axios';
import { getFileType } from '../utils/index';
const TS = require('tree-sitter-typescript');
const { URL } = require('url');

@Provide()
export class RagDocChunkService {
  @Inject()
  ctx;

  @Config('env')
  env: string;

  @Config(`assetsUrlWhiteList`)
  assetsUrlWhiteList: string[];

  @Logger()
  logger: ILogger;

  @Inject()
  RagMdParseService: RagMdParseService;

  async parseUrlHost(url: string) {
    const { hostname } = new URL(url);
    if (!this.assetsUrlWhiteList.includes(hostname))
      throw new Error(`hostname不合法[${hostname}]`);
    return url;
  }

  /**
   * 文本分割
   * @param splitConfig
   * @returns
   */
  async textSplitToChunks(splitConfig, dbInfo) {
    const {
      content,
      chunkSize = 1024,
      chunkOverlap = 200,
      splitter,
      extractImageInfo = false,
      separator = undefined,
      paragraphSeparator = undefined,
      backupSeparators = undefined,
    } = splitConfig;
    let textSplitter = null;
    // 文档对象列表
    let textDocumentList = [];
    // 创建文本分割器
    if (splitter === 'SentenceSplitter') {
      textSplitter = new SentenceSplitter({
        chunkSize: chunkSize,
        chunkOverlap: chunkOverlap,
        separator: separator || undefined,
        paragraphSeparator: paragraphSeparator || undefined,
      });
    }
    // token分割器
    else if (splitter === 'TokenTextSplitter') {
      textSplitter = new TokenTextSplitter({
        chunkSize: chunkSize,
        chunkOverlap: chunkOverlap,
        separator: separator || undefined,
        backupSeparators: backupSeparators?.length
          ? backupSeparators
          : undefined,
      });
    }
    // 代码分割器
    else if (splitter === 'CodeSplitter') {
      const CodeParser = new Parser();
      CodeParser.setLanguage(TS.typescript as Parser.Language);
      textSplitter = new CodeSplitter({
        maxChars: chunkSize,
        getParser: () => CodeParser,
      });
    }
    // Mardown分割器
    else if (splitter === 'MarkdownNodeParser') {
      let markdownContent = content;
      // 提取图片摘要
      if (extractImageInfo === true)
        markdownContent = await this.RagMdParseService.markdownImageFormat(
          content,
          dbInfo.vlmModel
        );
      const mdDocuments = await this.RagMdParseService.markdownToDocuments(
        markdownContent
      );
      textDocumentList = mdDocuments;
    } else {
      throw new Error('参数[splitter]错误');
    }

    // 分割逻辑执行
    if (
      ['SentenceSplitter', 'TokenTextSplitter', 'CodeSplitter'].includes(
        splitter
      )
    ) {
      // 分隔的文档片段
      const splitDocsList = await textSplitter.splitText(content);
      // 文档片段
      textDocumentList = splitDocsList.map(
        text =>
          new Document({
            text: text,
          })
      );
    }
    // 绑定上下文与元关系
    for (let index = 0; index < textDocumentList.length; index++) {
      const doc = textDocumentList[index];
      doc.relation = {};
      if (textDocumentList?.[index - 1]?.id_)
        doc.relation.preNode = textDocumentList?.[index - 1]?.id_;
      if (textDocumentList?.[index + 1]?.id_)
        doc.relation.nextNode = textDocumentList?.[index + 1]?.id_;
    }

    return textDocumentList;
  }

  /**
   * 文件分割
   * @param splitConfig
   * @param dbInfo
   */
  async fileSplitToChunks(splitConfig, dbInfo) {
    const { files } = splitConfig;
    return Promise.all(
      files.map(async ({ name, url }) => {
        const parseUrl = await this.parseUrlHost(url);

        const { data: fileStringContent } = await axios({
          method: 'get',
          url: parseUrl,
          responseType: 'text',
        });

        const documents = await this.textSplitToChunks(
          {
            ...splitConfig,
            content: fileStringContent,
          },
          dbInfo
        );

        return {
          name,
          fileType: getFileType(name) || 'file',
          documents,
          url,
        };
      })
    );
  }
}

2. Markdown文件处理类

import { Provide, Inject, Config, Logger } from '@midwayjs/core';
import { ILogger } from '@midwayjs/logger';
import { encrypt } from '../utils/SecCrypto';
import { ChatOpenAI } from '@langchain/openai';
import { HumanMessage } from '@langchain/core/messages';
import { sleep } from '../utils/index';
import { AiChatService } from '../service/AiChatService';
import { MarkdownNodeParser, Document } from 'llamaindex';

interface ImageInfo {
  type: 'markdown' | 'html';
  alt: string;
  url: string;
  index: number;
  originalText: string;
}

@Provide()
export class RagMdParseService {
  @Inject()
  ctx;

  @Config('env')
  env: string;

  @Logger()
  logger: ILogger;

  @Inject()
  AiChatService: AiChatService;

  @Config('openBaseHost')
  openBaseHost: string;

  @Config('openBaseAppKey')
  openBaseAppKey: string;

  @Config('openBaseAppName')
  openBaseAppName: string;

  // 内容提取最大重试次数
  maxRetryTimes: number = 2;
  // 提取内容最大长度
  maxExtractContLen: number = 150;

  /**
   * 匹配markdown中所有的图片
   * @param content
   * @returns
   */
  extractImages(content: string): ImageInfo[] {
    const images: ImageInfo[] = [];
    let index = 0;

    // 匹配 Markdown 图片语法
    const markdownRegex = /!\[(.*?)\]\((.*?)\)/g;
    let match;
    while ((match = markdownRegex.exec(content)) !== null) {
      images.push({
        type: 'markdown',
        alt: match[1],
        url: match[2],
        index: index++,
        originalText: match[0],
      });
    }

    // 匹配 HTML img 标签
    const htmlRegex = /<img[^>]+src="([^"]+)"[^>]*>/g;
    while ((match = htmlRegex.exec(content)) !== null) {
      const imgTag = match[0];
      const url = match[1];

      // 尝试提取已存在的 alt 属性
      const altMatch = imgTag.match(/alt="([^"]*)"/) || ['', ''];
      const alt = altMatch[1];

      images.push({
        type: 'html',
        alt,
        url,
        index: index++,
        originalText: imgTag,
      });
    }

    return images;
  }

  /**
   * 替换 Markdown 内容中的图片描述
   * @param content
   * @param images
   * @param descriptions
   * @returns
   */
  replaceImageDescriptions(
    content: string,
    images: ImageInfo[],
    descriptions: string[]
  ): string {
    let result = content;

    // 从后向前替换，避免位置变化影响
    for (let i = images.length - 1; i >= 0; i--) {
      const image = images[i];
      const description = descriptions[i];
      let newText: string;

      if (image.type === 'markdown') {
        // 替换 Markdown 图片
        newText = `![${description}](${image.url})`;
      } else {
        // 替换 HTML img 标签
        const styleMatch = image.originalText.match(/style="[^"]*"/);
        const style = styleMatch ? ` ${styleMatch[0]}` : '';
        newText = `<img src="${image.url}" alt="${description}"${style} />`;
      }

      result = result.replace(image.originalText, newText);
    }

    return result;
  }

  /**
   * Markdown图片信息统一处理
   * @param markdownContent
   */
  async markdownImageFormat(markdownContent: string, vlmModel: string) {
    // 1. 提取所有图片
    const images = this.extractImages(markdownContent);

    if (images.length === 0) return markdownContent;

    // 2. 为每个图片生成摘要
    const descriptions = await Promise.all(
      images.map(async img => {
        return await this.mardownImageInfoExtract(img.url, 0, vlmModel);
      })
    );

    // 3. 替换原文中的图片描述
    return this.replaceImageDescriptions(markdownContent, images, descriptions);
  }

  /**
   * 图片摘要信息提取
   * @param imgSrc
   */
  async mardownImageInfoExtract(
    imgUrl: string,
    retryTimes: number = 0,
    vlmModel
  ) {
    try {
      if (retryTimes > this.maxRetryTimes) return 'ImageInfoExtractTimeout';
      const appKey = this.openBaseAppKey;
      const appName = this.openBaseAppName;
      const timestemp = new Date().getTime();
      const signature = encrypt(`${timestemp}`, appKey);
      const chatInstance = new ChatOpenAI({
        model: `${vlmModel}`,
        apiKey: ``,
        configuration: {
          baseURL: `${this.openBaseHost}/open-base-api/`,
          defaultHeaders: {
            'app-name': appName,
            signature: signature,
          },
        },
      });
      const result = await chatInstance.invoke([
        new HumanMessage({
          content: [
            {
              type: 'text',
              text: `
              ## 角色
              你是一个图片内容识别助手，请帮我提取该图片所描述的内容，并生成一份简短的摘要信息
              ## 要求
              1. 摘要内容不要过长，请控制在${this.maxExtractContLen}字内
              2. 摘要内容不能包含任何换行符，请保证数据为最基本的字符串类型
              3. 摘要内容请以中文为主，代码与专有名词可以保持英文
              `,
            },
            {
              type: 'image_url',
              image_url: {
                url: imgUrl,
              },
            },
          ],
        }),
      ]);
      this.ctx.logger.info(`[mardownImageInfoExtract Result]`, result);
      // 无内容重试
      if (!result?.content) {
        await sleep(200);
        return this.mardownImageInfoExtract(imgUrl, retryTimes + 1, vlmModel);
      }
      // 记录使用日志
      if (result?.usage_metadata) {
        this.AiChatService.insertModelUseRecord({
          model: chatInstance.model,
          totalTokens: result?.usage_metadata?.total_tokens,
          origin: 'rag_image_extract',
          status: 'success',
          user: this.ctx?.user?.workid,
          outputTokens: result?.usage_metadata?.output_tokens,
          inputTokens: result?.usage_metadata?.input_tokens,
          apiPath: this.ctx.request.path,
        });
      }
      // 正常返回
      return result?.content;
    } catch (error) {
      this.ctx.logger.error(
        `[mardownImageInfoExtract Error]`,
        error.toString()
      );
      // 继续重试
      if (retryTimes < this.maxRetryTimes) {
        await sleep(200);
        return this.mardownImageInfoExtract(imgUrl, retryTimes + 1, vlmModel);
      }
      // 报错返回
      else return 'ImageInfoExtractError';
    }
  }

  /**
   * markdown内容转换为Document列表
   * @param markdownContent
   * @returns
   */
  async markdownToDocuments(markdownContent: string) {
    const MarkdownParser = new MarkdownNodeParser();
    const markdownNodes = MarkdownParser.getNodesFromDocuments([
      new Document({
        text: markdownContent,
      }),
    ]);

    return markdownNodes.map(
      ({ text, metadata }) =>
        new Document({
          metadata,
          text: `${Object.keys(metadata).reduce(
            (res, key) => res + `${key}: ${metadata[key]}\n`,
            ''
          )} ${text}`,
        })
    );
  }
}