Springing into AI - Part 9: Tool Calling

Problem

In GenAI applications, the LLM's are pre-trained models for information relative to a point in time. This does solve for cases when certain prompts are made for that information. However, they fail to provide adequate response to request for information outside its knowledge bank. The demand for this can vary per business use case and the problem we trying to solve to provide that information to LLM. Information can be static such as a resource, or a dynamic information obtained from your own knowledge bank from business information to wide variety of application use cases. 

Solution

For static resources such as document, we have a solution that we will discover in upcoming posts. For dynamic information such as retrieving information from our database or a custom business logic where we return the data back to the LLM for it to make sense of it, we can use a concept termed Tool Calling.  In its simplest form this is merely your custom business logic done as a function which can carry out the relevant operation and return the data back to the LLM. 

    In order for LLM to determine if it needs to call our custom tool, it is registrered with the LLM with some metadata that contains information about what its capability in natural language such as what it can do, how the parameters operate for the tool. This is absolutly crucial for us to be specific when we provide this metadata as the LLM evaluates the request against the metadata and determines if the tool provided is the right one for the request or not to be processed. It is to be noted, that not all LLM's support tool calling. Capabilities of different providers and their offerings can be found hereFigure below conceptualizes how it all works. 



In the figure above, each of the numbered items in their sequence of operation are described below:
  1. 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.
  2. 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: This is merely a service class with business methods that are additionally annotated with @Tool(description="Description of the tool") and @ToolParam(description="Information about the parameter")
        • @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.
      • SyncMcpToolCallbackAsyncMcpToolCallback: 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. 
  3. 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. 
  4. Tool Invocation: Once the ToolCallingManager has decided the tool to execute, it simply calls the underlying tool configured and awaits the response. 
  5. Tool Response: The response here is then processed and converted into format which the LLM can understand and act accordingly.
  6. 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.
  7. Chat Client Response: The final response obtained from above is simply propagated back
  8. User Response: The user is presented the response to the prompt. 

Playground

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
    @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. 
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
    @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.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
    @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.  

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
    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

Comments

Popular posts from this blog

Springing into AI - Part 4: LLM - Machine Setup

Springing into AI - Part 1: 1000 feet overview

Springing into AI - Part 2: Generative AI