Springing into AI - Part 9: Tool Calling

Welcome back, from wherever you are, I hope you've had a lovely week. One of the caveats that comes with using LLM is that they are pre-trained on information up to a certain date. Furthermore, they would not be able to provide information on custom business data unless trained to do so. In this post we look to address the latter, where we present our custom data (static, obtained from third party vendor, database, etc) to LLM through means of custom business functions termed tool-calling so that it may response accordingly to relevant prompts made by user. It is to be noted, that not all LLM's support tool calling. Capabilities of different providers and their offerings can be found here.

Tool-Calling Theory

Tool-Calling is one of many options (others being RAG, MCP, Agentic AI) that will follow in the series of posts that follow. Before we dive deep into those, understanding tool-calling would give a basic idea of how the solution is conceptualized to work with LLM to overcome their short-comings.

In the figure above, each of the numbered items in their sequence of operation are described below:

User Prompt: End user sends their prompts to our application by an exposed API. This is something we have already seen in past posts already.
Tool configured Chat Client: SpringAI provides us options to configure the chat client we have been using till now additionally with Tool(s) or ToolCallback(s). Under the hood if we configure it using "Tool", it would internally create a "ToolCallback".

Tool: It is here where we specify our own custom implementation for "tool(s)". The custom implementation can vary to provide data from wide sources such as database interaction, static data, third party API invocation, etc. SpringAI provides us four different implementations:

MethodToolCallback: Putting it simply this can be a service class with business methods that are additionally annotated with @Tool(description="blah blah") and @ToolParam(description="blah blah").

@Tool: It is of great importance and I cannot stress this enough that the words you choose to describe a tool/method is crucial for LLM to infer what tool/method it should request to be invoked.
@ToolParam: Should your tool/method require some parameters from user to fill their request, like @Tool above, the description here also should be well defined as LLM would infer these required input parameters from the prompt.
Limitations: MethodToolCallback does have some limitations that should be noted, and they can be seen here.

FunctionToolCallback: Here, we specify our custom business logic for data representation through means of Java "Function", "Supplier" and "Consumer". Each of these operations depending on your choice of implementation would have to be annotated with "@Description". Like mentioned the way we choose to word the description would be important for LLM to infer what function it should invoke for a given user prompt. These also have some limitations, that can be seen here.
SyncMcpToolCallback, AsyncMcpToolCallback: We leave this for now as discuss when we get to MCP (Model Context Protocol) section.

ToolCallback: This merely represents a way to define a tool so that it can be presented to LLM for execution. It does so by presenting information in form of "name", "description" and "information schema".

name: Represents a unique name for a tool
description: Description that we discussed in above section to let the LLM know what it is the tool does for it to infer which one to invoke accordingly.
information schema: This is merely schema of the parameters used to call tool.

Tool Calling Manager: The model decides what tool to invoke based on it's inference, it sends the response to the ToolCallingManager. It is the job of the manager to then correctly identify the tool to be invoked.
Tool Invocation: Once the ToolCallingManager has decided the tool to execute, it simply calls the underlying tool configured and awaits the response.
Tool Response: The response here is then processed and converted into format which the LLM can understand and act accordingly.
LLM with Tool Response: The response presented to LLM from our custom business logic above is then analyzed by the LLM against the prompt and we have it transform into a Natural Language Processing (NLP) human readable syntax that is propagated back.
Chat Client Response: The final response obtained from above is simply propagated back
User Response: The user is presented the response to the prompt.

Tool-Calling Playground

If the above theory confused you do not worry, let's get our hands dirty with the integration and hopefully it all starts to make sense. For our playground, we will simulate a situation where we try to present an end user our custom available mocked generated products from a database and the ability to search for a product based on given category. It is to be noted, such an information would not be known to LLM, so it makes a good use case for our playground to present custom data to LLM for response to user prompts. Before we see results, let's get some admin out of the way.

Source Code: Can be found here

Two individual classes ("ProductFunctionBasedService", "ProducMethodBasedService") to demonstrate function and method as a tool discussed above
System Prompt: We briefly touched this previously in the post where this helps us to steer the scope of LLM operation to a certain modelling behavior. We used system prompts here to ensure that our application only works with our business and nothing else.

Infrastructure Additions:

spring-boot-docker-compose
postgres docker container
sample mocked data: A sample dataset prepopulated that was generated using Mockaroo. In our use case this contains list of some random products along with quantity, category etc.

Endpoints:

http://localhost:8080/chat/method
http://localhost:8080/chat/function

Code Walkthrough - Method as a Tool

    @PostMapping("/method")
    public String productMethodSearch(@RequestBody ApplicationDomain.UserInputRequest userInputRequest) {

        return chatClient.prompt()
                .system(SYSTEM_PROMPT)
                .user(userInputRequest.input())
                .tools(productMethodBasedService)
                .call()
                .content();
    }

In the above code, besides the usual, we have introduced two new additions for our chat client:

Line 5, System Prompt: This helps as mentioned earlier to steer the context of LLM in a particular direction so that it may only cater to a specific need. We would not want in real world, an application where user can prompt for anything else besides the application scope.
Line 7, Method Tool: We simply pass a reference of the class which has some methods that are annotated with @Tool, @ToolParam accordingly. These methods would contain our business logic that we would want to perform for a given prompt and use the lifecycle described above to allow LLM to present it as a response back to user. Code snippet shown below is a representation of how we can use it.

    @Tool(description = "Retrieves list containing information about all available products")
    public List<ProductInventory> getProductList() {
        return (List<ProductInventory>) productRepository.findAll();
    }

    @Tool(description = "Retrieves product information for a specified category")
    public List<ProductInventory> getProductByCategory(
        @ToolParam(description = "Category for which product must be fetched") 
        final String category) {

        log.info("Product category {}", category);
        return productRepository.findAllByProductCategory(category);
    }

In above, for both the cases our methods are interacting with a custom database repository and providing information back as response that can be used by LLM. Explanation for key elements are described below:

Line 1, 6 - @Tool: I cannot stress it enough the description that we specify here is of great importance as again, it is what is used by LLM to infer the relevant tool to be invoked. For our demonstration we are offering two capabilities firstly being able to view list of products and second being able to search for products based on a given category.
Line 8 - @ToolParam: Since we need some user input for filter criteria, we absorb that information via ToolParam. The LLM infers the prompt and from it tries its best to provide the right value for the parameter based on the description we specify for the parameter.

Sample Run: Querying for products based on a given category

Sample Run: Querying for list of available products

Code Walkthrough - Function as a Tool

Like method as a tool, we have another endpoint to demonstrate the use of function as a tool.

    @PostMapping("/function")
    public String productFunctionSearch(@RequestBody ApplicationDomain.UserInputRequest userInputRequest) {

        return chatClient.prompt()
                .system(SYSTEM_PROMPT)
                .user(userInputRequest.input())
                .toolNames("productSearchByCategory", "productListSupplier")
                .call()
                .content();
    }

From the code above, we use the system prompt as we had done for method as a tool. The difference is in how have we configured chat client to use function as a tool by providing it the name of the Spring beans that are our functions.

    public record ProductSearch(
      @ToolParam(description = "Category for which the products must be searched") 
      String category) {}

    @Bean
    @Description("Based on the given category, this provides list of available products belonging to that category")
    public Function<ProductSearch, String> productSearchByCategory(
                                        final ProductRepository productRepository) {
        return productSearch -> {
            log.info("Searching for products by category {}", productSearch.category());
            List<ProductInventory> productInventoryList = productRepository.findAllByProductCategory(productSearch.category());
            String productList = null;
            try {
                productList = new ObjectMapper().writeValueAsString(productInventoryList);
            } catch (JsonProcessingException e) {
                log.error("Error occurred processing json", e);
            }
            return productList;
        };
    }

    @Bean
    @Description("Provides list of all the products available")
    public Supplier<String> productListSupplier(final ProductRepository productRepository) {
        return () -> {
            log.info("Getting list of products...");
            List<ProductInventory> productInventoryList = (List<ProductInventory>) productRepository.findAll();
            try {
                return new ObjectMapper().writeValueAsString(productInventoryList);
            } catch (JsonProcessingException e) {
                log.error("Error occurred processing json", e);
            }
            return null;
        };
    }

Woah, 35 lines of code !!. Bear with me, soon this will make sense as what we doing above is no different to what you have already seen in method as a tool section. Like before we offering two capabilities to the LLM for it to respond to list of available products and criteria based search.

Line 5, 22 - @Bean: These are simply declaring the two function names as Spring Beans, and it is this beans that you saw few section above that we configure into chat client as tools.
Line 6, 23 - @Description: Like @Tool that you saw in method as a tool section, Function as a tool make use of this annotation to provide description of its operation it performs.
Line 2 - @ToolParam: This is just providing a description of what that parameter is all about so that LLM can infer the description and extract the relevant information out from the user prompt.
Line 7 - Function: This is a type of Function declaration that takes an input and returns an output
Line 24 - Supplier: This is another type of Function that accepts no argument and provides an output response. The rest of the business logic code for both is just the usual code where we interact with the database and convert the object into String for overcoming limitations of Function as a tool mentioned in theory section.

Sample Run: Querying for products based on a given category

Sample Run: Querying for list of all available products

Log Analysis

Since both "method as a tool", and "function as a tool" works in similar fashion, please find the sample logs observer for each. The analysis of these logs summarizes the flow of events:

Sample Log - Method as a tool

Sample Log - Function as a tool

From above:

We present the system prompt along with the user prompt to the LLM
The LLM then tries to infer what tool to be invoked
This request is then processed by SpringAI tool resolver to ensure correct tool for invocation is identified
SpringAI invokes the custom tool using ToolCallingManager
The result obtained from the custom tool is then converted to JSON
This is then presented to LLM
LLM does its magic, and presents the response back to the end user

That should provide an understanding of how tool calling works. In the next post, we will look at Retrieval Augmented Generation (RAG) which like tool-calling allow us to enrich LLM with our custom data, but in particular in next post we will present a sample document such as PDF and ask LLM to analyze. Hopefully that should give you a teaser of exciting things to come. Stay tuned...

Search This Blog

Everything will be 200 OK :)