RagService java8环境 固定知识测试版本

package com.jiaozuo.wohuida.service;

import com.jiaozuo.wohuida.model.KnowledgeSnippet;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.*;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.*;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.SearchResultsWrapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import java.util.*;

@Slf4j
@Service
public class RagService {

    @Autowired
    private MilvusServiceClient milvusClient;

    @Autowired
    private LargeModelClient largeModelClient;

    private static final int VECTOR_DIM = 768;
    private static final String COLLECTION_NAME = "wohuida_knowledge";

    /**
     * 应用启动时，自动创建集合并初始化知识库数据
     */
    @PostConstruct
    public void initCollection() {
        System.out.println("=== 开始初始化 Milvus 集合 ===");

        // 【强制清理】：每次启动都先删除旧集合，确保索引能被正确创建！
        // (测试完成后，您可以把这5行删掉，或者注释掉)
        System.out.println("正在清理旧集合...");
        milvusClient.dropCollection(
                DropCollectionParam.newBuilder()
                        .withCollectionName(COLLECTION_NAME)
                        .build()
        );

        // 1. 重新检查集合是否存在
        boolean hasCollection = milvusClient.hasCollection(
                HasCollectionParam.newBuilder().withCollectionName(COLLECTION_NAME).build()
        ).getData();

        if (!hasCollection) {
            System.out.println("集合不存在，准备创建...");
            FieldType idField = FieldType.newBuilder().withName("id").withDataType(DataType.Int64).withPrimaryKey(true).withAutoID(true).build();
            FieldType contentField = FieldType.newBuilder().withName("content").withDataType(DataType.VarChar).withMaxLength(2048).build();
            FieldType vectorField = FieldType.newBuilder().withName("vector").withDataType(DataType.FloatVector).withDimension(VECTOR_DIM).build();

            CreateCollectionParam param = CreateCollectionParam.newBuilder()
                    .withCollectionName(COLLECTION_NAME)
                    .addFieldType(idField)
                    .addFieldType(contentField)
                    .addFieldType(vectorField)
                    .build();
            milvusClient.createCollection(param);
            System.out.println("Milvus 集合 [" + COLLECTION_NAME + "] 创建成功！");

            // 2. 创建索引 (终极方案：将 metric_type 放入 extraParam 的 JSON 字符串中)
            System.out.println("正在为向量字段创建索引...");
            CreateIndexParam indexParam = CreateIndexParam.newBuilder()
                    .withCollectionName(COLLECTION_NAME)
                    .withFieldName("vector")
                    .withIndexType(io.milvus.param.IndexType.IVF_FLAT)
                    // 【核心修复】：必须显式调用 withMetricType，并使用 io.milvus.param.MetricType
                    .withMetricType(io.milvus.param.MetricType.COSINE)
                    .withExtraParam("{\"nlist\":128}")
                    .withSyncMode(Boolean.TRUE)
                    .build();

            R<RpcStatus> response = milvusClient.createIndex(indexParam);
            if (response.getException() != null) {
                System.err.println("❌ 索引创建失败！原因: " + response.getException().getMessage());
                return;
            }
            System.out.println("✅ 索引创建成功！");

            // 3. 写入初始知识
            ingestKnowledge("融合业务拆机流程：1. 确认用户主卡无合约期限制；2. 结清当前所有欠费；3. 携带机主身份证原件到联通自有营业厅办理；4. 宽带光猫需一并归还，否则需赔偿设备费。");
            ingestKnowledge("流量超套收费标准：国内通用流量超出套餐后，按 0.03元/MB 计费，累计至 10元后自动叠加 10元/GB 流量包，不足 1GB 按 10元/GB 收取。");

            // 4. 刷盘
            milvusClient.flush(FlushParam.newBuilder().addCollectionName(COLLECTION_NAME).build());
            System.out.println("✅ 初始业务知识已成功写入并刷盘至 Milvus！");
        }

        // 5. 加载集合到内存
        System.out.println("正在加载集合到内存...");
        milvusClient.loadCollection(LoadCollectionParam.newBuilder().withCollectionName(COLLECTION_NAME).build());
        System.out.println("✅ 集合 [" + COLLECTION_NAME + "] 已成功加载到内存，可以开始检索！");
    }

    /**
     * 知识检索：根据用户问题，在向量库中寻找最相关的知识
     */
    public List<KnowledgeSnippet> searchKnowledge(String query) {
        List<Float> queryVector = largeModelClient.generateEmbedding(query);

        if (queryVector == null || queryVector.isEmpty()) {
            log.warn("警告：Embedding 生成失败，无法进行检索！");
            return Collections.emptyList();
        }

//        SearchParam searchParam = SearchParam.newBuilder()
//                .withCollectionName(COLLECTION_NAME)
//                // 使用字符串参数形式指定度量类型
//                .withParams("{\"metric_type\": \"COSINE\"}")
//                .withOutFields(Collections.singletonList("content"))
//                .withTopK(3) // 返回最相关的3条
//                .withVectors(Collections.singletonList(queryVector))
//                .withVectorFieldName("vector")
//                .build();
        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(COLLECTION_NAME)
                // 【核心修复】：使用 withMetricType 显式指定搜索时的度量类型
                .withMetricType(io.milvus.param.MetricType.COSINE)
                // 将 withParams 用于其他搜索参数，例如 nprobe
                .withParams("{\"nprobe\": 10}")
                .withOutFields(Collections.singletonList("content"))
                .withTopK(3) // 返回最相关的3条
                .withVectors(Collections.singletonList(queryVector))
                .withVectorFieldName("vector")
                .build();
        R<SearchResults> resp = milvusClient.search(searchParam);
        List<KnowledgeSnippet> results = new ArrayList<>();

        if (resp.getStatus() == R.Status.Success.getCode()) {
            SearchResultData resultsData = resp.getData().getResults();
            SearchResultsWrapper wrapper = new SearchResultsWrapper(resultsData);

            int scoresCount = resultsData.getScoresCount();
            if (scoresCount == 0) {
                System.out.println("Milvus 未检索到相关知识，返回空列表。");
                return results;
            }

            for (int i = 0; i < scoresCount; i++) {
                float score = wrapper.getIDScore(0).get(i).getScore();
                String content = wrapper.getFieldData("content", 0).get(i).toString();
                results.add(new KnowledgeSnippet(content, "沃慧答知识库", score));
            }
        }
        return results;
    }

    /**
     * 知识入库：将文本转化为向量并存入 Milvus
     */
    public void ingestKnowledge(String content) {
        // 1. 文本空校验
        if (content == null || content.trim().length() == 0) {
            log.warn("【知识入库】文本为空，跳过");
            return;
        }

        log.info("正在生成向量，文本：{}", content);
        List<Float> vector;
        try {
            vector = largeModelClient.generateEmbedding(content);
        } catch (Exception e) {
            log.error("【知识入库】调用大模型Embedding接口异常", e);
            return;
        }

        // 2. 向量判空校验
        if (vector == null || vector.isEmpty()) {
            log.error("【知识入库】生成向量为空，禁止写入Milvus");
            return;
        }

        // 3. 向量维度强校验（必须等于768）
        if (vector.size() != VECTOR_DIM) {
            log.error("【知识入库】向量维度不匹配！预期:{},实际:{}", VECTOR_DIM, vector.size());
            return;
        }

        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("content", Collections.singletonList(content)));
        fields.add(new InsertParam.Field("vector", Collections.singletonList(vector)));

        try {
            milvusClient.insert(InsertParam.newBuilder()
                    .withCollectionName(COLLECTION_NAME)
                    .withFields(fields)
                    .build());
            log.info("向量已写入 Milvus！");
        } catch (Exception e) {
            log.error("【知识入库】Milvus插入失败", e);
        }
    }
}