package com.cmc.agent.controller;

import com.ruoyi.common.core.controller.BaseController;
import com.ruoyi.common.core.domain.AjaxResult;
import com.cmc.agent.service.LangChainMilvusService;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.bgesmallzhv15.BgeSmallZhV15EmbeddingModel;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.IOException;
import java.util.List;

/**
 * cmc知识库Controller
 * 
 * @author cmc
 * @date 2025-04-08
 */
@RestController
@RequestMapping("/llm/rag")
public class RagController extends BaseController
{
    private static final EmbeddingModel embeddingModel = new BgeSmallZhV15EmbeddingModel();

    /**
     * 增强检索生成回答
     */
    @PostMapping("/answer")
    public AjaxResult answerWithRAG(String question, String collectionName) throws IOException {
        LangChainMilvusService langChainMilvusService = new LangChainMilvusService(
                "localhost",
                19530,
                collectionName,
                embeddingModel);

        // 2. Milvus检索
        List<String> contexts = langChainMilvusService.retrieveFromMilvus(question, 3);

        // 3. 调用本地LLM或HTTP服务
        return success(langChainMilvusService.generateAnswer(question, contexts));
    }

}
