Spring AI - Query databases with OpenAI

In our Spring AI framework guide we have already discussed how to implement both function and tool calls. Now it's time to see how to use what we've learned to ensure that the LLM model can independently generate and execute database query.

As you will surely have guessed by reading previous articles, the possibility of making an LLM model decide when the right time to invoke a function is, opens the way to infinite possibilities and constitutes a revolution in the field of AI. Furthermore, the model is not only able to decide when to make the call but is also capable of passing as input the variables necessary to make it work.

So let’s go into more detail with our practical tutorial.

Database query tutorial with Spring AI
Preparation of the DB
Spring configuration for DB access
Java implementation with Spring AI and OpenAI
Configuration
Function
RestController
The final test

Database query tutorial with Spring AI

Preparation of the DB

The first step of our tutorial involves setting up a local database with a simple table that will contain a list of books. As a database we will use Oracle XE which we have already seen how to start on docker in our previous article.

Once the DB is started, we create a user and a table by executing this query:

SQL

CREATE USER SPRING_AI_TEST IDENTIFIED BY SPRING_AI_PWD;
GRANT  ALL PRIVILEGES TO SPRING_AI_TEST;

CREATE TABLE "SPRING_AI_TEST"."BOOKS" (
	"TITLE" VARCHAR2(50 BYTE) NOT NULL ENABLE,
	"AUTHOR" VARCHAR2(50 BYTE) NOT NULL ENABLE,
	PRIMARY KEY ("TITLE")
);

We then created a user “SPRING_AI_TEST” with password “SPRING_AI_PWD” and a table “BOOKS” with two columns representing the title and author of the book.

In case of “ORA-65096” error while running the script, simply execute the following statement first:

SQL

ALTER SESSION SET "_ORACLE_SCRIPT"=TRUE;

Spring configuration for DB access

To be quick, let’s start from our example tutorial on how to create a project from scratch using Spring AI and OpenAI and add what is necessary to create the connection with the DB.

First we import the dependencies from Spring JDBC and the Oracle driver into the pom:

XML

<dependency>
	<groupId>org.springframework.boot</groupId>
	<artifactId>spring-boot-starter-jdbc</artifactId>
</dependency>
<dependency>
	<groupId>com.oracle.database.jdbc</groupId>
	<artifactId>ojdbc8</artifactId>
</dependency>

And then you need to set the datasource configurations in the application.properties:

Plaintext

spring.datasource.driver-class-name=oracle.jdbc.driver.OracleDriver
spring.datasource.url=jdbc:oracle:thin:@localhost:1521:xe
spring.datasource.username=SPRING_AI_TEST
spring.datasource.password=SPRING_AI_PWD

Java implementation with Spring AI and OpenAI

Configuration

We will use what we have already learned about the use of FunctionTools in Spring AI to define a tool that acts as a connector with our database. First it is necessary to define a Configuration class to define the function and the related tool:

Java

package com.example.aidemo.config;

import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.model.function.FunctionCallbackWrapper;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import com.example.aidemo.service.BooksService;

@Configuration
public class Tools {
	
	public static final String BOOKS_FUNCTION_NAME = "booksFunction";
	public static final String BOOKS_TOOL_NAME = "booksTool";
	public static final String BOOKS_FUNCTION_DESCRIPTION = "Use this function to execute an SQL query on the database.";
	
	@Bean(BOOKS_FUNCTION_NAME)
	public FunctionCallback booksFunction(BooksService booksService) {
		return FunctionCallbackWrapper.builder(booksService)
				.withName(BOOKS_FUNCTION_NAME).withDescription(BOOKS_FUNCTION_DESCRIPTION)
				.withResponseConverter(response -> response.result().toString())
				.build();
	}
	
	@Bean(BOOKS_TOOL_NAME)
	public OpenAiApi.FunctionTool booksFunctionTool() {
		String jsonTool = "{\n"
				+ "	\"type\": \"object\",\n"
				+ "	\"properties\": {\n"
				+ "		\"query\": {\n"
				+ "			\"type\": \"string\",\n"
				+ "			\"description\": \"Query to run in the database. Must be provided in plain text. The DB schema is defined by this DDL: CREATE TABLE 'SPRING_AI_TEST'.'BOOKS' ('TITLE' VARCHAR2(50 BYTE) NOT NULL ENABLE, 'AUTHOR' VARCHAR2(50 BYTE) NOT NULL ENABLE, PRIMARY KEY ('TITLE'))\" \n"
				+ "		}\n"
				+ "	},\n"
				+ "	\"required\": [\"query\"]\n"
				+ "}";
		
		OpenAiApi.FunctionTool.Function function = 
				new OpenAiApi.FunctionTool.Function(
						BOOKS_FUNCTION_DESCRIPTION, BOOKS_FUNCTION_NAME,
						ModelOptionsUtils.jsonToMap(jsonTool));
		
		return new OpenAiApi.FunctionTool(OpenAiApi.FunctionTool.Type.FUNCTION, function);
	}
	
}

As you can see, when defining the tool we provided the json as input which represents the parameters of the function. In particular, we defined a “query” parameter, telling the LLM model that it must be evaluated with a plain text SQL query.

We provided the DB model via the DLL we used to create the table. This way the LLM will be able to understand its structure and generate consistent SQL queries.

Function

Next we need to implement the Function that implements the logic:

Java

package com.example.aidemo.service;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.ColumnMapRowMapper;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
import org.springframework.stereotype.Service;

import com.example.aidemo.service.BooksService.Request;
import com.example.aidemo.service.BooksService.Response;

@Service
public class BooksService implements Function<Request, Response> {
	
	public record Request(String query) {}
	public record Response(List<Map<String, Object>> result) {}
	
	@Autowired
	protected NamedParameterJdbcTemplate namedJdbcTemplate;
	
	@Override
	public Response apply(Request r) {
		 List<Map<String, Object>> result = 
				 namedJdbcTemplate.query(r.query, new HashMap<>(), new ColumnMapRowMapper());
		 return new Response(result);
	}
	
}

We used the Spring JDBC NamedParameterJdbcTemplate class to execute the query passed as input in the function request.

RestController

The last necessary step is to create a RestController that makes use of the OpenAI model to handle requests:

Java

package com.example.aidemo.controller;

import java.util.List;

import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.api.OpenAiApi.FunctionTool;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

@RestController
public class BooksController {
	
	private final ChatModel chatModel;
	private final OpenAiApi.FunctionTool booksTool;
	
	@Autowired
	public BooksController(OpenAiChatModel chatModel, OpenAiApi.FunctionTool booksTool) {
		this.booksTool = booksTool;
		this.chatModel = chatModel;
	}
	
	@GetMapping("/ai/books")
	public Generation books(@RequestParam(value = "message") String message) {
		new UserMessage(message);
		
		List<FunctionTool> toolsList = List.of(booksTool);
		Prompt prompt = new Prompt(message, OpenAiChatOptions.builder()
				.withTools(toolsList)
				.build());
		
		ChatResponse response = chatModel.call(prompt);
		return response.getResult();
	}

}

For details on the code, refer to the previous articles in the guide.

The final test

To do a test, we will manually insert a couple of records into the database:

SQL

INSERT INTO "SPRING_AI_TEST"."BOOKS" ("TITLE", "AUTHOR") VALUES ('Frankenstein', 'Mary Shelley');
INSERT INTO "SPRING_AI_TEST"."BOOKS" ("TITLE", "AUTHOR") VALUES ('Dracula', 'Bram Stoker');

Next we try to invoke the endpoint we created to ask which books are present in the database:

http://localhost:8080/ai/books?message=Which books are in the database?

http://localhost:8080/ai/books?message=Which books are in the database?

OpenAi should be able to independently generate the SELECT query and return us a response similar to this:

JSON

{
  "metadata": {
    "contentFilterMetadata": null,
    "finishReason": "STOP"
  },
  "output": {
    "messageType": "ASSISTANT",
    "properties": {
      "role": "ASSISTANT",
      "finishReason": "STOP",
      "id": "chatcmpl-9NhfIH5qTVjVLD7SBI0FMseexpKoX"
    },
    "content": "The books in the database are \"Frankenstein\" by Mary Shelley and \"Dracula\" by Bram Stoker.",
    "media": [
      
    ]
  }
}

What we saw was just an example of the potential of the framework. Function calling can be used to implement any type of logic and transforms LLM into a true process orchestrator.