Tool Calling

package com.example.docsassistant;

import java.util.List;

import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.stereotype.Component;

@Component
public class DocsTools {

    public record DocHit(String sourcePath, String section, String content, double score) { }

    private final VectorStore vectorStore;

    public DocsTools(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }

    @Tool(description = """
            Search the DocsAssistant documentation and return the best matching passages. \
            Use it for every question about the documented libraries; do not answer from memory.""")
    public List<DocHit> searchDocs(
            @ToolParam(description = "What to look for, phrased as a natural-language query") String query,
            @ToolParam(description = "Documentation component to restrict the search to, for example "
                    + "irurueta-algebra; omit to search everything", required = false) String component) {
        SearchRequest.Builder request = SearchRequest.builder().query(query).topK(5);
        if (component != null && !component.isBlank()) {
            // the component comes from the model, so build the filter instead of concatenating a string
            request.filterExpression(new FilterExpressionBuilder().eq("component", component).build());
        }
        return this.vectorStore.similaritySearch(request.build()).stream()
                .map(d -> new DocHit(
                        (String) d.getMetadata().get("source_path"),
                        (String) d.getMetadata().get("section"),
                        d.getText(),
                        d.getScore() != null ? d.getScore() : 0.0))
                .toList();
    }
}

package com.example.docsassistant;

import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.stereotype.Component;

@Component
public class IssueTools {

    private final IssueApprovalQueue approvals;

    public IssueTools(IssueApprovalQueue approvals) {
        this.approvals = approvals;
    }

    @PreAuthorize("hasRole('MAINTAINER')")
    @Tool(description = """
            Draft a GitHub issue about a documentation problem. The issue is NOT created: a maintainer \
            must approve the draft first. Tell the user that the draft is waiting for approval.""")
    public String openIssue(
            @ToolParam(description = "Short imperative title, at most 80 characters") String title,
            @ToolParam(description = "Markdown description of the problem, with the page it affects") String body,
            ToolContext toolContext) {
        if (title.length() > 80) {
            // a message the model can act on: the default processor sends it back so the model retries
            throw new IllegalArgumentException("title is " + title.length() + " characters; shorten it to 80");
        }
        String requester = (String) toolContext.getContext().get("requester");
        if (requester == null) {
            throw new IllegalStateException("no requester in the tool context; refusing to draft an issue");
        }
        String draftId = this.approvals.submit(new IssueDraft(requester, title.strip(), body));
        return "Draft " + draftId + " is waiting for maintainer approval; the issue has not been created yet.";
    }
}

interface IssueApprovalQueue {

    String submit(IssueDraft draft);
}

record IssueDraft(String requester, String title, String body) { }

package com.example.docsassistant;

import org.springframework.context.annotation.Configuration;
import org.springframework.security.config.annotation.method.configuration.EnableMethodSecurity;

@Configuration(proxyBeanMethods = false)
@EnableMethodSecurity
class MethodSecurityConfig {
}

package com.example.docsassistant;

import java.util.Map;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.security.core.Authentication;
import org.springframework.stereotype.Service;

@Service
class DocsToolChat {

    private final ChatClient chatClient;
    private final IssueTools issueTools;

    DocsToolChat(ChatClient.Builder builder, DocsTools docsTools, IssueTools issueTools) {
        // read-only tools are safe as defaults; the write tool is added per request below
        this.chatClient = builder.defaultTools(docsTools).build();
        this.issueTools = issueTools;
    }

    String ask(Authentication authentication, String question) {
        boolean maintainer = authentication.getAuthorities().stream()
                .anyMatch(a -> "ROLE_MAINTAINER".equals(a.getAuthority()));
        ChatClient.ChatClientRequestSpec request = this.chatClient.prompt()
                .user(question)
                // never sent to the model; the tool reads it to know who is asking
                .toolContext(Map.of("requester", authentication.getName()));
        if (maintainer) {
            request = request.tools(this.issueTools);
        }
        return request.call().content();
    }
}

package com.example.docsassistant;

import java.util.function.Function;

import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.ai.tool.function.FunctionToolCallback;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration(proxyBeanMethods = false)
class CatalogTools {

    record ComponentRequest(
            @ToolParam(description = "Component name, for example irurueta-algebra") String component) { }

    record PageCount(String component, int pages) { }

    interface DocsCatalog {

        int pageCount(String component);
    }

    @Bean
    ToolCallback pageCount(DocsCatalog catalog) {
        Function<ComponentRequest, PageCount> function =
                request -> new PageCount(request.component(), catalog.pageCount(request.component()));
        return FunctionToolCallback.builder("pageCount", function)
                .description("Return how many documentation pages a component has")
                .inputType(ComponentRequest.class)
                .build();
    }
}

package com.example.docsassistant;

import java.lang.reflect.Method;

import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.ai.tool.method.MethodToolCallback;
import org.springframework.ai.tool.method.MethodToolCallbackProvider;
import org.springframework.ai.tool.support.ToolDefinitions;
import org.springframework.util.ReflectionUtils;

class ProgrammaticTools {

    // 1. every @Tool method of one or more objects, as an array
    ToolCallback[] fromObjects(DocsTools docsTools) {
        return ToolCallbacks.from(docsTools);
    }

    // 2. the same, as a provider that can be a bean or be passed to tools(...)
    ToolCallbackProvider asProvider(DocsTools docsTools) {
        return MethodToolCallbackProvider.builder().toolObjects(docsTools).build();
    }

    // 3. one method, with a definition you control (no @Tool needed on the method)
    ToolCallback oneMethod(DocsTools docsTools) {
        Method method = ReflectionUtils.findMethod(DocsTools.class, "searchDocs", String.class, String.class);
        return MethodToolCallback.builder()
                .toolDefinition(ToolDefinitions.builder(method)
                        .name("search_docs")
                        .description("Search the documentation")
                        .build())
                .toolMethod(method)
                .toolObject(docsTools)
                .build();
    }
}

package com.example.docsassistant;

import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

@Component
class PageLinkTools {

    private final String siteUrl;

    PageLinkTools(@Value("${docsassistant.site-url}") String siteUrl) {
        this.siteUrl = siteUrl.endsWith("/") ? siteUrl : siteUrl + "/";
    }

    // the result is the final answer: the model never sees it and no second model call is made
    @Tool(description = "Return the canonical URL of a documentation page", returnDirect = true)
    String pageUrl(@ToolParam(description = "Path of the page below the site root, for example "
            + "ai/spring-ai/advisors.html") String path) {
        if (!path.matches("[a-z0-9/_-]+\\.html")) {
            throw new IllegalArgumentException("path must look like ai/spring-ai/advisors.html");
        }
        return this.siteUrl + path;
    }
}

package com.example.docsassistant;

import java.util.List;

import org.springframework.ai.tool.execution.DefaultToolExecutionExceptionProcessor;
import org.springframework.ai.tool.execution.ToolExecutionExceptionProcessor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.security.access.AccessDeniedException;

@Configuration(proxyBeanMethods = false)
class ToolErrorPolicy {

    @Bean
    ToolExecutionExceptionProcessor toolExecutionExceptionProcessor() {
        return DefaultToolExecutionExceptionProcessor.builder()
                // false: other RuntimeExceptions are sent back to the model as a message it can react to
                .alwaysThrow(false)
                // authorisation failures are not the model's business: unwrap and rethrow to the caller
                .rethrowExceptions(List.of(AccessDeniedException.class))
                .build();
    }
}