Spring AI: Dynamic Tool Calling

package com.tungdadev.agent.tools;
 
import org.springframework.ai.document.Document;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbacks;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.ApplicationContext;
import org.springframework.context.event.EventListener;
import org.springframework.stereotype.Component;
 
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
 
@Component
public class DynamicToolRegistry {
 
    private final ApplicationContext applicationContext;
    private final VectorStore vectorStore;
    private final Map<String, ToolCallback> callbackMap = new ConcurrentHashMap<>();
 
    public DynamicToolRegistry(ApplicationContext applicationContext, VectorStore vectorStore) {
        this.applicationContext = applicationContext;
        this.vectorStore = vectorStore;
    }
 
    @EventListener(ApplicationReadyEvent.class)
    public void indexAllTools() {
        // Tìm toàn bộ beans có chứa phương thức gắn annotation @Tool
        String[] beanNames = applicationContext.getBeanDefinitionNames();
        List<Document> toolDocuments = new ArrayList<>();
 
        for (String beanName : beanNames) {
            Object bean = applicationContext.getBean(beanName);
            ToolCallback[] callbacks = ToolCallbacks.from(bean);
 
            for (ToolCallback callback : callbacks) {
                String toolName = callback.getToolDefinition().name();
                String description = callback.getToolDefinition().description();
 
                callbackMap.put(toolName, callback);
 
                // Lưu nội dung mô tả vào Vector Store để phục vụ tìm kiếm ngữ nghĩa
                Document doc = new Document(
                    description,
                    Map.of(
                        "toolName", toolName,
                        "type", "agent_tool"
                    )
                );
                toolDocuments.add(doc);
            }
        }
 
        if (!toolDocuments.isEmpty()) {
            vectorStore.add(toolDocuments);
        }
    }
 
    public ToolCallback getCallback(String toolName) {
        return callbackMap.get(toolName);
    }
}

package com.tungdadev.agent.advisor;
 
import com.tungdadev.agent.tools.DynamicToolRegistry;
import org.springframework.ai.chat.client.advisor.api.AdvisedRequest;
import org.springframework.ai.chat.client.advisor.api.AdvisedResponse;
import org.springframework.ai.chat.client.advisor.api.CallAroundAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAroundAdvisorChain;
import org.springframework.ai.document.Document;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.Ordered;
 
import java.util.*;
 
public class DynamicToolSelectionAdvisor implements CallAroundAdvisor {
 
    private final VectorStore vectorStore;
    private final DynamicToolRegistry toolRegistry;
    private final int topK;
    private final double similarityThreshold;
    private final List<ToolCallback> alwaysOnTools;
 
    public DynamicToolSelectionAdvisor(VectorStore vectorStore,
                                       DynamicToolRegistry toolRegistry,
                                       int topK,
                                       double similarityThreshold,
                                       List<ToolCallback> alwaysOnTools) {
        this.vectorStore = vectorStore;
        this.toolRegistry = toolRegistry;
        this.topK = topK;
        this.similarityThreshold = similarityThreshold;
        this.alwaysOnTools = alwaysOnTools != null ? alwaysOnTools : List.of();
    }
 
    @Override
    public AdvisedResponse aroundCall(AdvisedRequest advisedRequest, CallAroundAdvisorChain chain) {
        String userPrompt = advisedRequest.userText();
 
        // 1. Tìm kiếm ngữ nghĩa trong Vector Store
        SearchRequest searchRequest = SearchRequest.query(userPrompt)
                .withTopK(topK)
                .withSimilarityThreshold(similarityThreshold);
 
        List<Document> matchedDocs = vectorStore.similaritySearch(searchRequest);
 
        // 2. Chuyển đổi tên tool thành danh sách ToolCallback thực thi
        Set<ToolCallback> selectedTools = new HashSet<>(alwaysOnTools);
        for (Document doc : matchedDocs) {
            String toolName = (String) doc.getMetadata().get("toolName");
            ToolCallback callback = toolRegistry.getCallback(toolName);
            if (callback != null) {
                selectedTools.add(callback);
            }
        }
 
        // 3. Tái cấu trúc AdvisedRequest với danh sách tools đã được chọn lọc
        AdvisedRequest modifiedRequest = AdvisedRequest.from(advisedRequest)
                .toolCallbacks(selectedTools.toArray(new ToolCallback[0]))
                .build();
 
        // 4. Tiếp tục chuỗi xử lý
        return chain.nextAroundCall(modifiedRequest);
    }
 
    @Override
    public int getOrder() {
        return Ordered.HIGHEST_PRECEDENCE + 100;
    }
 
    @Override
    public String getName() {
        return "DynamicToolSelectionAdvisor";
    }
}

package com.tungdadev.agent.config;
 
import com.tungdadev.agent.advisor.DynamicToolSelectionAdvisor;
import com.tungdadev.agent.tools.DynamicToolRegistry;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
 
import java.util.List;
 
@Configuration
public class AgentChatConfiguration {
 
    @Bean
    public ChatClient enterpriseAgentChatClient(ChatClient.Builder builder,
                                                VectorStore vectorStore,
                                                DynamicToolRegistry toolRegistry) {
        return builder
                .defaultAdvisors(
                    new DynamicToolSelectionAdvisor(
                        vectorStore,
                        toolRegistry,
                        3,       // Top-3 tools liên quan nhất
                        0.72,    // Ngưỡng Cosine Similarity tối thiểu
                        List.of() // Always-on tools nếu có
                    )
                )
                .build();
    }
}