用 LangChain4j+Ollama 打造 Text-to-SQL AI Agent，数据库想问就问

package com.tsvetkov.db;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.sql.Statement;
public class DatabaseManager {
    // 数据库连接信息（就是之前让你记的那三个！）
    private static final String URL = "jdbc:postgresql://localhost:5432/tts-db";
    private static final String USER = "tts-user";
    private static final String PASSWORD = "tts-pass";
    // 连数据库并初始化表和数据
    public static Connection getInitialConnection() throws SQLException, IOException {
        // 1. 连数据库
        Connection connection = DriverManager.getConnection(URL, USER, PASSWORD);
        System.out.println("成功连到PostgreSQL啦！");
        // 2. 读init.sql文件，执行里面的SQL
        return new String(Files.readAllBytes(Paths.get("src/main/resources/init.sql")));
    }

import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.service.AiServices;
public class AIUtils {
    // 生成AI助手（能帮你转SQL、写答案）
    public static AIAssistant getAIAssistant() {
        return AiServices.builder(AIAssistant.class)
                .chatModel(initOllamaChatModel()) // 连Ollama模型
                .build();
    }
    // 启动Ollama的llama3.2:1b模型
    private static ChatModel initOllamaChatModel() {
        return OllamaChatModel.builder()
                .baseUrl("http://localhost:11434") // Ollama默认地址
                .modelName("llama3.2:1b") // 用的模型名
                .build();
    }
}

import dev.langchain4j.service.UserMessage;
import dev.langchain4j.service.V;
public interface AIAssistant {
    @UserMessage("""
                 You are a senior SQL engineer. Given the database schema and user question below, write a syntactically correct and schema-valid SQL SELECT query.
                 Database Schema (Use only columns and tables listed here)
                 {{schemaDescription}}
                 Rules:
                 Only valid syntax queries - meaning it must start with SELECT
                 Only use tables and columns from the schema above — do not guess
                 Only use SELECT statements (no INSERT, UPDATE, DELETE)
                 Use explicit JOINs, not subqueries unless necessary
                 Add LIMIT 100 to large result sets if not specified in the question
                 Use aggregate functions (COUNT, SUM, etc.) only if the question requires it
                 Return only the SQL query, no explanation, no comments
                 The query must be valid SQL and executable without syntax errors
                 User Question
                 {{question}}
                 """)
    String getQuery(@V("question") String question,  @V("schema") String schemaDescription);
}

import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
// 存字段信息（字段名+字段类型）
record ColumnInfo(String name, String type) {}
public class SchemaAnalyzer {
    // 每个表只取3条数据当例子，太多了AI看不过来
    private static final int EXAMPLES_LIMIT = 3;
    // 生成数据库结构描述（给AI看的“说明书”）
    public static String getSchemaDescription(Connection connection) throws SQLException {
        if (connection == null) {
            throw new IllegalArgumentException("数据库没连上，别瞎传！");
        }
        StringBuilder schemaDesc = new StringBuilder();
        // 1. 先获取所有表名
        List<String> tables = getTables(connection);
        // 2. 逐个表整理结构和示例数据
        for (String table : tables) {
            schemaDesc.append("表名：").append(table).append("\n");
            schemaDesc.append(getTableDetail(connection, table));
        }
        return schemaDesc.toString();
    }
    // 整理单个表的结构（字段名+类型）和示例数据
    private static String getTableDetail(Connection connection, String table) throws SQLException {
        StringBuilder detail = new StringBuilder();
        // 获取表的所有字段
        List<ColumnInfo> columns = getColumns(connection, table);
        for (ColumnInfo col : columns) {
            detail.append("  - ").append(col.name()).append("（类型：").append(col.type()).append("）\n");
        }
        // 获取表的示例数据
        List<List<String>> sampleRows = getSampleRows(connection, table, columns);
        // 把示例数据写成表格格式，AI看得更清楚
        detail.append("示例数据：\n");
        detail.append("  | ");
        for (ColumnInfo col : columns) {
            detail.append(col.name()).append(" | ");
        }
        detail.append("\n");
        for (List<String> row : sampleRows) {
            detail.append("  | ");
            for (String value : row) {
                detail.append(value).append(" | ");
            }
            detail.append("\n");
        }
        return detail.toString();
    }
    // 获取数据库里的所有表名
    private static List<String> getTables(Connection connection) throws SQLException {
        List<String> tables = new ArrayList<>();
        try (ResultSet rs = connection.getMetaData().getTables(null, "public", "%", new String[]{"TABLE"})) {
            while (rs.next()) {
                tables.add(rs.getString("TABLE_NAME"));
            }
        }
        return tables;
    }
    // 获取单个表的所有字段
    private static List<ColumnInfo> getColumns(Connection connection, String table) throws SQLException {
        List<ColumnInfo> columns = new ArrayList<>();
        try (ResultSet rs = connection.getMetaData().getColumns(null, "public", table, "%")) {
            while (rs.next()) {
                columns.add(new ColumnInfo(
                        rs.getString("COLUMN_NAME"),  // 字段名
                        rs.getString("TYPE_NAME")     // 字段类型
                ));
            }
        }
        return columns;
}        
   private static List<List<String>> getSampleRows(Connection connection, String table, List<ColumnInfo> columns) throws SQLException {
        var rows = new ArrayList<List<String>>();
        try (ResultSet rs = connection.createStatement().executeQuery("SELECT * FROM \"" + table + "\" LIMIT " + EXAMPLES_LIMIT)) {
            while (rs.next()) {
                var row = new ArrayList<String>();
                for (ColumnInfo col : columns) {
                    row.add(rs.getString(col.name()));
                }
                rows.add(row);
            }
        }
        return rows;
    }
 }

public static void validateQuery(String sql) {
// 基础校验——只允许SELECT语句
String upperSql = sql.toUpperCase().trim();
if(!upperSql.startsWith("SELECT") ||
   upperSql.contains("DROP") ||  // 不准删表
   upperSql.contains("DELETE") || // 不准删数据
   upperSql.contains("INSERT") || // 不准插数据
   upperSql.contains("UPDATE") || // 不准改数据
   upperSql.contains("ALTER") ||  // 不准改表结构
   upperSql.contains("CREATE")) { // 不准建表
    throw new RuntimeException("SQL有问题，不能执行：\n"+sql);
}
}

private static String getFormattedResultsFromQuery(Connection conn, String sql) throws SQLException {
    DatabaseManager.validateQuery(sql);
    var resultsDescription = new StringBuilder();
    try (Statement stmt = conn.createStatement();
         ResultSet rs = stmt.executeQuery(sql)) {
        int columnCount = rs.getMetaData().getColumnCount();
        // 先拼表头
        for (int i = 1; i <= columnCount; i++) {
            resultsDescription.append(rs.getMetaData().getColumnName(i));
            if (i < columnCount) resultsDescription.append(" | ");
        }
        resultsDescription.append("\n");
        // 再拼数据行
        while (rs.next()) {
            for (int i = 1; i <= columnCount; i++) {
                resultsDescription.append(rs.getString(i));
                if (i < columnCount) resultsDescription.append(" | ");
            }
            resultsDescription.append("\n");
        }
    }
    return resultsDescription.toString();
}

@UserMessage("""
      You are a data analyst. Based on the database schema, the user’s question, the SQL query and the SQL query results, generate a clear, concise, human-readable answer.
      Focus on answering the user's question directly using the data provided — do not describe the SQL or repeat the table structure. Give a straight answer.
      ### Database Schema
      {{schema}}
      ### User Question
      {{question}}
      ### SQL Query
      {{ query }}
      ### SQL Query Results
      {{results}}
      ### Answer
    """)
    String explainAnswer(@V("question") String question, @V("schema") String schema, @V("results") String results);

public static void main(String[] args) throws SQLException, IOException {
        System.out.println("Starting our text-to-sql application...");
        AIAssistant aiAssistant = AIUtils.getAIAssistant();
        try (Connection conn = DatabaseManager.getInitialConnection()) {
            String schemaDescription = SchemaAnalyzer.getSchemaDescription(conn);
            String question = "What are our top 3 most ordered products?";
            String query = aiAssistant.getQuery(question, schemaDescription);
            String formattedQueryResults = getFormattedResultsFromQuery(conn, query);
            String humanAnswer = aiAssistant.explainAnswer(question, schemaDescription, query, formattedQueryResults);
            System.out.println(humanAnswer);
        }
}