Spring AI - Eseguire query su database con OpenAI

Nella nostra guida sul framework Spring AI abbiamo già discusso su come implementare sia le chiamate alle funzioni che ai tool. Adesso è giunto il momento di vedere come usare quanto appreso per fare in modo che il modello LLM possa autonomamente generare ed eseguire query su database.

Come avrete sicuramente intuito leggendo gli articoli precedenti, la possibilità di far decidere ad un modello LLM quando sia il momento giusto di invocare una funzione apre la strada ad infinite possibilità e costituisce una rivoluzione nel campo dell’AI. Inoltre il modello non è solo in grado di decidere quando fare la chiamata ma è anche capace di passare in input le variabili necessarie per farla funzionare.

Andiamo dunque più nel dettaglio con il nostro tutorial pratico.

Tutorial query su database con Spring AI
Predisposizione del DB
Configurazione Spring per accesso al DB
Implementazione java con Spring AI ed OpenAI
Configuration
Function
RestController
Il test finale

Tutorial query su database con Spring AI

Predisposizione del DB

Il primo step del nostro tutorial prevede l’allestimento di un database in locale con una semplice tabella che conterrà un elenco di libri. Come database useremo Oracle XE che abbiamo già visto come avviare su docker nel nostro precedente articolo.

Una volta avviato il DB, creiamo un’utenza e una tabella eseguendo questa 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")
);

Abbiamo quindi creato un’utenza “SPRING_AI_TEST” con password “SPRING_AI_PWD” e una tabella “BOOKS” con due colonne che rappresentano il titolo e l’autore del libro.

Nel caso di errore “ORA-65096” durante l’esecuzione dello script, è sufficiente eseguire prima la seguente istruzione:

SQL

ALTER SESSION SET "_ORACLE_SCRIPT"=TRUE;

Configurazione Spring per accesso al DB

Per essere rapidi partiamo dal nostro tutorial di esempio su come creare un progetto da zero usando Spring AI e OpenAI e andiamo ad aggiungere quanto necessario per creare la connessione col DB.

Per prima cosa importiamo nel pom le dipendenze da Spring JDBC e dal driver Oracle:

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>

E in seguito è necessario impostare nell’application.properties le configurazioni del datasource:

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

Implementazione java con Spring AI ed OpenAI

Configuration

Useremo quanto già imparato in merito all’uso dei FunctionTool in Spring AI per definire un tool che funga da connettore con il nostro database. Per prima cosa è necessario definire una classe di tipo Configuration per definire la funzione e il relativo 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);
	}
	
}

Come si nota, alla definizione del tool abbiamo fornito in input il json che rappresenta i parametri della funzione. In particolare abbiamo definito un parametro “query” dicendo al modello LLM che deve essere valorizzato con una query SQL in plain text.

Abbiamo fornito il modello del DB mediante la DLL che abbiamo usato per creare la tabella. In questo modo l’LLM sarà in grado di comprenderne la struttura e generare query SQL consistenti.

Function

Successivamente dobbiamo implementare la Function che implementa la logica:

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);
	}
	
}

Abbiamo usato la classe NamedParameterJdbcTemplate di Spring JDBC per eseguire la query passata in input nella request della funzione.

RestController

L’ultimo passaggio necessario è la creazione di un RestController che faccia uso del modello OpenAI per gestire le richieste:

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();
	}

}

Per dettagli sul codice rimando agli articoli precedenti della guida.

Il test finale

Per fare un test, inseriremo manualmente un paio di record all’interno del 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');

Successivamente proviamo ad invocare l’endpoint che abbiamo creato per chiedergli quali libri sono presenti nel database:

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

OpenAi dovrebbe essere in grado di generare autonomamente la query di SELECT e restituirci una risposta simile a questa:

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": [
      
    ]
  }
}

Quanto visto era solo un esempio delle potenzialità del framework. Il function calling può essere usato per implementare qualsiasi tipo logica e trasforma LLM in un vero e proprio orchestratore di processo.