langchain4j动态创建agent

package com.kongjs.smo.ai.langchain4j.service;

import com.kongjs.smo.ai.langchain4j.agent.AiAgent;
import com.kongjs.smo.ai.langchain4j.store.PersistentChatMemoryStore;
import com.kongjs.smo.ai.dto.ParamDTO;
import com.kongjs.smo.ai.entity.Agent;
import com.kongjs.smo.ai.entity.KnowledgeBase;
import com.kongjs.smo.ai.service.KnowledgeBaseService;
import com.kongjs.smo.common.mybatis.flex.enums.EnableEnum;
import com.kongjs.smo.common.web.exception.BusinessException;
import com.mybatisflex.core.query.QueryWrapper;
import com.zaxxer.hikari.HikariDataSource;
import dev.langchain4j.agentic.AgenticServices;
import dev.langchain4j.agentic.agent.AgentBuilder;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.bgesmallzhv15q.BgeSmallZhV15QuantizedEmbeddingModel;
import dev.langchain4j.model.embedding.request.EmbeddingInputType;
import dev.langchain4j.rag.DefaultRetrievalAugmentor;
import dev.langchain4j.rag.RetrievalAugmentor;
import dev.langchain4j.rag.content.injector.ContentInjector;
import dev.langchain4j.rag.content.injector.DefaultContentInjector;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.router.DefaultQueryRouter;
import dev.langchain4j.rag.query.router.QueryRouter;
import dev.langchain4j.rag.query.transformer.CompressingQueryTransformer;
import dev.langchain4j.rag.query.transformer.QueryTransformer;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.pgvector.PgVectorEmbeddingStore;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.List;

@Service
public class AiAgentServiceImpl implements AiAgentService {

    private final Logger logger = LoggerFactory.getLogger("Model");
    @Resource
    private HikariDataSource dataSource;
    @Resource(name = "aiDataSource")
    private HikariDataSource aiDataSource;
    @Resource
    private PersistentChatMemoryStore persistentChatMemoryStore;
    @Resource
    private com.kongjs.smo.ai.service.AgentService agentService;
    @Resource
    private KnowledgeBaseService knowledgeBaseService;
    @Resource
    private AiModelService modelService;

    @Override
    public AiAgent aiAgent(Long agentId, ParamDTO param) {
        Agent agent = agentService.getById(agentId);
        if (agent == null) {
            throw new BusinessException("智能体不存在");
        }
        if (!agent.getStatus().equals(EnableEnum.ENABLE.getCode())) {
            throw new BusinessException("智能体未启用");
        }
        List<KnowledgeBase> knowledgeBases = getKnowledgeBases(param.getKnowledgeBaseIds());
        ChatModel chatModel = modelService.getChatModel(param.getModelId());
        StreamingChatModel streamingChatModel = modelService.getStreamingChatModel(param.getModelId());
        //  https://docs.langchain4j.dev/tutorials/chat-memory#memory-vs-history
        ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(50)
                .chatMemoryStore(persistentChatMemoryStore)
                .build();
        AgentBuilder<AiAgent, AgentBuilder<AiAgent, ?>> agentBuilder = AgenticServices.agentBuilder(AiAgent.class);
        if (!CollectionUtils.isEmpty(knowledgeBases)) {
            List<ContentRetriever> contentRetrievers = new ArrayList<>();
            for (KnowledgeBase knowledgeBase : knowledgeBases) {
                // 向量模型
                EmbeddingModel embeddingModel = new BgeSmallZhV15QuantizedEmbeddingModel();
                // 向量存储 https://docs.langchain4j.dev/integrations/embedding-stores/pgvector/
                // 1. **VECTOR（默认）**：仅余弦相似度向量检索
                // 2. **HYBRID**：向量相似度 + PostgreSQL 全文搜索，**RRF 倒数排名融合**多路结果
                EmbeddingStore<TextSegment> embeddingStore = PgVectorEmbeddingStore.datasourceBuilder()
                        .skipCreateVectorExtension(true)
                        .datasource(aiDataSource)
                        .table(knowledgeBase.getCollection())
                        .searchMode(PgVectorEmbeddingStore.SearchMode.VECTOR)
                        .dimension(knowledgeBase.getDimension() == null ? embeddingModel.dimension() : knowledgeBase.getDimension())
                        .build();
                // rag
                ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                        .embeddingStore(embeddingStore)
                        .embeddingModel(embeddingModel)
                        .embeddingInputType(EmbeddingInputType.QUERY)
                        .maxResults(20)
                        .minScore(0.75)
                        .build();
                contentRetrievers.add(contentRetriever);
            }
            ContentInjector contentInjector = DefaultContentInjector.builder()
                    .metadataKeysToInclude(List.of())
                    .build();
            QueryTransformer queryTransformer = new CompressingQueryTransformer(chatModel);
            QueryRouter queryRouter = new DefaultQueryRouter(contentRetrievers);
            RetrievalAugmentor retrievalAugmentor = DefaultRetrievalAugmentor.builder()
                    .queryTransformer(queryTransformer)
                    .contentInjector(contentInjector)
                    .queryRouter(queryRouter)
                    .build();
            agentBuilder.retrievalAugmentor(retrievalAugmentor);
        }
        return agentBuilder
                .name(agent.getAgentName())
                .description(agent.getDescription())
                .systemMessage(agent.getSystemPrompt())
                .userMessage(agent.getUserPrompt())
                .streamingChatModel(streamingChatModel)
                .chatMemoryProvider(chatMemoryProvider)
                .build();
    }

    public List<KnowledgeBase> getKnowledgeBases(List<Long> knowledgeBaseId) {
        if (CollectionUtils.isEmpty(knowledgeBaseId)) {
            return null;
        }
        return knowledgeBaseService.list(QueryWrapper.create()
                .in(KnowledgeBase::getId, knowledgeBaseId)
                .eq(KnowledgeBase::getStatus, EnableEnum.ENABLE.getCode()).limit(5));
    }
}