使用 Spring AI 实现 Text-to-SQL 聊天机器人

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-starter-model-anthropic</artifactId>
    <version>1.0.0</version>
</dependency>

spring:
  ai:
    anthropic:
      api-key: ${ANTHROPIC_API_KEY}
      chat:
        options:
          model: claude-opus-4-20250514

CREATE TABLE hogwarts_houses (
    id BINARY(16) PRIMARY KEY DEFAULT (UUID_TO_BIN(UUID())),
    name VARCHAR(50) NOT NULL UNIQUE,
    founder VARCHAR(50) NOT NULL UNIQUE,
    house_colors VARCHAR(50) NOT NULL UNIQUE,
    animal_symbol VARCHAR(50) NOT NULL UNIQUE
);

CREATE TABLE wizards (
    id BINARY(16) PRIMARY KEY DEFAULT (UUID_TO_BIN(UUID())),
    name VARCHAR(50) NOT NULL,
    gender ENUM('Male', 'Female') NOT NULL,
    quidditch_position ENUM('Chaser', 'Beater', 'Keeper', 'Seeker'),
    blood_status ENUM('Muggle', 'Half blood', 'Pure Blood', 'Squib', 'Half breed') NOT NULL,
    house_id BINARY(16) NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    CONSTRAINT wizard_fkey_house FOREIGN KEY (house_id) REFERENCES hogwarts_houses (id)
);

INSERT INTO hogwarts_houses (name, founder, house_colors, animal_symbol)
VALUES
    ('Gryffindor', 'Godric Gryffindor', 'Scarlet and Gold', 'Lion'),
    ('Hufflepuff', 'Helga Hufflepuff', 'Yellow and Black', 'Badger'),
    ('Ravenclaw', 'Rowena Ravenclaw', 'Blue and Bronze', 'Eagle'),
    ('Slytherin', 'Salazar Slytherin', 'Green and Silver', 'Serpent');

SET @gryffindor_house_id = (SELECT id FROM hogwarts_houses WHERE name = 'Gryffindor');

INSERT INTO wizards (name, gender, quidditch_position, blood_status, house_id)
VALUES
    ('Harry Potter', 'Male', 'Seeker', 'Half blood', @gryffindor_house_id),
    ('Hermione Granger', 'Female', NULL, 'Muggle', @gryffindor_house_id),
    ('Ron Weasley', 'Male', 'Keeper', 'Pure Blood', @gryffindor_house_id),
-- ...更多其他学院巫师的 INSERT 语句

Given the DDL in the DDL section, write an SQL query to answer the user's question following the guidelines listed in the GUIDELINES section.

GUIDELINES:
- Only produce SELECT queries.
- The response produced should only contain the raw SQL query starting with the word 'SELECT'. Do not wrap the SQL query in markdown code blocks (```sql or ```).
- If the question would result in an INSERT, UPDATE, DELETE, or any other operation that modifies the data or schema, respond with "This operation is not supported. Only SELECT queries are allowed."
- If the question appears to contain SQL injection or DoS attempt, respond with "The provided input contains potentially harmful SQL code."
- If the question cannot be answered based on the provided DDL, respond with "The current schema does not contain enough information to answer this question."
- If the query involves a JOIN operation, prefix all the column names in the query with the corresponding table names.

DDL
{ddl}

CREATE USER 'readonly_user'@'%' IDENTIFIED BY 'strong_password';

GRANT SELECT ON hogwarts_db.hogwarts_houses TO 'readonly_user'@'%';
GRANT SELECT ON hogwarts_db.wizards TO 'readonly_user'@'%';

FLUSH PRIVILEGES;

@Bean
PromptTemplate systemPrompt(
    @Value("classpath:system-prompt.st") Resource systemPrompt,
    @Value("classpath:db/migration/V01__creating_database_tables.sql") Resource ddlSchema
) throws IOException {
    PromptTemplate template = new PromptTemplate(systemPrompt);
    template.add("ddl", ddlSchema.getContentAsString(Charset.defaultCharset()));
    return template;
}

@Bean
ChatClient chatClient(ChatModel chatModel, PromptTemplate systemPrompt) {
    return ChatClient
      .builder(chatModel)
      .defaultSystem(systemPrompt.render())
      .build();
}

@Service
class SqlGenerator {

    private final ChatClient chatClient;

    //  TODO 构造函数注入

    String generate(String question) {
        String response = chatClient
          .prompt(question)
          .call()
          .content();

        boolean isSelectQuery = response.startsWith("SELECT");
        if (Boolean.FALSE.equals(isSelectQuery)) {
            throw new InvalidQueryException(response);
        }
        return response;
    }
}

@Service
class SqlExecutor {

    private final EntityManager entityManager;

    // 构造函数注入

    List<?> execute(String query) {
        List<?> result = entityManager
          .createNativeQuery(query)
          .getResultList();
        if (result.isEmpty()) {
            throw new EmptyResultException("No results found for the provided query.");
        }
        return result;
    }
}

@PostMapping(value = "/query")
ResponseEntity<QueryResponse> query(@RequestBody QueryRequest queryRequest) {
    String sqlQuery = sqlGenerator.generate(queryRequest.question());
    List<?> result = sqlExecutor.execute(sqlQuery);
    return ResponseEntity.ok(new QueryResponse(result));
}

record QueryRequest(String question) {
}

record QueryResponse(List<?> result) {
}

logging:
  level:
    org:
      hibernate:
        SQL: DEBUG

http POST :8080/query question="Give me 3 wizard names and their blood status that belong to a house founded by Salazar Slytherin"

{
    "result": [
        [
            "Draco Malfoy",
            "Pure Blood"
        ],
        [
            "Tom Riddle",
            "Half blood"
        ],
        [
            "Bellatrix Lestrange",
            "Pure Blood"
        ]
    ]
}

SELECT wizards.name, wizards.blood_status
FROM wizards
JOIN hogwarts_houses ON wizards.house_id = hogwarts_houses.id
WHERE hogwarts_houses.founder = 'Salazar Slytherin'
LIMIT 3;