文档多分块与检索

• 项目目标
• 系统核心概念
• 功能实现步骤
• 分层向量与检索
  • 功能背景
  • 实现流程
  • 文档拆分代码
  • 向量检索
  • 拼接提示词
• 修改 知识库的 命中测试
• 修改 应用的 命中测试
• 结论

项目目标

本项目旨在构建一个功能完备的 RAG（Retrieval-Augmented Generation）系统，主要目标包括：

• 知识库管理：支持创建、更新和删除知识库，便于用户高效维护内容。
• 文档处理：包括文档的拆分、片段的向量化处理，以提升检索效率和准确性。
• 问答系统：提供高效的向量检索和实时生成回答的能力，支持复杂汇总类问题的处理。
• 系统优化：通过统计分析和推理问答调试，不断优化系统性能和用户体验。

系统核心概念

在 RAG 系统中，以下是几个核心概念：

• 应用：知识库的集合。每个应用可以自定义提示词，以满足不同的个性化需求。
• 知识库：由多个文档组成，便于用户对内容进行分类和管理。
• 文档：系统中对应的真实文档内容。
• 片段：文档经过拆分后的最小内容单元，用于更高效的处理和检索。

功能实现步骤

1. 数据库设计 查看 01.md
   设计并实现项目所需的数据表结构与数据库方案，为后续的数据操作打下坚实基础。

2. 用户登录 查看 02.md
   实现了安全可靠的用户认证系统，保护用户数据并限制未经授权的访问。

3. 模型管理 查看 03.md
   支持针对不同平台的模型（如 OpenAI、Google Gemini、Claude）进行管理与配置。

4. 知识库管理 查看 04.md
   提供创建、更新及删除知识库的功能，方便用户维护与管理文档内容。

5. 文档拆分 查看 05.md
   可将文档拆分为多个片段，便于后续向量化和检索操作。

6. 片段向量 查看 06.md
   将文本片段进行向量化处理，以便进行语义相似度计算及高效检索。

7. 命中率测试 查看 07.md
   通过语义相似度和 Top-N 算法，检索并返回与用户问题最相关的文档片段，用于评估检索的准确性。

8. 文档管理 查看 08.md
   提供上传和管理文档的功能，上传后可自动拆分为片段便于进一步处理。

9. 片段管理 查看 09.md
   允许对已拆分的片段进行增、删、改、查等操作，确保内容更新灵活可控。

10. 问题管理 查看 10.md
    为片段指定相关问题，以提升检索时的准确性与关联度。

11. 应用管理 查看 11.md
    提供创建和配置应用（智能体）的功能，并可关联指定模型和知识库。

12. 向量检索 查看 12.md
    基于语义相似度，在知识库中高效检索与用户问题最匹配的片段。

13. 推理问答调试 查看 13.md
    提供检索与问答性能的评估工具，帮助开发者进行系统优化与调试。

14. 对话问答 查看 14.md
    为用户提供友好的人机交互界面，结合检索到的片段与用户问题实时生成回答。

15. 统计分析 查看 15.md
    对用户的提问与系统回答进行数据化分析，并以可视化图表的形式呈现系统使用情况。

16. 用户管理 查看 16.md
    提供多用户管理功能，包括用户的增删改查及权限控制。

17. API 管理 查看 17.md
    对外提供标准化 API，便于外部系统集成和调用本系统的功能。

18. 存储文件到 S3 查看 18.md
    将用户上传的文件存储至 S3 等对象存储平台，提升文件管理的灵活性与可扩展性。

19. 文档解析优化 查看 19.md
    介绍与对比常见的文档解析方案，并提供提升文档解析速度和准确性的优化建议。

20. 片段汇总 查看 20.md
    对片段内容进行汇总，以提升总结类问题的查询与回答效率。

21. 文档多分块与检索 查看 21.md
    将片段进一步拆分为句子并进行向量检索，提升检索的准确度与灵活度。

22. 多文档支持 查看 22.md
    兼容多种文档格式，包括 .doc, .docx, .xls, .xlsx, .ppt, .pptx 等。

23. 对话日志 查看 23.md
    记录并展示对话日志，用于后续分析和问题回溯。

24. 检索性能优化 查看 24.md
    提供整库扫描和分区检索等多种方式，进一步提高检索速度和效率。

25. Milvus 查看 25.md
    将向量数据库切换至 Milvus，以在大规模向量检索场景中获得更佳的性能与可扩展性。

26. 文档解析方案和费用对比 查看 26.md
    对比不同文档解析方案在成本、速度、稳定性等方面的差异，为用户提供更加经济高效的选择。

27. 爬取网页数据 查看 27.md
    支持从网页中抓取所需内容，后续处理流程与本地文档一致：分段、向量化、存储与检索。

分层向量与检索

功能背景

在 RAG 系统中，向量检索对于单个问题的检索准确度较高，但是如何拆分的段落较大再进行向量的相似度计算时 得出的相似度会较低。为了解决这一问题，本步骤引入了分层向量与检索的方法，通过对片段进行更细粒度的拆分和管理和在检索时结合 sentence 和 paragraph 的内容，提高检索的精准度和灵活性。

实现流程

1. 文档拆分

文档拆分是将原始文档分解为更小的单元，以便于向量化和高效检索。本系统将文档拆分为段落级别和句子级别两种粒度。

片段级别（Paragraph）

• Chunk Size（块大小）： 1000 个 token
• Chunk Overlap（块重叠）： 400 个 token
• Embedding Model（嵌入模型）： text-embedding-3-large

段落级别的拆分适用于较长的文本内容，通过较大的块大小和一定的重叠，确保每个片段包含足够的上下文信息，以便后续的向量化处理。

句子级别（Sentence）

• Chunk Size（块大小）： 150 个 token
• Chunk Overlap（块重叠）： 50 个 token
• Embedding Model（嵌入模型）： text-embedding-3-large

句子级别的拆分进一步细化段落内容，将其拆分为更小的句子单元。这种拆分方式有助于提高检索的精准度，特别是在处理复杂的汇总类问题时，能够更有效地提取相关信息。

2. 检索流程

分层向量与检索的具体流程如下：

1. 问题向量化

   • 用户输入问题后，首先将其转换为向量表示，即问题向量。

2. 句子级检索

   • 使用问题向量在 max_kb_sentence 表中进行相似性检索，筛选出最相关的句子。

3. 片段获取

   • 根据检索出的句子信息，提取其对应的 paragraph_id，并从 paragraph 表中查询完整的片段（Paragraph）。

4. 上下文构建

   • 从检索得到的片段中选取10 个最相关的片段，作为大模型推理时的上下文。

5. 大模型推理

   • 结合检索到的片段和设计的提示词（Prompt），将构建的上下文输入大模型，进行推理并生成最终答案。

文档拆分代码

在本系统的设计中，文档拆分分为两步：

1. 用户生成文档

   • 使用大模型对文档进行识别，对识别后的文档进行段落拆分，并将拆分后的段落返回给前端。

2. 段落修改与提交

   • 用户可以对拆分后的段落进行修改，修改后提交到后端。

3. 句子级拆分与存储

   • 通过 DatasetDocumentVectorService 的 batch 方法解析前端上传的段落，再拆分成句子级别的向量并存储入库。随后对段落进行汇总并存储入库。

以下是相关的代码实现：

MaxKbSentenceService 类

该类负责将段落摘要转换为句子级别的记录，并进行向量化处理，最终存储到数据库中。

package com.litongjava.maxkb.service.kb;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Future;

import org.postgresql.util.PGobject;

import com.litongjava.db.activerecord.Db;
import com.litongjava.db.activerecord.Row;
import com.litongjava.jfinal.aop.Aop;
import com.litongjava.maxkb.constant.TableNames;
import com.litongjava.maxkb.model.MaxKbSentence;
import com.litongjava.maxkb.utils.ExecutorServiceUtils;
import com.litongjava.tio.utils.crypto.Md5Utils;
import com.litongjava.tio.utils.snowflake.SnowflakeIdUtils;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class MaxKbSentenceService {
  MaxKbEmbeddingService maxKbEmbeddingService = Aop.get(MaxKbEmbeddingService.class);
  CompletionService<MaxKbSentence> completionServiceSentence = new ExecutorCompletionService<>(ExecutorServiceUtils.getExecutorService());

  CompletionService<Row> completionServiceRow = new ExecutorCompletionService<>(ExecutorServiceUtils.getExecutorService());
  MaxKbParagraphSummaryService maxKbParagraphSummaryService = Aop.get(MaxKbParagraphSummaryService.class);

  public boolean summaryToSentenceAndSave(Long dataset_id, String modelName, List<Row> paragraphRecords, Long documentIdFinal) {
    //省略多余的代码
  }

  public boolean splitToSentenceAndSave(Long dataset_id, String modelName, List<Row> paragraphRecords, Long documentIdFinal) {
    boolean transactionSuccess;
    List<MaxKbSentence> sentences = new ArrayList<>();
    for (Row paragraph : paragraphRecords) {
      String paragraphContent = paragraph.getStr("content");

      //继续拆分片段 为句子
      Document document = new Document(paragraphContent);
      // 使用较大的块大小（150）和相同的重叠（50）
      DocumentSplitter splitter = DocumentSplitters.recursive(150, 50, new OpenAiTokenizer());
      List<TextSegment> segments = splitter.split(document);
      for (TextSegment segment : segments) {
        String sentenceContent = segment.text();

        MaxKbSentence maxKbSentence = new MaxKbSentence();
        maxKbSentence.setId(SnowflakeIdUtils.id()).setType(1).setHitNum(0)
            //
            .setMd5(Md5Utils.getMD5(sentenceContent)).setContent(sentenceContent)
            //
            .setDatasetId(dataset_id).setDocumentId(documentIdFinal).setParagraphId(paragraph.getLong("id"));

        sentences.add(maxKbSentence);
      }

    }

    List<Future<Row>> futures = new ArrayList<>(sentences.size());
    for (MaxKbSentence sentence : sentences) {
      futures.add(completionServiceRow.submit(() -> {
        PGobject vector = maxKbEmbeddingService.getVector(sentence.getContent(), modelName);
        Row record = sentence.toRecord();
        record.set("embedding", vector);
        return record;
      }));
    }

    List<Row> sentenceRows = new ArrayList<>();
    for (int i = 0; i < sentences.size(); i++) {
      try {
        Future<Row> future = completionServiceRow.take();
        Row record = future.get();
        if (record != null) {
          sentenceRows.add(record);
        }
      } catch (Exception e) {
        log.error(e.getMessage(), e);
      }
    }

    transactionSuccess = Db.tx(() -> {
      Db.delete(TableNames.max_kb_sentence, Row.by("document_id", documentIdFinal));
      Db.batchSave(TableNames.max_kb_sentence, sentenceRows, 2000);
      return true;
    });
    return transactionSuccess;
  }
}

代码说明：

• 并行处理：利用 CompletionService 实现异步处理，提升系统的处理效率。

  • 生成摘要：对每个段落内容调用 MaxKbParagraphSummaryService 生成摘要。
  • 生成向量：对生成的摘要内容调用 MaxKbEmbeddingService 进行向量化处理。

• 事务管理：在数据库操作中使用事务，确保数据的一致性和完整性。在保存新的句子记录前，删除已有的相关记录，避免数据重复。

• 错误处理：对异步任务中的异常进行捕获和日志记录，确保系统的稳定性。

DatasetDocumentVectorService 类

该类负责处理用户上传的文档，将其拆分为段落和句子级别的片段，并进行向量化处理和存储。

package com.litongjava.maxkb.service;

import java.util.ArrayList;
import java.util.List;

import com.jfinal.kit.Kv;
import com.litongjava.db.TableInput;
import com.litongjava.db.TableResult;
import com.litongjava.db.activerecord.Db;
import com.litongjava.db.activerecord.Row;
import com.litongjava.jfinal.aop.Aop;
import com.litongjava.maxkb.constant.TableNames;
import com.litongjava.maxkb.service.kb.MaxKbEmbeddingService;
import com.litongjava.maxkb.service.kb.MaxKbSentenceService;
import com.litongjava.maxkb.vo.DocumentBatchVo;
import com.litongjava.maxkb.vo.Paragraph;
import com.litongjava.model.result.ResultVo;
import com.litongjava.table.services.ApiTable;
import com.litongjava.tio.utils.crypto.Md5Utils;
import com.litongjava.tio.utils.hutool.FilenameUtils;
import com.litongjava.tio.utils.snowflake.SnowflakeIdUtils;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class DatasetDocumentVectorService {

  MaxKbEmbeddingService maxKbEmbeddingService = Aop.get(MaxKbEmbeddingService.class);
  MaxKbSentenceService maxKbSentenceService = Aop.get(MaxKbSentenceService.class);

  public ResultVo batch(Long userId, Long dataset_id, List<DocumentBatchVo> list) {

    TableInput tableInput = new TableInput();
    tableInput.set("id", dataset_id);
    if (!userId.equals(1L)) {
      tableInput.set("user_id", userId);
    }

    TableResult<Row> result = ApiTable.get(TableNames.max_kb_dataset, tableInput);

    Row dataset = result.getData();
    if (dataset == null) {
      return ResultVo.fail("Dataset not found.");
    }

    Long embedding_mode_id = dataset.getLong("embedding_mode_id");
    String sqlModelName = String.format("SELECT model_name FROM %s WHERE id = ?", TableNames.max_kb_model);
    String modelName = Db.queryStr(sqlModelName, embedding_mode_id);
    String sqlDocumentId = String.format("SELECT id FROM %s WHERE user_id = ? AND file_id = ?", TableNames.max_kb_document);

    List<Kv> kvs = new ArrayList<>();

    for (DocumentBatchVo documentBatchVo : list) {
      Long fileId = documentBatchVo.getId();
      String filename = documentBatchVo.getName();
      Long documentId = Db.queryLong(sqlDocumentId, userId, fileId);

      List<Paragraph> paragraphs = documentBatchVo.getParagraphs();
      int char_length = 0;
      int size = 0;
      if (paragraphs != null) {
        size = paragraphs.size();
        for (Paragraph p : paragraphs) {
          if (p.getContent() != null) {
            char_length += p.getContent().length();
          }
        }
      }

      String type = FilenameUtils.getSuffix(filename);

      if (documentId == null) {
        documentId = SnowflakeIdUtils.id();
        Row record = Row.by("id", documentId)
            //
            .set("file_id", fileId).set("user_id", userId).set("name", filename)
            //
            .set("char_length", char_length).set("status", "1").set("is_active", true)
            //
            .set("type", type).set("dataset_id", dataset_id).set("paragraph_count", size)
            //
            .set("hit_handling_method", "optimization").set("directly_return_similarity", 0.9);
        Db.save(TableNames.max_kb_document, record);
        Kv kv = record.toKv();
        kvs.add(kv);
      } else {
        Row existingRecord = Db.findById(TableNames.max_kb_document, documentId);
        if (existingRecord != null) {
          Kv kv = existingRecord.toKv();
          kvs.add(kv);
        } else {
          // Handle the case where documentId is provided but the record does not exist
          return ResultVo.fail("Document not found for ID: " + documentId);
        }
      }

      List<Row> paragraphRecords = new ArrayList<>();
      Long documentIdFinal = documentId;
      boolean transactionSuccess = saveToParagraph(dataset_id, fileId, documentId, paragraphs, type, paragraphRecords);

      if (!transactionSuccess) {
        return ResultVo.fail("Transaction failed while saving paragraphs for document ID: " + documentIdFinal);
      }

      transactionSuccess = Aop.get(MaxKbSentenceService.class).summaryToSentenceAndSave(dataset_id, modelName, paragraphRecords, documentIdFinal);

      if (!transactionSuccess) {
        return ResultVo.fail("Transaction failed while summary paragraph for document ID: " + documentIdFinal);
      }

      transactionSuccess = maxKbSentenceService.splitToSentenceAndSave(dataset_id, modelName, paragraphRecords, documentIdFinal);

      if (!transactionSuccess) {
        return ResultVo.fail("Transaction failed while saving senttents for document ID: " + documentIdFinal);
      }

    }

    return ResultVo.ok(kvs);
  }

  private boolean saveToParagraph(Long dataset_id, Long fileId, Long documentId, List<Paragraph> paragraphs,
      //
      String type, List<Row> batchRecord) {
    final long documentIdFinal = documentId;
    if (paragraphs != null) {
      for (Paragraph p : paragraphs) {

        String title = p.getTitle();
        String content = p.getContent();

        Row row = Row.by("id", SnowflakeIdUtils.id())
            //
            .set("source_id", fileId)
            //
            .set("source_type", type)
            //
            .set("title", title)
            //
            .set("content", content)
            //
            .set("md5", Md5Utils.getMD5(content))
            //
            .set("status", "1")
            //
            .set("hit_num", 0)
            //
            .set("is_active", true).set("dataset_id", dataset_id).set("document_id", documentIdFinal);
        batchRecord.add(row);

      }
    }
    return Db.tx(() -> {
      Db.delete(TableNames.max_kb_paragraph, Row.by("document_id", documentIdFinal));
      Db.batchSave(TableNames.max_kb_paragraph, batchRecord, 2000);
      return true;
    });
  }
}

代码说明：

• 批量处理：该类通过 batch 方法批量处理用户上传的文档，将其拆分为段落和句子级别的片段，并进行向量化处理和存储。

• 事务管理：在保存段落和句子记录时，使用数据库事务确保操作的原子性，避免数据不一致的情况。

• 错误处理：在处理过程中，若任何一步操作失败，将返回相应的错误信息，确保系统的稳定性和可靠性。

向量检索

在系统的初期设计中，向量检索是针对 max_kb_paragraph 表进行的。然而，随着系统的发展，发现对 max_kb_sentence 表进行向量相似度计算能够提供更高的检索准确性。因此，现将向量检索的目标从段落级别调整为句子级别。

检索 SQL 代码

以下是针对 max_kb_sentence 表进行向量相似度计算的 SQL 查询：

--# kb.search_sentense_related_paragraph__with_dataset_ids
SELECT DISTINCT
  sub.id,
  sub.content,
  sub.title,
  sub.status,
  sub.hit_num,
  sub.is_active,
  sub.dataset_id,
  sub.document_id,
  sub.dataset_name,
  sub.document_name,
  sub.similarity,
  sub.similarity AS comprehensive_score
FROM (
  SELECT
    d.name AS document_name,
    ds.name AS dataset_name,
    s.id,
    s.content,
    CASE
      WHEN s.type = 1 THEN p.title
      ELSE ''
    END AS title,
    CASE
      WHEN s.type = 1 THEN p.status
      ELSE '1'
    END AS status,
    s.hit_num,
    CASE
      WHEN s.type = 1 THEN p.is_active
      ELSE true
    END AS is_active,
    s.dataset_id,
    s.document_id,
    (1 - (s.embedding <=> c.v)) AS similarity
  FROM
    max_kb_sentence s
  JOIN
    max_kb_embedding_cache c ON c.id = ?
  JOIN
    max_kb_dataset ds ON s.dataset_id = ds.id
  JOIN
    max_kb_document d ON s.document_id = d.id
  LEFT JOIN
    max_kb_paragraph p ON s.type = 1 AND s.paragraph_id = p.id
  WHERE
    (
      (s.type = 1 AND p.is_active = TRUE)
      OR
      (s.type = 2 AND d.is_active = TRUE)
    )
    AND s.deleted = 0
    AND ds.deleted = 0
    AND s.dataset_id = ANY (?)
) sub
WHERE
  sub.similarity > ?
ORDER BY
  sub.similarity DESC
LIMIT ?;

SQL 说明：

• 相似度计算：通过 (1 - (s.embedding <=> c.v)) AS similarity 计算句子向量与问题向量之间的相似度。

• 条件过滤：仅检索未删除且活跃的句子和文档，同时根据提供的 dataset_id 进行过滤。

• 结果排序与限制：根据相似度降序排序，并限制返回的结果数量，以提升检索效率。

调用 SQL 的检索代码

以下是调用上述 SQL 进行向量检索的 Java 代码实现：

package com.litongjava.maxkb.service.kb;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import com.litongjava.db.activerecord.Db;
import com.litongjava.db.activerecord.Row;
import com.litongjava.jfinal.aop.Aop;
import com.litongjava.maxkb.vo.MaxKbSearchStep;
import com.litongjava.maxkb.vo.ParagraphSearchResultVo;
import com.litongjava.openai.constants.OpenAiModels;
import com.litongjava.template.SqlTemplates;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class MaxKbParagraphSearchService {

  public MaxKbSearchStep search(Long[] datasetIdArray, Float similarity, Integer top_n, String quesiton) {
    MaxKbSearchStep maxKbSearchStep = new MaxKbSearchStep();
    long start = System.currentTimeMillis();
    List<ParagraphSearchResultVo> results = search0(datasetIdArray, similarity, top_n, quesiton);
    long end = System.currentTimeMillis();
    maxKbSearchStep.setStep_type("step_type")
        //
        .setModel_name("text_embedding_3_large").setProblem_text("problem_text")
        //
        .setCost(0).setRun_time(((double) (end - start)) / 1000)
        //
        .setParagraph_list(results);
    return maxKbSearchStep;
  }

  private List<ParagraphSearchResultVo> search0(Long[] datasetIdArray, Float similarity, Integer top_n, String quesiton) {
    Long vectorId = Aop.get(MaxKbEmbeddingService.class).getVectorId(quesiton, OpenAiModels.text_embedding_3_large);
    String sql = SqlTemplates.get("kb.search_sentense_related_paragraph__with_dataset_ids");

    List<Row> records = Db.find(sql, vectorId, datasetIdArray, similarity, top_n);

    log.info("search_paragraph:{},{},{},{},{}", vectorId, Arrays.toString(datasetIdArray), similarity, top_n, records.size());
    List<ParagraphSearchResultVo> results = new ArrayList<>();
    for (Row record : records) {
      ParagraphSearchResultVo vo = record.toBean(ParagraphSearchResultVo.class);
      results.add(vo);
    }
    return results;
  }

  public MaxKbSearchStep searchV1(Long[] datasetIdArray, Float similarity, Integer top_n, String quesiton) {
    MaxKbSearchStep maxKbSearchStep = new MaxKbSearchStep();
    long start = System.currentTimeMillis();
    List<ParagraphSearchResultVo> results = searchV10(datasetIdArray, similarity, top_n, quesiton);
    long end = System.currentTimeMillis();
    maxKbSearchStep.setStep_type("step_type")
        //
        .setModel_name("text_embedding_3_large").setProblem_text("problem_text")
        //
        .setCost(0).setRun_time(((double) (end - start)) / 1000)
        //
        .setParagraph_list(results);
    return maxKbSearchStep;
  }

  public List<ParagraphSearchResultVo> searchV10(Long[] datasetIdArray, Float similarity, Integer top_n,
      //
      String quesiton) {
    Long vectorId = Aop.get(MaxKbEmbeddingService.class).getVectorId(quesiton, OpenAiModels.text_embedding_3_large);
    String sql = SqlTemplates.get("kb.search_paragraph_with_dataset_ids");

    List<Row> records = Db.find(sql, vectorId, datasetIdArray, similarity, top_n);

    log.info("search_paragraph:{},{},{},{},{}", vectorId, Arrays.toString(datasetIdArray), similarity, top_n, records.size());
    List<ParagraphSearchResultVo> results = new ArrayList<>();
    for (Row record : records) {
      ParagraphSearchResultVo vo = record.toBean(ParagraphSearchResultVo.class);
      results.add(vo);
    }
    return results;
  }
}

代码说明：

• 向量化问题：通过 MaxKbEmbeddingService 将用户问题转换为向量表示，以便进行相似度计算。

• SQL 调用：利用预定义的 SQL 模板 kb.search_sentense_related_paragraph__with_dataset_ids 进行句子级别的向量检索。

• 结果处理：将查询结果转换为 ParagraphSearchResultVo 对象列表，供后续的上下文构建和大模型推理使用。

拼接提示词

在向量检索后，需要将检索到的相关内容拼接成适合大模型推理的提示词格式。以下是相关的代码实现：

MaxKbChatDataXMLGenerator 类

该类负责将检索结果转换为 XML 格式的提示词，以便于大模型进行推理。

package com.litongjava.maxkb.service.kb;

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import com.litongjava.maxkb.vo.ParagraphSearchResultVo;

public class MaxKbChatDataXMLGenerator {

  public static String generateXML(List<ParagraphSearchResultVo> records) {
    StringBuilder data = new StringBuilder();
    data.append("<data>");

    // 将记录按 source 列表进行分组
    Map<Long, List<ParagraphSearchResultVo>> groupedBySource = records.stream().collect(Collectors.groupingBy(ParagraphSearchResultVo::getDocument_id));

    int i = 0;
    for (Map.Entry<Long, List<ParagraphSearchResultVo>> entry : groupedBySource.entrySet()) {
      List<ParagraphSearchResultVo> list = entry.getValue();
      data.append("<record>");

      // source
      data.append("<source>").append(i + 1).append("</source>");

      String documentNames = list.get(0).getDocument_name();

      String contents = list.stream().map(ParagraphSearchResultVo::getContent).map(MaxKbChatDataXMLGenerator::escapeXml).collect(Collectors.joining("\r\n"));

      data.append("<document_name>").append(documentNames).append("</document_name>");
      data.append("<contents>").append(contents).append("</contents>");

      data.append("</record>");
      i++;
    }

    data.append("</data>");
    return data.toString();
  }

  // XML 转义方法，确保特殊字符被正确处理
  private static String escapeXml(String input) {
    if (input == null) {
      return "";
    }
    return input.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace("\"", "&quot;").replace("'", "&apos;");
  }
}

代码说明：

• 分组处理：将检索到的句子按所属文档进行分组，确保每个文档对应一个记录。

• XML 构建：为每个分组构建 <record> 元素，包含 <source>, <document_name>, 和 <contents> 三个子元素。

• XML 转义：通过 escapeXml 方法对内容进行转义，确保生成的 XML 格式正确，避免特殊字符导致的解析错误。

MaxKbChatDataXMLGeneratorTest 类

该类为 MaxKbChatDataXMLGenerator 的测试类，用于验证 XML 生成的正确性。

package com.litongjava.maxkb.service.kb;

import java.util.ArrayList;
import java.util.List;

import org.junit.Test;

import com.litongjava.maxkb.vo.ParagraphSearchResultVo;

public class MaxKbChatDataXMLGeneratorTest {

  @Test
  public void test() {
    List<ParagraphSearchResultVo> records = getRecords(); // 假设此方法获取记录
    String xmlData = MaxKbChatDataXMLGenerator.generateXML(records);
    System.out.println(xmlData);
  }

  // 模拟获取记录的方法
  private static List<ParagraphSearchResultVo> getRecords() {
    // 您可以根据实际情况填充数据
    List<ParagraphSearchResultVo> records = new ArrayList<>();

    // 示例数据
    records.add(new ParagraphSearchResultVo(1l, "内容1", "Deloria and Olsen, AMS User's Guide, AMS 129.pdf", 474233893756760066L));
    records.add(new ParagraphSearchResultVo(2L, "内容2", "Deloria and Olsen, AMS User's Guide, AMS 129.pdf", 474233893756760066L));
    records.add(new ParagraphSearchResultVo(2L, "内容3", "Another Document.pdf", 474233893756760070L));

    return records;
  }
}

代码说明：

• 测试方法：通过模拟数据，测试 MaxKbChatDataXMLGenerator 类的 generateXML 方法，确保生成的 XML 格式符合预期。

• 示例数据：提供了三个示例记录，分别来自两个不同的文档，验证分组和内容拼接的正确性。

示例生成的 XML

<data>
<record>
	<source>1</source><document_name>Deloria and Olsen, AMS User's Guide, AMS 129.pdf</document_name>
	<contents>内容1内容2</contents>
</record>
<record>
	<source>2</source><document_name>Another Document.pdf</document_name>
	<contents>内容3</contents>
</record>
</data>

XML 说明：

• 结构清晰：每个 <record> 元素包含了源编号、文档名称以及对应的内容，便于大模型理解和处理。

• 内容聚合：同一文档下的多个句子内容被聚合在一起，形成连续的文本块，提高上下文的连贯性。

修改 知识库的 命中测试

在之前的 知识库的 命中测试 中 我们是对 max_kb_paragraph 表进行向量相对度检索,由于我们采用了分段检索,现在我们进行对 max_kb_sentence 进行检索 修改 com.litongjava.maxkb.service.kb.MaxKbDatasetHitTestService.embeddingSearch(Long, Long, String, Double, Integer) 中的调用 sql 的代码,修改成下面的形式

    String sql = SqlTemplates.get("kb.search_sentense_related_paragraph__with_dataset_ids");
    Long[] datasetIdArray = { datasetId };
    List<Row> records = Db.find(sql, vectorId, datasetIdArray, similarity, top_number);
    log.info("search_paragraph:{},{},{},{},{}", vectorId, Arrays.toString(datasetIdArray), similarity, top_number, records.size());

修改 应用的 命中测试

修改后的代码如下

package com.litongjava.maxkb.service.kb;

import java.util.Arrays;
import java.util.List;

import com.jfinal.kit.Kv;
import com.litongjava.db.TableInput;
import com.litongjava.db.TableResult;
import com.litongjava.db.activerecord.Db;
import com.litongjava.db.activerecord.Row;
import com.litongjava.jfinal.aop.Aop;
import com.litongjava.kit.RowUtils;
import com.litongjava.maxkb.constant.TableNames;
import com.litongjava.model.result.ResultVo;
import com.litongjava.openai.constants.OpenAiModels;
import com.litongjava.table.services.ApiTable;
import com.litongjava.template.SqlTemplates;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class MaxKbApplicationHitTestService {

  public ResultVo hitTest(Long userId, Long applicationId, String query_text, Double similarity, Integer top_number, String search_mode) {
    if ("embedding".equals(search_mode)) {
      return embeddingSearch(userId, applicationId, query_text, similarity, top_number);
    } else if ("blend".equals(search_mode)) {
      return blendSearch(userId, applicationId, query_text, similarity, top_number);
    }
    return null;
  }

  private ResultVo blendSearch(Long userId, Long applicationId, String query_text, Double similarity, Integer top_number) {
    // 获取数据集信息
    TableInput tableInput = new TableInput();
    tableInput.set("id", applicationId);
    if (userId.equals(1L)) {

    } else {
      tableInput.set("user_id", userId);
    }

    TableResult<Row> datasetResult = ApiTable.get(TableNames.max_kb_dataset, tableInput);

    Row dataset = datasetResult.getData();
    // 获取模型名称
    Long embeddingModeId = dataset.getLong("embedding_mode_id");
    String modelName = null;
    if (embeddingModeId != null) {
      modelName = Db.queryStr(String.format("SELECT model_name FROM %s WHERE id = ?", TableNames.max_kb_model), embeddingModeId);
      if (modelName == null) {
        modelName = OpenAiModels.text_embedding_3_large;
      }

    } else {
      modelName = OpenAiModels.text_embedding_3_large;
    }

    String sql = SqlTemplates.get("kb.list_database_id_by_application_id");
    List<Long> datasetIds = Db.queryListLong(sql, applicationId);
    if (datasetIds.size() < 1) {
      return ResultVo.fail("not found database of application id:", applicationId);
    }
    Long vectorId = Aop.get(MaxKbEmbeddingService.class).getVectorId(query_text, modelName);

    sql = SqlTemplates.get("kb.search_sentense_related_paragraph__with_dataset_ids");
    Long[] array = datasetIds.toArray(new Long[0]);
    List<Row> records = Db.find(sql, vectorId, array, similarity, top_number);
    List<Kv> kvs = RowUtils.recordsToKv(records, false);
    return ResultVo.ok(kvs);
  }

  private ResultVo embeddingSearch(Long userId, Long applicationId, String query_text, Double similarity, Integer top_number) {
    // 获取数据集信息
    TableInput tableInput = new TableInput();
    tableInput.set("id", applicationId);
    if (userId.equals(1L)) {

    } else {
      tableInput.set("user_id", userId);
    }

    TableResult<Row> datasetResult = ApiTable.get(TableNames.max_kb_dataset, tableInput);

    Row dataset = datasetResult.getData();
    // 获取模型名称
    String modelName = null;
    Long embeddingModeId = null;
    if (dataset != null) {
      embeddingModeId = dataset.getLong("embedding_mode_id");
      if (embeddingModeId != null) {
        modelName = Db.queryStr(String.format("SELECT model_name FROM %s WHERE id = ?", TableNames.max_kb_model), embeddingModeId);
      }
      if (modelName == null) {
        modelName = OpenAiModels.text_embedding_3_large;
      }

    } else {
      modelName = OpenAiModels.text_embedding_3_large;
    }

    String sql = SqlTemplates.get("kb.list_database_id_by_application_id");
    List<Long> datasetIds = Db.queryListLong(sql, applicationId);
    if (datasetIds.size() < 1) {
      return ResultVo.fail("not found database of application id:", applicationId);
    }
    Long vectorId = Aop.get(MaxKbEmbeddingService.class).getVectorId(query_text, modelName);
    //String ids = datasetIds.stream().map(id -> "?").collect(Collectors.joining(", "));

    //    sql = SqlTemplates.get("kb.hit_test_by_dataset_ids_with_max_kb_embedding_cache");
    //    sql = sql.replace("#(in_list)", ids);
    //
    //    int paramSize = 3 + datasetIds.size();
    //    Object[] params = new Object[paramSize];
    //
    //    params[0] = vectorId;
    //    for (int i = 0; i < datasetIds.size(); i++) {
    //      params[i + 1] = datasetIds.get(0);
    //    }
    //    params[paramSize - 2] = similarity;
    //    params[paramSize - 1] = top_number;
    //    log.info("sql:{},params:{}", sql, params);
    //    List<Row> records = Db.find(sql, params);

    sql = SqlTemplates.get("kb.search_sentense_related_paragraph__with_dataset_ids");
    Long[] array = datasetIds.toArray(new Long[0]);
    List<Row> records = Db.find(sql, vectorId, array, similarity, top_number);

    log.info("search_paragraph:{},{},{},{},{}", vectorId, Arrays.toString(array), similarity, top_number, records.size());

    List<Kv> kvs = RowUtils.recordsToKv(records, false);
    return ResultVo.ok(kvs);
  }

}

结论

通过引入分层向量与检索的方法，本项目显著提升了 RAG 系统在处理复杂汇总类问题时的准确性和效率。通过将文档拆分为段落级别和句子级别的片段，并进行向量化处理，系统能够更精准地检索相关信息，减少上下文冗余，提高回答质量。未来，随着系统的不断优化和功能的扩展，RAG 系统将能够更好地服务于各类用户，满足其在信息检索和问答方面的多样化需求。