RAG/src/main/java/com/wok/supportbot/config/EmbeddingConfigFixer.java

package com.wok.supportbot.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.ApplicationListener;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

import java.util.List;
import java.util.Map;

/**
 * Embedding 配置检查器
 * 在应用启动完成后，检查 EMBEDDING 配置的合理性并记录日志。
 * 注意：不再强制修正非 DashScope 配置，尊重用户在 DB 中配置的提供商和模型。
 * 新增：检测 EmbeddingModel 返回的实际维度与 yml 配置的向量维度是否一致。
 */
@Component
@Slf4j
public class EmbeddingConfigFixer implements ApplicationListener<ApplicationReadyEvent> {

    @Autowired
    private JdbcTemplate jdbcTemplate;

    @Autowired
    private EmbeddingModelFactory embeddingModelFactory;

    @Value("${spring.ai.dashscope.api-key:}")
    private String dashscopeApiKey;

    @Value("${knowledge.vector.dimension:1024}")
    private int fallbackDimension;

    @Override
    public void onApplicationEvent(ApplicationReadyEvent event) {
        try {
            checkEmbeddingConfig();
            checkDimensionConsistency();
        } catch (Exception e) {
            log.warn("Embedding 配置检查异常（不影响启动）: {}", e.getMessage());
        }
    }

    private void checkEmbeddingConfig() {
        // 打印当前 EMBEDDING 所有配置
        List<Map<String, Object>> allConfigs = jdbcTemplate.queryForList(
                "SELECT id, name, provider, model_name, api_key, base_url, is_active, is_delete " +
                        "FROM ai_model_config WHERE app_type = 'EMBEDDING' ORDER BY is_active DESC");
        log.info("========== EMBEDDING 配置检查 ==========");
        for (Map<String, Object> row : allConfigs) {
            String apiKey = (String) row.get("api_key");
            String maskedKey = apiKey != null && apiKey.length() > 8
                    ? apiKey.substring(0, 4) + "****" + apiKey.substring(apiKey.length() - 4)
                    : "****";
            log.info("  id={} name={} provider={} model={} apiKey={} baseUrl={} active={} deleted={}",
                    row.get("id"), row.get("name"), row.get("provider"),
                    row.get("model_name"), maskedKey, row.get("base_url"),
                    row.get("is_active"), row.get("is_delete"));
        }

        // 查找活跃配置
        Map<String, Object> activeConfig = null;
        for (Map<String, Object> row : allConfigs) {
            if (Boolean.TRUE.equals(row.get("is_active"))
                    && Boolean.FALSE.equals(row.get("is_delete"))) {
                activeConfig = row;
                break;
            }
        }

        if (activeConfig == null) {
            log.warn("EMBEDDING 类型无活跃配置，请在前端管理页面配置");
            return;
        }

        String apiKey = (String) activeConfig.get("api_key");
        String modelName = (String) activeConfig.get("model_name");
        String provider = (String) activeConfig.get("provider");
        String baseUrl = (String) activeConfig.get("base_url");

        // 仅记录警告，不再强制改写（用户可能有意使用非 DashScope 的 Embedding 提供商）
        if (!"dashscope".equalsIgnoreCase(provider)) {
            log.info("EMBEDDING 使用非 DashScope 提供商 [{}]，model={}。"
                    + "请确保前端「向量维度」与实际模型输出一致，否则需重建 vector_store 表", provider, modelName);
        }

        // DashScope 用户的 API Key 一致性提示（不强制修正，尊重用户 DB 配置）
        if ("dashscope".equalsIgnoreCase(provider)
                && dashscopeApiKey != null && !dashscopeApiKey.isBlank()
                && !dashscopeApiKey.equals(apiKey)) {
            log.warn("EMBEDDING DashScope 配置的 API Key 与 application.yml 不一致，以 DB 配置为准。"
                    + "如需统一管理，请在管理页面修改。");
        }

        if (baseUrl != null && !baseUrl.isBlank()) {
            log.debug("EMBEDDING 配置了自定义 baseUrl={}", baseUrl);
        }

        log.info("EMBEDDING 配置检查完成，当前使用 provider={} model={}", provider, modelName);
    }

    /**
     * 检测 EmbeddingModel 实际返回的维度与配置的向量维度是否一致。
     * 配置维度优先级：DB extraConfig.dimensions > application.yml knowledge.vector.dimension。
     */
    private void checkDimensionConsistency() {
        try {
            int configuredDimension = resolveConfiguredDimension();
            EmbeddingModel model = embeddingModelFactory.getEmbeddingModel();
            int actualDimension = model.dimensions();
            if (actualDimension > 0 && actualDimension != configuredDimension) {
                log.warn("========== ⚠️ 向量维度不匹配！==========");
                log.warn("Embedding 模型 {} 返回维度: {}", model, actualDimension);
                log.warn("当前配置的向量维度: {}", configuredDimension);
                log.warn("请执行以下步骤修复：");
                log.warn("  1. 在前端「AI 大模型配置管理」中修改 EMBEDDING 配置的向量维度为 {}", actualDimension);
                log.warn("  2. 在 PostgreSQL 中执行: DROP TABLE IF EXISTS vector_store CASCADE");
                log.warn("  3. 重启服务（PgVectorStore 会自动重建表）");
                log.warn("  4. 重新上传所有知识库文档");
                log.warn("========================================");
            } else {
                log.info("向量维度校验通过: Embedding 模型返回 {} 维，配置 {} 维，一致",
                        actualDimension, configuredDimension);
            }
        } catch (Exception e) {
            log.warn("无法校验向量维度（EmbeddingModel 初始化可能失败）: {}", e.getMessage());
        }
    }

    /**
     * 解析配置维度：DB extraConfig 优先，yml 兜底。
     */
    private int resolveConfiguredDimension() {
        try {
            List<Map<String, Object>> activeConfigs = jdbcTemplate.queryForList(
                    "SELECT extra_config FROM ai_model_config WHERE app_type = 'EMBEDDING' AND is_active = true AND is_delete = false LIMIT 1");
            if (!activeConfigs.isEmpty()) {
                Map<String, Object> row = activeConfigs.get(0);
                String extraConfigJson = (String) row.get("extra_config");
                // 简单解析 JSONB 中的 dimensions 字段
                if (extraConfigJson != null && extraConfigJson.contains("\"dimensions\"")) {
                    java.util.regex.Pattern p = java.util.regex.Pattern.compile("\"dimensions\"\\s*:\\s*(\\d+)");
                    java.util.regex.Matcher m = p.matcher(extraConfigJson);
                    if (m.find()) {
                        int dim = Integer.parseInt(m.group(1));
                        log.info("从 DB EMBEDDING extraConfig 读取维度: {}", dim);
                        return dim;
                    }
                }
            }
        } catch (Exception e) {
            log.warn("从 DB 读取 EMBEDDING 维度失败: {}", e.getMessage());
        }
        log.info("使用 application.yml fallback 维度: {}", fallbackDimension);
        return fallbackDimension;
    }
}