Skip to content

段落分块与检索 ​

本章句子拆分和早期 SQL 是可选算法示例。当前应用主检索链路按知识库模型检索段落,向量为 1024 维,见12。不要直接将旧维度 SQL 作为当前迁移脚本。

分层向量与检索 ​

功能背景 ​

在 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. 句子级检索

    • 使用问题向量在 moss_kb_sentence 表中进行相似性检索,筛选出最相关的句子。
  3. 片段获取

    • 根据检索出的句子信息,提取其对应的 paragraph_id,并从 paragraph 表中查询完整的片段(Paragraph)。
  4. 上下文构建

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

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

文档拆分代码 ​

在本系统的设计中,文档拆分分为两步:

  1. 用户生成文档

    • 使用大模型对文档进行识别,对识别后的文档进行段落拆分,并将拆分后的段落返回给前端。
  2. 段落修改与提交

    • 用户可以对拆分后的段落进行修改,修改后提交到后端。
  3. 句子级拆分与存储

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

以下是相关的代码实现:

MossKbSentenceService 类 ​

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

java
package nexus.io.mosskb.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 nexus.io.db.activerecord.Db;
import nexus.io.db.activerecord.Row;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.constant.TableNames;
import nexus.io.mosskb.model.MossKbSentence;
import nexus.io.mosskb.utils.ExecutorServiceUtils;
import nexus.io.tio.utils.crypto.Md5Utils;
import nexus.io.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 MossKbSentenceService {
  MossKbEmbeddingService mossKbEmbeddingService = Aop.get(MossKbEmbeddingService.class);
  CompletionService<MossKbSentence> completionServiceSentence = new ExecutorCompletionService<>(ExecutorServiceUtils.getExecutorService());

  CompletionService<Row> completionServiceRow = new ExecutorCompletionService<>(ExecutorServiceUtils.getExecutorService());
  MossKbParagraphSummaryService mossKbParagraphSummaryService = Aop.get(MossKbParagraphSummaryService.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<MossKbSentence> 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();

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

        sentences.add(mossKbSentence);
      }

    }

    List<Future<Row>> futures = new ArrayList<>(sentences.size());
    for (MossKbSentence sentence : sentences) {
      futures.add(completionServiceRow.submit(() -> {
        PGobject vector = mossKbEmbeddingService.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.moss_kb_sentence, Row.by("document_id", documentIdFinal));
      Db.batchSave(TableNames.moss_kb_sentence, sentenceRows, 2000);
      return true;
    });
    return transactionSuccess;
  }
}

代码说明:

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

    • 生成摘要:对每个段落内容调用 MossKbParagraphSummaryService 生成摘要。
    • 生成向量:对生成的摘要内容调用 MossKbEmbeddingService 进行向量化处理。
  • 事务管理:在数据库操作中使用事务,确保数据的一致性和完整性。在保存新的句子记录前,删除已有的相关记录,避免数据重复。

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

MossKbParagraphSplitService 类 ​

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

java
package nexus.io.mosskb.service.kb;

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

import com.jfinal.kit.Kv;
import nexus.io.db.TableInput;
import nexus.io.db.TableResult;
import nexus.io.db.activerecord.Db;
import nexus.io.db.activerecord.Row;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.constant.TableNames;
import nexus.io.mosskb.vo.ParagraphBatchVo;
import nexus.io.mosskb.vo.Paragraph;
import nexus.io.model.result.ResultVo;
import nexus.io.table.services.ApiTable;
import nexus.io.tio.utils.crypto.Md5Utils;
import nexus.io.tio.utils.hutool.FilenameUtils;
import nexus.io.tio.utils.snowflake.SnowflakeIdUtils;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class MossKbParagraphSplitService {

  MossKbEmbeddingService mossKbEmbeddingService = Aop.get(MossKbEmbeddingService.class);
  MossKbSentenceService mossKbSentenceService = Aop.get(MossKbSentenceService.class);

  public ResultVo batch(Long userId, Long dataset_id, List<ParagraphBatchVo> 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.moss_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.moss_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.moss_kb_document);

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

    for (ParagraphBatchVo 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.moss_kb_document, record);
        Kv kv = record.toKv();
        kvs.add(kv);
      } else {
        Row existingRecord = Db.findById(TableNames.moss_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(MossKbSentenceService.class).summaryToSentenceAndSave(dataset_id, modelName, paragraphRecords, documentIdFinal);

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

      transactionSuccess = mossKbSentenceService.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.moss_kb_paragraph, Row.by("document_id", documentIdFinal));
      Db.batchSave(TableNames.moss_kb_paragraph, batchRecord, 2000);
      return true;
    });
  }
}

代码说明:

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

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

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

向量检索 ​

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

检索 SQL 代码 ​

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

sql
--# kb.search_sentense_related_paragraph__with_dataset_ids
SELECT DISTINCT
  sub.id,
  sub.paragraph_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.document_type,
  sub.document_url,
  sub.similarity,
  sub.similarity AS comprehensive_score
FROM (
  SELECT
    d.name AS document_name,
    d.type AS document_type,
    d.url AS document_url,
    ds.name AS dataset_name,
    s.id,
    s.paragraph_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
    moss_kb_sentence s
  JOIN
    moss_kb_embedding_cache c ON c.id = ?
  JOIN
    moss_kb_dataset ds ON s.dataset_id = ds.id
  JOIN
    moss_kb_document d ON s.document_id = d.id
  LEFT JOIN
    moss_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 代码实现:

java
package nexus.io.mosskb.service.kb;

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

import nexus.io.db.activerecord.Db;
import nexus.io.db.activerecord.Row;
import nexus.io.jfinal.aop.Aop;
import nexus.io.mosskb.vo.MossKbRetrieveResult;
import nexus.io.mosskb.vo.ParagraphSearchResultVo;
import nexus.io.openai.consts.OpenAiModels;
import nexus.io.template.SqlTemplates;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class MossKbParagraphRetrieveService {

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

  private List<ParagraphSearchResultVo> search0(Long[] datasetIdArray, Float similarity, Integer top_n, String quesiton) {
    Long vectorId = Aop.get(MossKbEmbeddingService.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 MossKbRetrieveResult searchV1(Long[] datasetIdArray, Float similarity, Integer top_n, String quesiton) {
    MossKbRetrieveResult mossKbSearchStep = new MossKbRetrieveResult();
    long start = System.currentTimeMillis();
    List<ParagraphSearchResultVo> results = searchV10(datasetIdArray, similarity, top_n, quesiton);
    long end = System.currentTimeMillis();
    mossKbSearchStep.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 mossKbSearchStep;
  }

  public List<ParagraphSearchResultVo> searchV10(Long[] datasetIdArray, Float similarity, Integer top_n,
      //
      String quesiton) {
    Long vectorId = Aop.get(MossKbEmbeddingService.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;
  }
}

代码说明:

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

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

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

拼接提示词 ​

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

MossKbParagraphXMLGenerator 类 ​

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

java
package nexus.io.mosskb.service.kb;

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

import nexus.io.mosskb.vo.ParagraphSearchResultVo;

public class MossKbParagraphXMLGenerator {

  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(MossKbParagraphXMLGenerator::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 格式正确,避免特殊字符导致的解析错误。

MossKbParagraphXMLGeneratorTest 类 ​

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

java
package nexus.io.mosskb.service.kb;

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

import org.junit.Test;

import nexus.io.mosskb.vo.ParagraphSearchResultVo;

public class MossKbParagraphXMLGeneratorTest {

  @Test
  public void test() {
    List<ParagraphSearchResultVo> records = getRecords(); // 假设此方法获取记录
    String xmlData = MossKbParagraphXMLGenerator.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;
  }
}

代码说明:

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

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

示例生成的 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> 元素包含了源编号、文档名称以及对应的内容,便于大模型理解和处理。

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

修改 知识库的 命中测试 ​

在之前的 知识库的 命中测试 中 我们是对 moss_kb_paragraph 表进行向量相对度检索,由于我们采用了分段检索,现在我们进行对 moss_kb_sentence 进行检索 修改 nexus.io.mosskb.service.kb.MossKbDatasetHitTestService.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());

修改 应用的 命中测试 ​

修改后的代码如下

java
package nexus.io.mosskb.service.kb;

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

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

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class MossKbApplicationHitTestService {

  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.moss_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.moss_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(MossKbEmbeddingService.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.moss_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.moss_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(MossKbEmbeddingService.class).getVectorId(query_text, modelName);
    //String ids = datasetIds.stream().map(id -> "?").collect(Collectors.joining(", "));

    //    sql = SqlTemplates.get("kb.hit_test_by_dataset_ids_with_moss_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 系统将能够更好地服务于各类用户,满足其在信息检索和问答方面的多样化需求。

MIT Licensed