Implementing an AI Assistant with Spring AI

1. Overview

In this tutorial, we’ll discuss the concepts of Spring AI that can help create an AI assistant using LLMs like ChatGPT, Ollama, Mistreal, etc.

Enterprises are increasingly adopting AI assistants to enhance the user experience across a wide range of existing business functionalities:

• Answering user queries
• Performing transactions based on user input
• Summarizing lengthy sentences and documents

While these are just a few fundamental capabilities of LLMs, their potential extends far beyond these functionalities.

2. Spring AI Features

The Spring AI framework offers a range of exciting features to deliver AI-driven functionalities:

• interfaces that can integrate seamlessly with the underlying LLM services and Vector DBs
• context-aware response generation and action execution using RAG and the Function Calling APIs
• structured output converters to transform LLM responses into POJOs or machine-readable formats like JSON
• enrich prompts and apply guardrails through interceptors provided by the Advisor API
• enhanced user engagement by maintaining conversation states

We can visualize it as well:

To illustrate some of these features, we’re going to build a chatbot in a legacy Order Management System (OMS):

Typical functionalities of an OMS consists of:

• Create order
• Get user orders

3. Prerequisites

First, we’ll need an OpenAI subscription to use its LLM service. Then in the Spring Boot application, we’ll add the Maven dependencies for Spring AI libraries. We’ve already covered the prerequisites in our other articles in great detail, therefore we’ll not elaborate further on this topic.

Furthermore, we’ll use the in-memory HSQLDB database to get started quickly. Let’s create the requisite tables and insert some data in it:

CREATE TABLE User_Order (
    order_id BIGINT NOT NULL PRIMARY KEY,
    user_id VARCHAR(20) NOT NULL,
    quantity INT
);

INSERT INTO User_Order (order_id, user_id, quantity) VALUES (1, 'Jenny', 2);
INSERT INTO User_Order (order_id, user_id, quantity) VALUES (2, 'Mary', 5);
INSERT INTO User_Order (order_id, user_id, quantity) VALUES (3, 'Alex', 1);
INSERT INTO User_Order (order_id, user_id, quantity) VALUES (4, 'John', 3);
INSERT INTO User_Order (order_id, user_id, quantity) VALUES (5, 'Sophia', 4);
--and so on..

We’ll use a few standard configuration properties related to initializing the HSQLDB and OpenAI client in the Spring application.properties file:

spring.datasource.driver-class-name=org.hsqldb.jdbc.JDBCDriver
spring.datasource.url=jdbc:hsqldb:mem:testdb;DB_CLOSE_DELAY=-1
spring.datasource.username=sa
spring.datasource.password=
spring.jpa.hibernate.ddl-auto=none

spring.ai.openai.chat.options.model=gpt-4o-mini
spring.ai.openai.api-key=xxxxxxx

Selecting the right model fit for a use case is a complex iterative process involving a lot of trial and error. Nevertheless, for a simple demo application for this article, the cost-effective GPT-4o mini model will suffice.

4. Function Calling API

This feature is one of the pillars of the popular agentic concept in LLM-based applications. It enables applications to perform a complex collection of precise tasks and make decisions independently.

For instance, it can help immensely to build a chatbot in the legacy Order Management application.  The chatbot can help users raise order requests, retrieve order history, and do more through natural language.

These are specialized skills powered by one or more application functions. We define the algorithm in a prompt sent to the LLM with the supporting function schema. The LLM receives this schema and identifies the correct function to perform the requested skill. Later, it sends its decision to the application.

Finally, the application executes the function and sends back more information to the LLM:

First, let’s look at the main components of the legacy application:

The OrderManagementService class has two important functions: creating orders and fetching user order history. Both functions use the OrderRepository bean to integrate with the DB:

@Service
public class OrderManagementService {
    @Autowired
    private OrderRepository orderRepository;

    public Long createOrder(OrderInfo orderInfo) {
        return orderRepository.save(orderInfo).getOrderID();
    }

    public Optional<List<OrderInfo>> getAllUserOrders(String userID) {
       return orderRepository.findByUserID(userID);
    }
}

Now, let’s see how Spring AI can help implement a chatbot in the legacy application:

In the class diagram, the OmAiAssistantConfiguration class is a Spring configuration bean. It registers the function callback beans, createOrderFn and getUserOrderFn:

@Configuration
public class OmAiAssistantConfiguration {
    @Bean
    @Description("Create an order. The Order ID is identified with orderID. "
      + The order quantity is identified by orderQuantity."
      + "The user is identified by userID. "
      + "The order quantity should be a positive whole number."
      + "If any of the parameters like user id and the order quantity is missing"
      + "then ask the user to provide the missing information.")
    public Function<CreateOrderRequest, Long> createOrderFn(OrderManagementService orderManagementService) {
        return createOrderRequest -> orderManagementService.createOrder(createOrderRequest.orderInfo());
    }
    @Bean
    @Description("get all the orders of an user. The user ID is identified with userID.")
    public Function<GetOrderRequest, List<OrderInfo>> getUserOrdersFn(OrderManagementService orderManagementService) {
        return getOrderRequest -> orderManagementService.getAllUserOrders(getOrderRequest.userID()).get();
    }
}

The @Description annotation helps generate the function schema. Subsequently, the application sends this schema to the LLM as part of the prompt. Since the functions use the existing methods of the OrderManagementService class, it promotes reusability.

Additionally, CreateOrderRequest and GetOrderRequests are Record classes, that help Spring AI generate the POJOs for the downstream service calls:

record GetOrderRequest(String userID) {}

record CreateOrderRequest(OrderInfo orderInfo) {}

Finally, let’s take a look at the new OrderManagementAIAssistant class that will send the user queries to the LLM service:

@Service
public class OrderManagementAIAssistant {
    @Autowired
    private ChatModel chatClient;

    public ChatResponse callChatClient(Set<String> functionNames, String promptString) {
        Prompt prompt  = new Prompt(promptString, OpenAiChatOptions
          .builder()
          .withFunctions(functionNames)
          .build()
        );
        return chatClient.call(prompt);
    }
}

The callChatClient() method helps register the functions in the Prompt object. Later, it invokes the ChatModel#call() method to get the response from the LLM service.

5. Function Calling Scenarios

For a user query or instruction given to an AI assistant, we’ll cover a few basic scenarios:

• LLM decides and identifies one or more functions to execute
• LLM complains about incomplete information for executing the function
• LLM executes statements conditionally

We’ve discussed the concepts so far; therefore, moving ahead, we’ll use this feature to build the chatbot.

5.1. Execute a Callback Function Once or More Times

Now, let’s study the behavior of the LLM when we invoke it with a prompt consisting of the user query and the function schema.

We’ll begin with an example that creates an order:

void whenOrderInfoProvided_thenSaveInDB(String promptString) {
    ChatResponse response = this.orderManagementAIAssistant
      .callChatClient(Set.of("createOrderFn"), promptString);
    String resultContent = response.getResult().getOutput().getText();
    logger.info("The response from the LLM service: {}", resultContent);
}

Surprisingly, with natural language, we get the desired results:

Moving on, let’s see if the LLM can understand the prompt that requests fetching order details of a user:

void whenUserIDProvided_thenFetchUserOrders(String promptString) {
    ChatResponse response = this.orderManagementAIAssistant
      .callChatClient(Set.of("getUserOrdersFn"), promptString);
    String resultContent = response.getResult().getOutput().getText();
    logger.info("The response from the LLM service: {}", resultContent);
}

The results are similar to the previous case. The LLM successfully identifies the registered function to execute. Later, the Spring Boot application invokes the requested function and returns the order details of one or more users.

5.2. Conditionally Execute Callback Functions

Let’s consider the program that would create order only when certain criteria are met:

void whenUserIDProvided_thenCreateOrderIfUserHasLessThanTwoOrders(String promptString) {
    ChatResponse response = this.orderManagementAIAssistant
      .callChatClient(Set.of("getUserOrdersFn", "createOrderFn"), promptString);
    String resultContent = response.getResult()
      .getOutput()
      .getText();
    logger.info("The response from the LLM service: {}", resultContent);
}

We’ve registered the getUserOrderFn and CreateOrderFn functions in the ChatModel bean. Further, the LLM must determine which function will produce the best results for the requested action.

Just like before, the results are encouraging:

6. Spring AI Advisors API

In the previous sections, we discussed the application’s functional aspects. However, there are a few common concerns across all the functionalities, such as:

• Prevent users from entering sensitive information
• Logging and auditing of the user queries
• Maintain conversational states
• Enriching prompts

Fortunately, the Advisors APIs can help consistently address these concerns. We’ve already explained this in detail in one of our articles.

7. Spring AI Structured Output API and Spring AI RAG

The LLMs mostly generate response messages to user queries in natural languages. However, the downstream services mostly understand the messages in machine-readable formats like POJOs, JSON, etc. This is where the Spring AI’s capability to generate structured output plays an important role.

Like Advisors APIs, we’ve explained this in great detail in one of our articles. Hence, we’ll not discuss more on this and move to the next topic.

Sometimes, the applications have to query a Vector DB to perform a semantic search on the stored data to retrieve additional information. Later, the fetched result from the Vector DB is used in the prompt to provide contextualized information to the LLM. This is called the RAG technique, which can also be implemented using Spring AI.

8. Conclusion

In this article, we explored the key Spring AI features that can help create an AI assistant. Spring AI is evolving with lots of out-of-the-box features. However, the right selection of the underlying LLM service and Vector DB is crucial, irrespective of the programming framework. Additionally, achieving the optimum configuration of these services can be tricky and demands a lot of effort. Nevertheless, it’s important for the application’s wide adoption.