refs #735: add few-shot NL2SQL MCP tool
This commit is contained in:
@@ -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) {}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user