添加重排序功能

This commit is contained in:
杨振
2026-04-20 01:02:09 +08:00
parent 2c6ff66830
commit 80ca76ea37
15 changed files with 938 additions and 2 deletions

View File

@@ -77,10 +77,33 @@ public class KnowledgeInfoBo extends BaseEntity {
*/
private String embeddingModel;
/**
* 是否启用重排序0 否 1是
*/
private Integer enableRerank;
/**
* 重排序模型名称
*/
private String rerankModel;
/**
* 重排序后返回的文档数量
*/
private Integer rerankTopN;
/**
* 重排序相关性分数阈值
*/
private Double rerankScoreThreshold;
/**
* 备注
*/
private String remark;
}

View File

@@ -0,0 +1,44 @@
package org.ruoyi.domain.bo.rerank;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import java.util.List;
/**
* 重排序请求参数
*
* @author yang
* @date 2026-04-19
*/
@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public class RerankRequest {
/**
* 查询文本
*/
private String query;
/**
* 候选文档列表
*/
private List<String> documents;
/**
* 返回的文档数量topN
* 如果不指定,默认返回所有文档
*/
private Integer topN;
/**
* 是否返回原始文档内容
* 默认为 true
*/
@Builder.Default
private Boolean returnDocuments = true;
}

View File

@@ -0,0 +1,72 @@
package org.ruoyi.domain.bo.rerank;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import java.util.List;
/**
* 重排序结果
*
* @author yang
* @date 2026-04-19
*/
@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public class RerankResult {
/**
* 重排序后的文档结果列表
*/
private List<RerankDocument> documents;
/**
* 原始请求中的文档总数
*/
private Integer totalDocuments;
/**
* 重排序耗时(毫秒)
*/
private Long durationMs;
/**
* 单个重排序文档结果
*/
@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public static class RerankDocument {
/**
* 文档在原始列表中的索引位置
*/
private Integer index;
/**
* 相关性分数(通常 0-1 之间,越高越相关)
*/
private Double relevanceScore;
/**
* 文档内容
*/
private String document;
}
/**
* 创建空结果
*/
public static RerankResult empty() {
return RerankResult.builder()
.documents(List.of())
.totalDocuments(0)
.durationMs(0L)
.build();
}
}

View File

@@ -51,4 +51,30 @@ public class QueryVectorBo {
*/
private String baseUrl;
// ========== 重排序相关参数 ==========
/**
* 是否启用重排序
* 默认为 false
*/
private Boolean enableRerank = false;
/**
* 重排序模型名称
*/
private String rerankModelName;
/**
* 重排序后返回的文档数量topN
* 如果不指定,默认与 maxResults 相同
*/
private Integer rerankTopN;
/**
* 重排序相关性分数阈值
* 低于此阈值的文档将被过滤
*/
private Double rerankScoreThreshold;
}

View File

@@ -0,0 +1,48 @@
package org.ruoyi.domain.dto.request;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.util.List;
/**
* 智谱AI重排序请求DTO
*
* @author yang
* @date 2026-04-19
*/
public record ZhipuRerankRequest(
String model,
String query,
List<String> documents,
Integer top_n,
Boolean return_documents
) {
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
/**
* 创建智谱重排序请求
*/
public static ZhipuRerankRequest create(String modelName, String query,
List<String> documents, Integer topN,
Boolean returnDocuments) {
return new ZhipuRerankRequest(
modelName,
query,
documents,
topN != null ? topN : documents.size(),
returnDocuments != null ? returnDocuments : true
);
}
/**
* 转换为JSON字符串
*/
public String toJson() {
try {
return OBJECT_MAPPER.writeValueAsString(this);
} catch (JsonProcessingException e) {
throw new IllegalArgumentException("序列化智谱重排序请求失败", e);
}
}
}

View File

@@ -0,0 +1,65 @@
package org.ruoyi.domain.dto.response;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.ruoyi.domain.bo.rerank.RerankResult;
import java.util.List;
import java.util.stream.Collectors;
/**
* 智谱AI重排序响应DTO
*
* @author yang
* @date 2026-04-19
*/
public record ZhipuRerankResponse(
String model,
String object,
List<ResultItem> results,
UsageInfo usage
) {
/**
* 单个重排序结果项
*/
public record ResultItem(
Integer index,
@JsonProperty("relevance_score")
Double relevanceScore,
String document
) {}
/**
* Token使用信息
*/
public record UsageInfo(
@JsonProperty("total_tokens")
Integer totalTokens,
@JsonProperty("input_tokens")
Integer inputTokens,
@JsonProperty("output_tokens")
Integer outputTokens
) {}
/**
* 转换为通用RerankResult
*/
public RerankResult toRerankResult(int totalDocs, long durationMs) {
if (results == null || results.isEmpty()) {
return RerankResult.empty();
}
List<RerankResult.RerankDocument> documents = results.stream()
.map(item -> RerankResult.RerankDocument.builder()
.index(item.index())
.relevanceScore(item.relevanceScore())
.document(item.document())
.build())
.collect(Collectors.toList());
return RerankResult.builder()
.documents(documents)
.totalDocuments(totalDocs)
.durationMs(durationMs)
.build();
}
}

View File

@@ -78,6 +78,26 @@ public class KnowledgeInfo extends BaseEntity {
*/
private String embeddingModel;
/**
* 是否启用重排序0 否 1是
*/
private Integer enableRerank;
/**
* 重排序模型名称
*/
private String rerankModel;
/**
* 重排序后返回的文档数量
*/
private Integer rerankTopN;
/**
* 重排序相关性分数阈值
*/
private Double rerankScoreThreshold;
/**
* 备注
*/

View File

@@ -94,6 +94,30 @@ public class KnowledgeInfoVo implements Serializable {
@ExcelProperty(value = "向量模型")
private String embeddingModel;
/**
* 是否启用重排序0 否 1是
*/
@ExcelProperty(value = "是否启用重排序")
private Integer enableRerank;
/**
* 重排序模型名称
*/
@ExcelProperty(value = "重排序模型")
private String rerankModel;
/**
* 重排序后返回的文档数量
*/
@ExcelProperty(value = "重排序返回数量")
private Integer rerankTopN;
/**
* 重排序相关性分数阈值
*/
@ExcelProperty(value = "重排序分数阈值")
private Double rerankScoreThreshold;
/**
* 备注
*/

View File

@@ -0,0 +1,106 @@
package org.ruoyi.factory;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo;
import org.ruoyi.common.chat.service.chat.IChatModelService;
import org.ruoyi.service.rerank.RerankModelService;
import org.springframework.beans.factory.NoSuchBeanDefinitionException;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Service;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
/**
* 重排序模型工厂服务类
* 参考设计模式EmbeddingModelFactory
* 负责创建和管理重排序模型实例
*
* @author yang
* @date 2026-04-19
*/
@Service
@RequiredArgsConstructor
@Slf4j
public class RerankModelFactory {
private final ApplicationContext applicationContext;
private final IChatModelService chatModelService;
/**
* 模型缓存使用ConcurrentHashMap保证线程安全
*/
private final Map<String, RerankModelService> modelCache = new ConcurrentHashMap<>();
/**
* 创建重排序模型实例
* 如果模型已存在于缓存中,则直接返回;否则创建新的实例
*
* @param rerankModelName 重排序模型名称
*/
public RerankModelService createModel(String rerankModelName) {
return modelCache.computeIfAbsent(rerankModelName, name -> {
ChatModelVo modelConfig = chatModelService.selectModelByName(rerankModelName);
if (modelConfig == null) {
throw new IllegalArgumentException("未找到重排序模型配置name=" + name);
}
return createModelInstance(modelConfig.getProviderCode(), modelConfig);
});
}
/**
* 刷新模型缓存
* 根据给定的模型ID从缓存中移除对应的模型
*
* @param modelId 模型的唯一标识ID
*/
public void refreshModel(Long modelId) {
modelCache.remove(modelId);
}
/**
* 获取所有支持模型工厂的列表
*
* @return 支持的模型工厂名称列表
*/
public List<String> getSupportedFactories() {
return new ArrayList<>(applicationContext.getBeansOfType(RerankModelService.class)
.keySet());
}
/**
* 创建具体的模型实例
* 根据提供的工厂名称和配置信息创建并配置模型实例
*
* @param factory 工厂名称用于标识模型类型providerCode
* @param config 模型配置信息
* @return RerankModelService 配置好的模型实例
* @throws IllegalArgumentException 当无法获取指定的模型实例时抛出
*/
private RerankModelService createModelInstance(String factory, ChatModelVo config) {
try {
// 优先尝试使用 providerCode + "Rerank" 作为 Bean 名称
// 例如zhipu -> zhipuRerankjina -> jinaRerank
String rerankBeanName = factory + "Rerank";
RerankModelService model = applicationContext.getBean(rerankBeanName, RerankModelService.class);
model.configure(config);
log.info("成功创建重排序模型: factory={}, modelName={}", rerankBeanName, config.getModelName());
return model;
} catch (NoSuchBeanDefinitionException e) {
// 如果找不到,尝试使用原始的 providerCode
try {
RerankModelService model = applicationContext.getBean(factory, RerankModelService.class);
model.configure(config);
log.info("成功创建重排序模型: factory={}, modelName={}", factory, config.getModelName());
return model;
} catch (NoSuchBeanDefinitionException ex) {
throw new IllegalArgumentException("获取不到重排序模型: " + factory + "" + factory + "Rerank", ex);
}
}
}
}

View File

@@ -54,6 +54,7 @@ import org.ruoyi.service.chat.AbstractChatService;
import org.ruoyi.service.chat.IChatMessageService;
import org.ruoyi.service.chat.impl.memory.PersistentChatMemoryStore;
import org.ruoyi.service.knowledge.IKnowledgeInfoService;
import org.ruoyi.service.retrieval.KnowledgeRetrievalService;
import org.ruoyi.service.vector.VectorStoreService;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
@@ -89,6 +90,8 @@ public class ChatServiceFacade implements IChatService {
private final VectorStoreService vectorStoreService;
private final KnowledgeRetrievalService knowledgeRetrievalService;
private final SseEmitterManager sseEmitterManager;
private final IChatMessageService chatMessageService;
@@ -452,8 +455,8 @@ public class ChatServiceFacade implements IChatService {
// 构建向量查询参数
QueryVectorBo queryVectorBo = buildQueryVectorBo(chatRequest, knowledgeInfoVo, chatModel);
// 获取向量查询结果(知识库内容作为系统上下文,放在历史消息之后
List<String> nearestList = vectorStoreService.getQueryVector(queryVectorBo);
// 使用知识库检索服务(支持重排序
List<String> nearestList = knowledgeRetrievalService.retrieveTexts(queryVectorBo);
for (String prompt : nearestList) {
// 知识库内容作为系统上下文添加
messages.add(new AiMessage(prompt));
@@ -480,6 +483,13 @@ public class ChatServiceFacade implements IChatService {
queryVectorBo.setVectorModelName(knowledgeInfoVo.getVectorModel());
queryVectorBo.setEmbeddingModelName(knowledgeInfoVo.getEmbeddingModel());
queryVectorBo.setMaxResults(knowledgeInfoVo.getRetrieveLimit());
// 设置重排序参数
queryVectorBo.setEnableRerank(knowledgeInfoVo.getEnableRerank() != null && knowledgeInfoVo.getEnableRerank() == 1);
queryVectorBo.setRerankModelName(knowledgeInfoVo.getRerankModel());
queryVectorBo.setRerankTopN(knowledgeInfoVo.getRerankTopN());
queryVectorBo.setRerankScoreThreshold(knowledgeInfoVo.getRerankScoreThreshold());
return queryVectorBo;
}

View File

@@ -0,0 +1,70 @@
package org.ruoyi.service.rerank;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.model.scoring.ScoringModel;
import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo;
import org.ruoyi.domain.bo.rerank.RerankRequest;
import org.ruoyi.domain.bo.rerank.RerankResult;
import java.util.List;
/**
* 重排序模型服务接口
* 继承 langchain4j 的 ScoringModel 接口
* 参考设计模式BaseEmbedModelService
*
* @author Yzm
* @date 2026-04-19
*/
public interface RerankModelService extends ScoringModel {
/**
* 根据配置信息配置重排序模型
*
* @param config 包含模型配置信息的 ChatModelVo 对象
*/
void configure(ChatModelVo config);
/**
* 执行重排序(批量文档)
* 这是业务层使用的便捷方法
*
* @param rerankRequest 重排序请求,包含查询文本和候选文档列表
* @return 重排序结果,包含排序后的文档和相关性分数
*/
RerankResult rerank(RerankRequest rerankRequest);
/**
* 实现 ScoringModel 接口的 scoreAll 方法
* 将 ScoringModel 的调用转换为重排序调用
*/
@Override
default Response<List<Double>> scoreAll(List<TextSegment> segments, String query) {
// 将 TextSegment 转换为文档字符串列表
List<String> documents = segments.stream()
.map(TextSegment::text)
.toList();
RerankRequest request = RerankRequest.builder()
.query(query)
.documents(documents)
.topN(documents.size())
.returnDocuments(false)
.build();
RerankResult result = rerank(request);
// 提取分数列表,按原始顺序排列
List<Double> scores = new java.util.ArrayList<>(
java.util.Collections.nCopies(documents.size(), 0.0));
for (RerankResult.RerankDocument doc : result.getDocuments()) {
if (doc.getIndex() != null && doc.getIndex() < documents.size()) {
scores.set(doc.getIndex(), doc.getRelevanceScore());
}
}
return Response.from(scores);
}
}

View File

@@ -0,0 +1,107 @@
package org.ruoyi.service.rerank.impl;
import com.fasterxml.jackson.databind.ObjectMapper;
import dev.langchain4j.http.client.*;
import lombok.extern.slf4j.Slf4j;
import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo;
import org.ruoyi.domain.bo.rerank.RerankRequest;
import org.ruoyi.domain.bo.rerank.RerankResult;
import org.ruoyi.service.rerank.RerankModelService;
import org.springframework.stereotype.Component;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
/**
* Jina AI 重排序模型实现
* 参考设计模式OpenAiEmbeddingProvider
* 使用 Jina 官方重排序API
*
* @author yang
* @date 2026-04-19
*/
@Slf4j
@Component("jina")
public class JinaRerankModelService implements RerankModelService {
protected ChatModelVo chatModelVo;
protected HttpClient httpClient;
protected final ObjectMapper objectMapper = new ObjectMapper();
@Override
public void configure(ChatModelVo config) {
this.chatModelVo = config;
HttpClientBuilder httpClientBuilder = HttpClientBuilderLoader.loadHttpClientBuilder();
this.httpClient = httpClientBuilder.build();
}
@Override
public RerankResult rerank(RerankRequest rerankRequest) {
long startTime = System.currentTimeMillis();
try {
// 构建Jina重排序请求体
Map<String, Object> requestBody = new HashMap<>();
requestBody.put("model", chatModelVo.getModelName());
requestBody.put("query", rerankRequest.getQuery());
requestBody.put("documents", rerankRequest.getDocuments());
requestBody.put("top_n", rerankRequest.getTopN() != null ?
rerankRequest.getTopN() : rerankRequest.getDocuments().size());
requestBody.put("return_documents", rerankRequest.getReturnDocuments());
// 构建HTTP请求
HttpRequest httpRequest = HttpRequest.builder()
.url(chatModelVo.getApiHost())
.method(HttpMethod.POST)
.addHeader("Content-Type", "application/json")
.addHeader("Authorization", "Bearer " + chatModelVo.getApiKey())
.body(objectMapper.writeValueAsString(requestBody))
.build();
// 发送请求
SuccessfulHttpResponse httpResponse = httpClient.execute(httpRequest);
// 解析响应
@SuppressWarnings("unchecked")
Map<String, Object> response = objectMapper.readValue(httpResponse.body(), Map.class);
return parseResponse(response, rerankRequest.getDocuments().size(),
System.currentTimeMillis() - startTime);
} catch (Exception e) {
log.error("Jina重排序失败: {}", e.getMessage(), e);
throw new RuntimeException("重排序服务调用失败: " + e.getMessage(), e);
}
}
/**
* 解析Jina API响应
*/
@SuppressWarnings("unchecked")
private RerankResult parseResponse(Map<String, Object> response, int totalDocs, long durationMs) {
List<Map<String, Object>> results = (List<Map<String, Object>>) response.get("results");
if (results == null || results.isEmpty()) {
return RerankResult.empty();
}
List<RerankResult.RerankDocument> documents = new ArrayList<>();
for (Map<String, Object> result : results) {
Integer index = (Integer) result.get("index");
Double score = ((Number) result.get("relevance_score")).doubleValue();
String document = (String) result.get("document");
documents.add(RerankResult.RerankDocument.builder()
.index(index)
.relevanceScore(score)
.document(document)
.build());
}
return RerankResult.builder()
.documents(documents)
.totalDocuments(totalDocs)
.durationMs(durationMs)
.build();
}
}

View File

@@ -0,0 +1,162 @@
package org.ruoyi.service.rerank.impl;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.security.MacAlgorithm;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.ruoyi.common.chat.domain.vo.chat.ChatModelVo;
import org.ruoyi.domain.bo.rerank.RerankRequest;
import org.ruoyi.domain.bo.rerank.RerankResult;
import org.ruoyi.domain.dto.request.ZhipuRerankRequest;
import org.ruoyi.domain.dto.response.ZhipuRerankResponse;
import org.ruoyi.service.rerank.RerankModelService;
import org.springframework.stereotype.Component;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.lang.reflect.Constructor;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.TimeUnit;
/**
* 智谱AI 重排序模型实现
* 参考设计模式AliBaiLianMultiEmbeddingProvider
*
* @author yang
* @date 2026-04-19
*/
@Slf4j
@Component("zhipuRerank")
public class ZhiPuRerankModelService implements RerankModelService {
private final OkHttpClient okHttpClient;
private final ObjectMapper objectMapper = new ObjectMapper();
private ChatModelVo chatModelVo;
public ZhiPuRerankModelService() {
this.okHttpClient = new OkHttpClient.Builder()
.connectTimeout(30, TimeUnit.SECONDS)
.readTimeout(60, TimeUnit.SECONDS)
.writeTimeout(30, TimeUnit.SECONDS)
.build();
}
@Override
public void configure(ChatModelVo config) {
this.chatModelVo = config;
}
@Override
public RerankResult rerank(RerankRequest rerankRequest) {
long startTime = System.currentTimeMillis();
try {
// 构建请求
ZhipuRerankRequest request = buildRequest(rerankRequest);
ZhipuRerankResponse response = executeRequest(request);
return response.toRerankResult(
rerankRequest.getDocuments().size(),
System.currentTimeMillis() - startTime
);
} catch (Exception e) {
log.error("智谱重排序失败: {}", e.getMessage(), e);
throw new RuntimeException("重排序服务调用失败: " + e.getMessage(), e);
}
}
/**
* 构建请求对象
*/
private ZhipuRerankRequest buildRequest(RerankRequest rerankRequest) {
return ZhipuRerankRequest.create(
chatModelVo.getModelName(),
rerankRequest.getQuery(),
rerankRequest.getDocuments(),
rerankRequest.getTopN(),
rerankRequest.getReturnDocuments()
);
}
/**
* 执行HTTP请求并解析响应
*/
private ZhipuRerankResponse executeRequest(ZhipuRerankRequest request) throws IOException {
String jsonBody = request.toJson();
RequestBody body = RequestBody.create(jsonBody, MediaType.get("application/json"));
// 生成智谱认证Token
String token = generateToken(chatModelVo.getApiKey());
String url = chatModelVo.getApiHost() + "/" + chatModelVo.getModelName();
Request httpRequest = new Request.Builder()
.url(url)
.addHeader("Authorization", token)
.post(body)
.build();
try (okhttp3.Response response = okHttpClient.newCall(httpRequest).execute()) {
if (!response.isSuccessful()) {
String err = response.body() != null ? response.body().string() : "无错误信息";
throw new IllegalArgumentException("智谱API调用失败: " + response.code() + " - " + err);
}
ResponseBody responseBody = response.body();
if (responseBody == null) {
throw new IllegalArgumentException("响应体为空");
}
return parseResponse(responseBody.string());
}
}
/**
* 解析响应
*/
private ZhipuRerankResponse parseResponse(String responseBody) throws IOException {
return objectMapper.readValue(responseBody, ZhipuRerankResponse.class);
}
/**
* 生成智谱JWT Token
*/
private String generateToken(String apiKey) {
try {
String[] apiKeyParts = apiKey.split("\\.");
String keyId = apiKeyParts[0];
String secret = apiKeyParts[1];
long expireMillis = 1000L * 60 * 30; // 30分钟
java.util.Map<String, Object> payload = new java.util.HashMap<>();
payload.put("api_key", keyId);
payload.put("exp", System.currentTimeMillis() + expireMillis);
payload.put("timestamp", System.currentTimeMillis());
// 使用反射创建 MacAlgorithm兼容不同版本的 jjwt
MacAlgorithm macAlgorithm;
try {
Class<?> c = Class.forName("io.jsonwebtoken.impl.security.DefaultMacAlgorithm");
Constructor<?> ctor = c.getDeclaredConstructor(String.class, String.class, int.class);
ctor.setAccessible(true);
macAlgorithm = (MacAlgorithm) ctor.newInstance("HS256", "HmacSHA256", 128);
} catch (Exception e) {
macAlgorithm = Jwts.SIG.HS256;
}
String token = Jwts.builder()
.header()
.add("alg", "HS256")
.add("sign_type", "SIGN")
.and()
.content(objectMapper.writeValueAsString(payload))
.signWith(new SecretKeySpec(secret.getBytes(StandardCharsets.UTF_8), "HmacSHA256"), macAlgorithm)
.compact();
return "Bearer " + token;
} catch (Exception e) {
throw new RuntimeException("生成智谱Token失败: " + e.getMessage(), e);
}
}
}

View File

@@ -0,0 +1,24 @@
package org.ruoyi.service.retrieval;
import org.ruoyi.domain.bo.vector.QueryVectorBo;
import java.util.List;
/**
* 知识库检索服务接口
* 整合粗召回(向量检索)和重排序流程
*
* @author yang
* @date 2026-04-19
*/
public interface KnowledgeRetrievalService {
/**
* 执行知识库检索,返回文本内容
* 流程:向量粗召回 -> 重排序(可选) -> 返回结果
*
* @param queryVectorBo 查询参数
* @return 文本内容列表
*/
List<String> retrieveTexts(QueryVectorBo queryVectorBo);
}

View File

@@ -0,0 +1,135 @@
package org.ruoyi.service.retrieval.impl;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.ruoyi.domain.bo.rerank.RerankRequest;
import org.ruoyi.domain.bo.rerank.RerankResult;
import org.ruoyi.domain.bo.vector.QueryVectorBo;
import org.ruoyi.factory.RerankModelFactory;
import org.ruoyi.service.rerank.RerankModelService;
import org.ruoyi.service.retrieval.KnowledgeRetrievalService;
import org.ruoyi.service.vector.VectorStoreService;
import org.springframework.stereotype.Service;
import java.util.ArrayList;
import java.util.List;
/**
* 知识库检索服务实现
* 整合粗召回(向量检索)和重排序流程
*
* @author yang
* @date 2026-04-19
*/
@Slf4j
@Service
@RequiredArgsConstructor
public class KnowledgeRetrievalServiceImpl implements KnowledgeRetrievalService {
private final VectorStoreService vectorStoreService;
private final RerankModelFactory rerankModelFactory;
/**
* 粗召回默认扩大倍数
* 如果启用重排序,粗召回会获取更多结果供重排序筛选
*/
private static final int RERANK_EXPANSION_FACTOR = 3;
@Override
public List<String> retrieveTexts(QueryVectorBo queryVectorBo) {
log.info("开始知识库检索, kid={}, query={}", queryVectorBo.getKid(), queryVectorBo.getQuery());
// 1. 粗召回阶段 - 向量检索
List<String> coarseResults = coarseRetrieval(queryVectorBo);
log.debug("粗召回返回 {} 条结果", coarseResults.size());
if (coarseResults.isEmpty()) {
return coarseResults;
}
// 2. 重排序阶段(可选)
if (Boolean.TRUE.equals(queryVectorBo.getEnableRerank()) &&
queryVectorBo.getRerankModelName() != null) {
return rerank(queryVectorBo, coarseResults);
}
return coarseResults;
}
/**
* 粗召回阶段 - 向量检索
*/
private List<String> coarseRetrieval(QueryVectorBo queryVectorBo) {
// 如果启用重排序,扩大粗召回数量
int originalMaxResults = queryVectorBo.getMaxResults();
int expandedResults = originalMaxResults;
if (Boolean.TRUE.equals(queryVectorBo.getEnableRerank()) &&
queryVectorBo.getRerankModelName() != null) {
expandedResults = originalMaxResults * RERANK_EXPANSION_FACTOR;
log.debug("启用重排序,粗召回数量从 {} 扩大到 {}", originalMaxResults, expandedResults);
}
// 临时修改查询数量
queryVectorBo.setMaxResults(expandedResults);
try {
return vectorStoreService.getQueryVector(queryVectorBo);
} finally {
// 恢复原始值
queryVectorBo.setMaxResults(originalMaxResults);
}
}
/**
* 重排序阶段
*/
private List<String> rerank(QueryVectorBo queryVectorBo, List<String> coarseResults) {
long startTime = System.currentTimeMillis();
try {
// 1. 通过工厂获取重排序模型
RerankModelService rerankModel = rerankModelFactory.createModel(queryVectorBo.getRerankModelName());
// 2. 构建重排序请求
int topN = queryVectorBo.getRerankTopN() != null ?
queryVectorBo.getRerankTopN() : queryVectorBo.getMaxResults();
RerankRequest rerankRequest = RerankRequest.builder()
.query(queryVectorBo.getQuery())
.documents(coarseResults)
.topN(topN)
.returnDocuments(true)
.build();
log.info("执行重排序, model={}, documents={}, topN={}",
queryVectorBo.getRerankModelName(), coarseResults.size(), topN);
// 3. 执行重排序
RerankResult rerankResult = rerankModel.rerank(rerankRequest);
// 4. 转换重排序结果
List<String> finalResults = new ArrayList<>();
for (RerankResult.RerankDocument doc : rerankResult.getDocuments()) {
// 应用分数阈值过滤
if (queryVectorBo.getRerankScoreThreshold() != null &&
doc.getRelevanceScore() < queryVectorBo.getRerankScoreThreshold()) {
continue;
}
if (doc.getDocument() != null) {
finalResults.add(doc.getDocument());
}
}
long duration = System.currentTimeMillis() - startTime;
log.info("重排序完成, 返回 {} 条结果, 耗时 {}ms", finalResults.size(), duration);
return finalResults;
} catch (Exception e) {
log.error("重排序失败: {}", e.getMessage(), e);
// 重排序失败时返回原始粗召回结果(截取到期望数量)
int limit = Math.min(queryVectorBo.getMaxResults(), coarseResults.size());
return new ArrayList<>(coarseResults.subList(0, limit));
}
}
}