段落分块与检索
本章句子拆分和早期 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. 检索流程
分层向量与检索的具体流程如下:
问题向量化
- 用户输入问题后,首先将其转换为向量表示,即问题向量。
句子级检索
- 使用问题向量在
moss_kb_sentence表中进行相似性检索,筛选出最相关的句子。
- 使用问题向量在
片段获取
- 根据检索出的句子信息,提取其对应的
paragraph_id,并从paragraph表中查询完整的片段(Paragraph)。
- 根据检索出的句子信息,提取其对应的
上下文构建
- 从检索得到的片段中选取10 个最相关的片段,作为大模型推理时的上下文。
大模型推理
- 结合检索到的片段和设计的提示词(Prompt),将构建的上下文输入大模型,进行推理并生成最终答案。
文档拆分代码
在本系统的设计中,文档拆分分为两步:
用户生成文档
- 使用大模型对文档进行识别,对识别后的文档进行段落拆分,并将拆分后的段落返回给前端。
段落修改与提交
- 用户可以对拆分后的段落进行修改,修改后提交到后端。
句子级拆分与存储
- 通过
MossKbParagraphSplitService的batch方法解析前端上传的段落,再拆分成句子级别的向量并存储入库。随后对段落进行汇总并存储入库。
- 通过
以下是相关的代码实现:
MossKbSentenceService 类
该类负责将段落摘要转换为句子级别的记录,并进行向量化处理,最终存储到数据库中。
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 类
该类负责处理用户上传的段落,将其拆分为句子级别的片段,并进行向量化处理和存储。
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 查询:
--# 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 代码实现:
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 格式的提示词,以便于大模型进行推理。
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("&", "&").replace("<", "<").replace(">", ">").replace("\"", """).replace("'", "'");
}
}
代码说明:
分组处理:将检索到的句子按所属文档进行分组,确保每个文档对应一个记录。
XML 构建:为每个分组构建
<record>元素,包含<source>,<document_name>, 和<contents>三个子元素。XML 转义:通过
escapeXml方法对内容进行转义,确保生成的 XML 格式正确,避免特殊字符导致的解析错误。
MossKbParagraphXMLGeneratorTest 类
该类为 MossKbParagraphXMLGenerator 的测试类,用于验证 XML 生成的正确性。
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
<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());
修改 应用的 命中测试
修改后的代码如下
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 系统将能够更好地服务于各类用户,满足其在信息检索和问答方面的多样化需求。
