package com.cmc.agent.controller;

import com.ruoyi.common.config.RuoYiConfig;
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.onnx.bgesmallzhv15.BgeSmallZhV15EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;

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

    /**
     * 新建知识库
     */
    @PostMapping("/create")
    public AjaxResult createKnowLedgeBase(String collectionName)
    {
        LangChainMilvusService langChainMilvusService = new LangChainMilvusService(
                "localhost",
                19530,
                collectionName,
                embeddingModel);
        langChainMilvusService.createCollection(512);
        return success();
    }

    /**
     * 插入知识库文件
     */
    @PostMapping("/insert")
    public AjaxResult insertKnowledge(MultipartFile file, String collectionName) throws IOException {
        LangChainMilvusService langChainMilvusService = new LangChainMilvusService(
                "localhost",
                19530,
                collectionName,
                embeddingModel);

        File profilePath = new File( RuoYiConfig.getProfile() + "/upload/knowledge");
        if (!profilePath.exists())
            profilePath.mkdirs();
        File transferFile = new File( profilePath + "/" + file.getOriginalFilename());
        if (!transferFile.exists())
            file.transferTo(transferFile);
        langChainMilvusService.insertLangchainEmbeddingDocument(transferFile);
        return success();
    }

}
