Spring AI: Dynamic Tool Calling

package com.tungdadev.agent.tools;
 
import org.springframework.ai.document.Document;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.annotation.Tool;
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 org.springframework.util.ClassUtils;
import org.springframework.util.ReflectionUtils;
 
import java.nio.charset.StandardCharsets;
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() {
        List<Document> toolDocuments = new ArrayList<>();
 
        for (String beanName : applicationContext.getBeanDefinitionNames()) {
            // Lọc theo kiểu trước khi getBean(): ToolCallbacks.from() ném IllegalStateException
            // nếu object không có method @Tool nào, còn getBean() mọi bean sẽ khởi tạo cả bean lazy
            Class<?> beanType = applicationContext.getType(beanName);
            if (beanType == null || !hasToolMethod(beanType)) {
                continue;
            }
 
            Object bean = applicationContext.getBean(beanName);
            for (ToolCallback callback : ToolCallbacks.from(bean)) {
                String toolName = callback.getToolDefinition().name();
                String description = callback.getToolDefinition().description();
 
                callbackMap.put(toolName, callback);
 
                // ID cố định theo tên tool: restart thì ghi đè document cũ thay vì thêm bản trùng
                String id = UUID.nameUUIDFromBytes(toolName.getBytes(StandardCharsets.UTF_8)).toString();
                toolDocuments.add(new Document(id, description, Map.of(
                        "toolName", toolName,
                        "type", "agent_tool"
                )));
            }
        }
 
        if (!toolDocuments.isEmpty()) {
            vectorStore.add(toolDocuments);
        }
    }
 
    private static boolean hasToolMethod(Class<?> beanType) {
        return Arrays.stream(ReflectionUtils.getDeclaredMethods(ClassUtils.getUserClass(beanType)))
                .anyMatch(method -> method.isAnnotationPresent(Tool.class));
    }
 
    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.ChatClientRequest;
import org.springframework.ai.chat.client.ChatClientResponse;
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
import org.springframework.ai.document.Document;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
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 CallAdvisor {
 
    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 ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
        // Model không hỗ trợ tool calling thì để nguyên request
        if (!(request.prompt().getOptions() instanceof ToolCallingChatOptions options)) {
            return chain.nextCall(request);
        }
 
        String userPrompt = request.prompt().getUserMessage().getText();
 
        // 1. Tìm kiếm ngữ nghĩa trong Vector Store, chỉ trong các document mô tả tool
        SearchRequest searchRequest = SearchRequest.builder()
                .query(userPrompt)
                .topK(topK)
                .similarityThreshold(similarityThreshold)
                .filterExpression("type == 'agent_tool'")
                .build();
 
        List<Document> matchedDocs = vectorStore.similaritySearch(searchRequest);
 
        // 2. Chuyển đổi tên tool thành ToolCallback thực thi (giữ thứ tự, không trùng tên)
        Map<String, ToolCallback> selectedTools = new LinkedHashMap<>();
        for (ToolCallback tool : alwaysOnTools) {
            selectedTools.put(tool.getToolDefinition().name(), tool);
        }
        for (Document doc : matchedDocs) {
            String toolName = (String) doc.getMetadata().get("toolName");
            ToolCallback callback = toolRegistry.getCallback(toolName);
            if (callback != null) {
                selectedTools.putIfAbsent(toolName, callback);
            }
        }
 
        // 3. Sao chép ChatOptions, chỉ giữ các tool đã chọn, rồi gắn vào Prompt mới
        ToolCallingChatOptions newOptions = options.copy();
        newOptions.setToolCallbacks(new ArrayList<>(selectedTools.values()));
 
        ChatClientRequest modifiedRequest = request.mutate()
                .prompt(request.prompt().mutate().chatOptions(newOptions).build())
                .build();
 
        // 4. Tiếp tục chuỗi xử lý
        return chain.nextCall(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();
    }
}