mirror of
https://gitcode.com/ageerle/ruoyi-ai.git
synced 2026-04-21 07:43:40 +00:00
Merge pull request #292 from yangzhen233/feature/rerank-model
Feature/rerank model
This commit is contained in:
46
docs/script/sql/update/updat-0420.sql
Normal file
46
docs/script/sql/update/updat-0420.sql
Normal file
@@ -0,0 +1,46 @@
|
||||
/*
|
||||
Navicat Premium Dump SQL
|
||||
|
||||
Source Server : localhost-mysql
|
||||
Source Server Type : MySQL
|
||||
Source Server Version : 80045 (8.0.45)
|
||||
Source Host : localhost:3306
|
||||
Source Schema : ruoyi-ai
|
||||
|
||||
Target Server Type : MySQL
|
||||
Target Server Version : 80045 (8.0.45)
|
||||
File Encoding : 65001
|
||||
|
||||
Date: 20/04/2026 15:30:00
|
||||
*/
|
||||
|
||||
SET NAMES utf8mb4;
|
||||
SET FOREIGN_KEY_CHECKS = 0;
|
||||
|
||||
-- ----------------------------
|
||||
-- 新增:重排序模型(chat_model)
|
||||
-- ----------------------------
|
||||
INSERT INTO `chat_model`
|
||||
(id, category, model_name, provider_code, model_describe, model_dimension, model_show, api_host, api_key, create_dept, create_by, create_time, update_by, update_time, remark, tenant_id)
|
||||
VALUES(2045071617578237953, 'rerank', 'rerank', 'zhipu', '智谱重排序', NULL, 'Y', 'https://open.bigmodel.cn', 'e9xx', 103, 1, '2026-04-17 17:27:24', 1, '2026-04-20 15:21:48', '智谱重排序', 0);
|
||||
|
||||
INSERT INTO `chat_model`
|
||||
(id, category, model_name, provider_code, model_describe, model_dimension, model_show, api_host, api_key, create_dept, create_by, create_time, update_by, update_time, remark, tenant_id)
|
||||
VALUES(2046119803482902530, 'rerank', 'qwen3-rerank', 'qianwen', '千问3重排序', NULL, NULL, 'https://dashscope.aliyuncs.com', 'sk-xx', 103, 1, '2026-04-20 14:52:31', 1, '2026-04-20 15:03:13', '千问3文本重排序', 0);
|
||||
|
||||
-- ----------------------------
|
||||
-- 新增:字典类型 - 重排序模型分类
|
||||
-- ----------------------------
|
||||
INSERT INTO `sys_dict_data`
|
||||
(dict_code, tenant_id, dict_sort, dict_label, dict_value, dict_type, css_class, list_class, is_default, create_dept, create_by, create_time, update_by, update_time, remark)
|
||||
VALUES(2045070879435259905, '000000', 4, '重排序', 'rerank', 'chat_model_category', NULL, '#000000', 'N', 103, 1, '2026-04-17 17:24:28', 1, '2026-04-19 01:02:20', '重排序模型');
|
||||
|
||||
-- ----------------------------
|
||||
-- 修改表:knowledge_info 增加重排序相关字段
|
||||
-- ----------------------------
|
||||
ALTER TABLE `knowledge_info` ADD COLUMN `enable_rerank` tinyint DEFAULT 0 NULL COMMENT '是否启用重排序(0否 1是)';
|
||||
ALTER TABLE `knowledge_info` ADD COLUMN `rerank_score_threshold` double NULL COMMENT '重排序相关性分数阈值';
|
||||
ALTER TABLE `knowledge_info` ADD COLUMN `rerank_top_n` int NULL COMMENT '重排序后返回的文档数量';
|
||||
ALTER TABLE `knowledge_info` ADD COLUMN `rerank_model` varchar(100) NULL COMMENT '重排序模型名称';
|
||||
|
||||
SET FOREIGN_KEY_CHECKS = 1;
|
||||
@@ -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;
|
||||
|
||||
|
||||
|
||||
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
package org.ruoyi.domain.dto.request;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonInclude;
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 阿里百炼重排序请求DTO(OpenAI兼容格式)
|
||||
*
|
||||
* @author yang
|
||||
* @date 2026-04-20
|
||||
*/
|
||||
@JsonInclude(JsonInclude.Include.NON_NULL)
|
||||
public record AliBaiLianRerankRequest(
|
||||
String model,
|
||||
List<String> documents,
|
||||
String query,
|
||||
@JsonProperty("top_n")
|
||||
Integer topN,
|
||||
String instruct,
|
||||
@JsonProperty("return_documents")
|
||||
Boolean returnDocuments
|
||||
) {
|
||||
private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
|
||||
|
||||
/**
|
||||
* 创建文本重排序请求
|
||||
*/
|
||||
public static AliBaiLianRerankRequest create(String modelName, String query,
|
||||
List<String> documents, Integer topN,
|
||||
Boolean returnDocuments) {
|
||||
return new AliBaiLianRerankRequest(
|
||||
modelName,
|
||||
documents,
|
||||
query,
|
||||
topN != null ? topN : documents.size(),
|
||||
null,
|
||||
returnDocuments != null ? returnDocuments : true
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 转换为JSON字符串
|
||||
*/
|
||||
public String toJson() {
|
||||
try {
|
||||
return OBJECT_MAPPER.writeValueAsString(this);
|
||||
} catch (JsonProcessingException e) {
|
||||
throw new IllegalArgumentException("序列化阿里百炼重排序请求失败", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
package org.ruoyi.domain.dto.response;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import org.ruoyi.domain.bo.rerank.RerankResult;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
/**
|
||||
* 阿里百炼重排序响应DTO(OpenAI兼容格式)
|
||||
*
|
||||
* @author yang
|
||||
* @date 2026-04-20
|
||||
*/
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
public record AliBaiLianRerankResponse(
|
||||
String id,
|
||||
String object,
|
||||
List<ResultItem> results,
|
||||
UsageInfo usage
|
||||
) {
|
||||
/**
|
||||
* 单个重排序结果项
|
||||
*/
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
public record ResultItem(
|
||||
Integer index,
|
||||
@JsonProperty("relevance_score")
|
||||
Double relevanceScore,
|
||||
Object document
|
||||
) {
|
||||
/**
|
||||
* 获取文档文本内容
|
||||
*/
|
||||
public String getDocumentText() {
|
||||
if (document == null) return null;
|
||||
if (document instanceof String) return (String) document;
|
||||
if (document instanceof Map) {
|
||||
Object text = ((Map<?, ?>) document).get("text");
|
||||
return text != null ? text.toString() : null;
|
||||
}
|
||||
return document.toString();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Token使用信息
|
||||
*/
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
public record UsageInfo(
|
||||
@JsonProperty("total_tokens")
|
||||
Integer totalTokens,
|
||||
@JsonProperty("prompt_tokens")
|
||||
Integer promptTokens
|
||||
) {}
|
||||
|
||||
/**
|
||||
* 转换为通用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.getDocumentText())
|
||||
.build())
|
||||
.collect(Collectors.toList());
|
||||
|
||||
return RerankResult.builder()
|
||||
.documents(documents)
|
||||
.totalDocuments(totalDocs)
|
||||
.durationMs(durationMs)
|
||||
.build();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package org.ruoyi.domain.dto.response;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
|
||||
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
|
||||
*/
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
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使用信息
|
||||
*/
|
||||
@JsonIgnoreProperties(ignoreUnknown = true)
|
||||
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();
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
/**
|
||||
* 备注
|
||||
*/
|
||||
|
||||
@@ -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;
|
||||
|
||||
/**
|
||||
* 备注
|
||||
*/
|
||||
|
||||
@@ -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 -> zhipuRerank,jina -> 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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
package org.ruoyi.service.rerank.impl;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
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.AliBaiLianRerankRequest;
|
||||
import org.ruoyi.domain.dto.response.AliBaiLianRerankResponse;
|
||||
import org.ruoyi.service.rerank.RerankModelService;
|
||||
import org.springframework.stereotype.Component;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
|
||||
/**
|
||||
* 阿里百炼重排序模型实现
|
||||
* 参考设计模式:AliBaiLianMultiEmbeddingProvider
|
||||
*
|
||||
* @author yang
|
||||
* @date 2026-04-20
|
||||
*/
|
||||
@Slf4j
|
||||
@Component("qianwenRerank")
|
||||
public class AliBaiLianRerankModelService implements RerankModelService {
|
||||
|
||||
private final OkHttpClient okHttpClient;
|
||||
private final ObjectMapper objectMapper = new ObjectMapper();
|
||||
private ChatModelVo chatModelVo;
|
||||
|
||||
public AliBaiLianRerankModelService() {
|
||||
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 {
|
||||
// 构建请求
|
||||
AliBaiLianRerankRequest request = buildRequest(rerankRequest);
|
||||
AliBaiLianRerankResponse 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 AliBaiLianRerankRequest buildRequest(RerankRequest rerankRequest) {
|
||||
return AliBaiLianRerankRequest.create(
|
||||
chatModelVo.getModelName(),
|
||||
rerankRequest.getQuery(),
|
||||
rerankRequest.getDocuments(),
|
||||
rerankRequest.getTopN(),
|
||||
rerankRequest.getReturnDocuments()
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 执行HTTP请求并解析响应
|
||||
*/
|
||||
private AliBaiLianRerankResponse executeRequest(AliBaiLianRerankRequest request) throws IOException {
|
||||
String jsonBody = request.toJson();
|
||||
RequestBody body = RequestBody.create(jsonBody, MediaType.get("application/json"));
|
||||
|
||||
// 阿里百炼重排序 OpenAI兼容端点
|
||||
String url = chatModelVo.getApiHost() + "/compatible-api/v1/reranks";
|
||||
Request httpRequest = new Request.Builder()
|
||||
.url(url)
|
||||
.addHeader("Authorization", "Bearer " + chatModelVo.getApiKey())
|
||||
.addHeader("Content-Type", "application/json")
|
||||
.post(body)
|
||||
.build();
|
||||
|
||||
try (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 AliBaiLianRerankResponse parseResponse(String responseBody) throws IOException {
|
||||
return objectMapper.readValue(responseBody, AliBaiLianRerankResponse.class);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
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() + "/api/paas/v4/rerank";
|
||||
Request httpRequest = new Request.Builder()
|
||||
.url(url)
|
||||
.addHeader("Authorization", token)
|
||||
.post(body)
|
||||
.build();
|
||||
|
||||
try (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);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package org.ruoyi.service.rerank.impl;
|
||||
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
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.Arrays;
|
||||
import java.util.List;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
/**
|
||||
* 阿里百炼重排序模型测试类
|
||||
* 运行前请设置环境变量 DASHSCOPE_API_KEY 或直接修改 apiKey
|
||||
*/
|
||||
class AliBaiLianRerankModelServiceTest {
|
||||
|
||||
private AliBaiLianRerankModelService service;
|
||||
|
||||
// 请替换为你的 API Key
|
||||
private static final String API_KEY = System.getenv("DASHSCOPE_API_KEY");
|
||||
private static final String API_HOST = "https://dashscope.aliyuncs.com";
|
||||
private static final String MODEL_NAME = "qwen3-rerank";
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
service = new AliBaiLianRerankModelService();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testConfigure() {
|
||||
ChatModelVo config = createConfig();
|
||||
service.configure(config);
|
||||
assertNotNull(service);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRerank() {
|
||||
// 跳过测试如果没有配置 API Key
|
||||
if (API_KEY == null || API_KEY.isEmpty()) {
|
||||
System.out.println("跳过测试: 未设置环境变量 DASHSCOPE_API_KEY");
|
||||
return;
|
||||
}
|
||||
|
||||
ChatModelVo config = createConfig();
|
||||
service.configure(config);
|
||||
|
||||
List<String> documents = Arrays.asList(
|
||||
"文本排序模型广泛用于搜索引擎和推荐系统中,它们根据文本相关性对候选文本进行排序",
|
||||
"量子计算是计算科学的一个前沿领域",
|
||||
"预训练语言模型的发展给文本排序模型带来了新的进展"
|
||||
);
|
||||
|
||||
RerankRequest request = RerankRequest.builder()
|
||||
.query("什么是文本排序模型")
|
||||
.documents(documents)
|
||||
.topN(2)
|
||||
.returnDocuments(true)
|
||||
.build();
|
||||
|
||||
RerankResult result = service.rerank(request);
|
||||
|
||||
System.out.println("=== 重排序结果 ===");
|
||||
System.out.println("总文档数: " + result.getTotalDocuments());
|
||||
System.out.println("耗时: " + result.getDurationMs() + "ms");
|
||||
|
||||
result.getDocuments().forEach(doc -> {
|
||||
System.out.println("索引: " + doc.getIndex() +
|
||||
", 相关性分数: " + doc.getRelevanceScore() +
|
||||
", 文档: " + doc.getDocument());
|
||||
});
|
||||
|
||||
assertNotNull(result);
|
||||
assertNotNull(result.getDocuments());
|
||||
assertFalse(result.getDocuments().isEmpty());
|
||||
assertEquals(2, result.getDocuments().size());
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRerankWithFullDocuments() {
|
||||
if (API_KEY == null || API_KEY.isEmpty()) {
|
||||
System.out.println("跳过测试: 未设置环境变量 DASHSCOPE_API_KEY");
|
||||
return;
|
||||
}
|
||||
|
||||
ChatModelVo config = createConfig();
|
||||
service.configure(config);
|
||||
|
||||
List<String> documents = Arrays.asList(
|
||||
"Java是一种广泛使用的编程语言",
|
||||
"Python是人工智能领域最流行的语言",
|
||||
"Go语言由Google开发,适合并发编程"
|
||||
);
|
||||
|
||||
RerankRequest request = RerankRequest.builder()
|
||||
.query("哪种语言适合AI开发")
|
||||
.documents(documents)
|
||||
.build();
|
||||
|
||||
RerankResult result = service.rerank(request);
|
||||
|
||||
System.out.println("=== 重排序结果2 ===");
|
||||
result.getDocuments().forEach(doc -> {
|
||||
System.out.println("索引: " + doc.getIndex() +
|
||||
", 分数: " + doc.getRelevanceScore() +
|
||||
", 文档: " + doc.getDocument());
|
||||
});
|
||||
|
||||
assertNotNull(result);
|
||||
assertEquals(3, result.getDocuments().size());
|
||||
|
||||
// Python相关文档应该排在前面
|
||||
assertEquals(1, result.getDocuments().get(0).getIndex());
|
||||
assertTrue(result.getDocuments().get(0).getRelevanceScore() > 0.5);
|
||||
}
|
||||
|
||||
private ChatModelVo createConfig() {
|
||||
ChatModelVo config = new ChatModelVo();
|
||||
config.setApiHost(API_HOST);
|
||||
config.setApiKey(API_KEY != null ? API_KEY : "test-api-key");
|
||||
config.setModelName(MODEL_NAME);
|
||||
return config;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
package org.ruoyi.service.rerank.impl;
|
||||
|
||||
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.Arrays;
|
||||
import java.util.List;
|
||||
|
||||
/**
|
||||
* 阿里百炼重排序模型测试 - Main方法直接运行
|
||||
* 运行前请设置 API_KEY
|
||||
*/
|
||||
public class AliBaiLianRerankTestMain {
|
||||
|
||||
// 请替换为你的 API Key
|
||||
private static final String API_KEY = "sk-your-api-key-here";
|
||||
private static final String API_HOST = "https://dashscope.aliyuncs.com";
|
||||
private static final String MODEL_NAME = "qwen3-rerank";
|
||||
|
||||
public static void main(String[] args) {
|
||||
AliBaiLianRerankModelService service = new AliBaiLianRerankModelService();
|
||||
|
||||
// 配置
|
||||
ChatModelVo config = new ChatModelVo();
|
||||
config.setApiHost(API_HOST);
|
||||
config.setApiKey(API_KEY);
|
||||
config.setModelName(MODEL_NAME);
|
||||
service.configure(config);
|
||||
|
||||
// 测试数据
|
||||
List<String> documents = Arrays.asList(
|
||||
"文本排序模型广泛用于搜索引擎和推荐系统中,它们根据文本相关性对候选文本进行排序",
|
||||
"量子计算是计算科学的一个前沿领域",
|
||||
"预训练语言模型的发展给文本排序模型带来了新的进展"
|
||||
);
|
||||
|
||||
RerankRequest request = RerankRequest.builder()
|
||||
.query("什么是文本排序模型")
|
||||
.documents(documents)
|
||||
.topN(2)
|
||||
.returnDocuments(true)
|
||||
.build();
|
||||
|
||||
System.out.println("=== 开始测试阿里百炼重排序 ===");
|
||||
System.out.println("API Host: " + API_HOST);
|
||||
System.out.println("Model: " + MODEL_NAME);
|
||||
System.out.println("Query: 什么是文本排序模型");
|
||||
System.out.println();
|
||||
|
||||
try {
|
||||
RerankResult result = service.rerank(request);
|
||||
|
||||
System.out.println("=== 重排序结果 ===");
|
||||
System.out.println("总文档数: " + result.getTotalDocuments());
|
||||
System.out.println("耗时: " + result.getDurationMs() + "ms");
|
||||
System.out.println();
|
||||
|
||||
result.getDocuments().forEach(doc -> {
|
||||
System.out.println("索引: " + doc.getIndex());
|
||||
System.out.println("相关性分数: " + doc.getRelevanceScore());
|
||||
System.out.println("文档: " + doc.getDocument());
|
||||
System.out.println("---");
|
||||
});
|
||||
|
||||
System.out.println("=== 测试成功 ===");
|
||||
|
||||
} catch (Exception e) {
|
||||
System.err.println("=== 测试失败 ===");
|
||||
e.printStackTrace();
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user