refs #735: add few-shot NL2SQL MCP tool

This commit is contained in:
devmrko
2026-07-24 16:21:01 +09:00
parent b738708528
commit caa7d55085
6 changed files with 142 additions and 11 deletions

View File

@@ -78,6 +78,7 @@ public class SelectAiService {
response.put("originalPrompt", normalizedPrompt);
response.put("fewShotStatus", enrichedPrompt.status());
response.put("fewShotExampleCount", enrichedPrompt.exampleCount());
addFewShotExamples(response, enrichedPrompt.examples());
response.put("generatedSql", normalizedSql);
response.put("execution", "READ_ONLY_EXECUTED");
response.put("rowCount", execution.items().size());
@@ -89,6 +90,21 @@ public class SelectAiService {
return response;
}
private void addFewShotExamples(ObjectNode response, List<QaVectorService.VectorExample> examples) {
ArrayNode items = response.putArray("fewShotExamples");
for (QaVectorService.VectorExample example : examples) {
ObjectNode item = items.addObject();
item.put("exampleId", example.exampleId());
item.put("question", example.question());
item.put("answerSql", example.answerSql());
if (example.answer() != null) {
item.put("answer", example.answer());
}
item.put("embeddingModel", example.embeddingModel());
item.put("cosineDistance", example.cosineDistance());
}
}
/** Returns the prompt Select AI assembled for SQL generation without executing generated SQL. */
public JsonNode generatePrompt(String bearerToken, String prompt) {
requireActiveToken(bearerToken);
@@ -171,20 +187,21 @@ public class SelectAiService {
String prompt
) {
if (!fewShotEnabled(selectAi) || qaVectorService == null) {
return new EnrichedPrompt(prompt, "DISABLED", 0);
return new EnrichedPrompt(prompt, "DISABLED", 0, List.of());
}
try {
List<QaVectorService.VectorExample> examples = qaVectorService
.search(bearerToken, prompt, fewShotTopK(selectAi))
.examples();
if (examples.isEmpty()) {
return new EnrichedPrompt(composePolicyPrompt(prompt), "NO_MATCH", 0);
return new EnrichedPrompt(composePolicyPrompt(prompt), "NO_MATCH", 0, List.of());
}
return new EnrichedPrompt(
composeFewShotPrompt(prompt, examples), "APPLIED", Math.min(examples.size(), MAX_FEW_SHOT_EXAMPLES));
composeFewShotPrompt(prompt, examples), "APPLIED", Math.min(examples.size(), MAX_FEW_SHOT_EXAMPLES),
examples.subList(0, Math.min(examples.size(), MAX_FEW_SHOT_EXAMPLES)));
} catch (Exception ignored) {
// Vector retrieval is an optional prompt aid; preserve the normal Text2SQL path on failure.
return new EnrichedPrompt(composePolicyPrompt(prompt), "UNAVAILABLE", 0);
return new EnrichedPrompt(composePolicyPrompt(prompt), "UNAVAILABLE", 0, List.of());
}
}
@@ -325,5 +342,6 @@ public class SelectAiService {
private record QueryExecution(ArrayNode items, boolean truncated) {}
private record EnrichedPrompt(String prompt, String status, int exampleCount) {}
private record EnrichedPrompt(
String prompt, String status, int exampleCount, List<QaVectorService.VectorExample> examples) {}
}