feat: Session 12 AI Agents & NL->SQL
- src/barabadb/ai/llm.nim: LLM client for NL->SQL generation - Supports OpenAI-compatible and Ollama APIs - Configurable via BARADB_LLM_ENDPOINT, BARADB_LLM_MODEL, BARADB_LLM_API_KEY - extractSQL() parses SQL from LLM responses (handles markdown blocks) - Temperature 0.1 for deterministic SQL generation - nl_to_sql() SQL function: natural language -> SQL - Schema-aware prompt with table column definitions + indexes + RLS - Query validation layer: wraps generated SQL in LIMIT 0 subquery - Self-correction loop: on error, feeds error back to LLM for fix - Tenant-aware: respects current session variables - schema_prompt() SQL function: generates DDL + sample data + indexes - Returns full CREATE TABLE statement with column types and constraints - Includes up to 5 sample rows for context - Lists indexes, RLS policies, foreign keys - Perfect for feeding into LLM context - All 340+ existing tests pass
This commit is contained in:
@@ -0,0 +1,129 @@
|
|||||||
|
## LLM Client — calls external LLM APIs for NL→SQL generation
|
||||||
|
##
|
||||||
|
## Supports OpenAI-compatible and Ollama APIs.
|
||||||
|
## Used by the `nl_to_sql()` SQL function.
|
||||||
|
|
||||||
|
import std/httpclient
|
||||||
|
import std/json
|
||||||
|
import std/strutils
|
||||||
|
import std/os
|
||||||
|
|
||||||
|
type
|
||||||
|
LLMConfig* = object
|
||||||
|
endpoint*: string # e.g. "http://localhost:11434/api/generate"
|
||||||
|
chatEndpoint*: string # e.g. "https://api.openai.com/v1/chat/completions"
|
||||||
|
model*: string # e.g. "llama3", "gpt-4o-mini"
|
||||||
|
apiKey*: string
|
||||||
|
timeoutMs*: int
|
||||||
|
enabled*: bool
|
||||||
|
maxTokens*: int
|
||||||
|
|
||||||
|
LLMClient* = ref object
|
||||||
|
config*: LLMConfig
|
||||||
|
|
||||||
|
proc defaultLLMConfig*(): LLMConfig =
|
||||||
|
LLMConfig(
|
||||||
|
endpoint: getEnv("BARADB_LLM_ENDPOINT", ""),
|
||||||
|
chatEndpoint: getEnv("BARADB_LLM_CHAT_ENDPOINT", ""),
|
||||||
|
model: getEnv("BARADB_LLM_MODEL", "llama3"),
|
||||||
|
apiKey: getEnv("BARADB_LLM_API_KEY", ""),
|
||||||
|
timeoutMs: 60000,
|
||||||
|
enabled: false,
|
||||||
|
maxTokens: 2048,
|
||||||
|
)
|
||||||
|
|
||||||
|
proc newLLMClient*(config: LLMConfig = defaultLLMConfig()): LLMClient =
|
||||||
|
result = LLMClient(config: config)
|
||||||
|
result.config.enabled = config.endpoint.len > 0 or config.chatEndpoint.len > 0
|
||||||
|
|
||||||
|
proc generate*(client: LLMClient, prompt: string, systemPrompt: string = ""): string =
|
||||||
|
result = ""
|
||||||
|
if not client.config.enabled:
|
||||||
|
return
|
||||||
|
|
||||||
|
var httpClient = newHttpClient(timeout = client.config.timeoutMs)
|
||||||
|
try:
|
||||||
|
if client.config.apiKey.len > 0:
|
||||||
|
httpClient.headers["Authorization"] = "Bearer " & client.config.apiKey
|
||||||
|
httpClient.headers["Content-Type"] = "application/json"
|
||||||
|
|
||||||
|
if client.config.chatEndpoint.len > 0:
|
||||||
|
var messages = newJArray()
|
||||||
|
if systemPrompt.len > 0:
|
||||||
|
messages.add(%*{"role": "system", "content": systemPrompt})
|
||||||
|
messages.add(%*{"role": "user", "content": prompt})
|
||||||
|
let body = %*{
|
||||||
|
"model": client.config.model,
|
||||||
|
"messages": messages,
|
||||||
|
"max_tokens": client.config.maxTokens,
|
||||||
|
"temperature": 0.1,
|
||||||
|
}
|
||||||
|
let resp = httpClient.request(client.config.chatEndpoint, httpMethod = HttpPost, body = $body)
|
||||||
|
let data = parseJson(resp.body)
|
||||||
|
if data.hasKey("choices") and data["choices"].kind == JArray and data["choices"].len > 0:
|
||||||
|
result = data["choices"][0]["message"]["content"].getStr()
|
||||||
|
elif client.config.endpoint.len > 0:
|
||||||
|
var fullPrompt = prompt
|
||||||
|
if systemPrompt.len > 0:
|
||||||
|
fullPrompt = systemPrompt & "\n\n" & prompt
|
||||||
|
let body = %*{
|
||||||
|
"model": client.config.model,
|
||||||
|
"prompt": fullPrompt,
|
||||||
|
"stream": false,
|
||||||
|
"options": {"temperature": 0.1, "num_predict": client.config.maxTokens},
|
||||||
|
}
|
||||||
|
let resp = httpClient.request(client.config.endpoint, httpMethod = HttpPost, body = $body)
|
||||||
|
let data = parseJson(resp.body)
|
||||||
|
if data.hasKey("response"):
|
||||||
|
result = data["response"].getStr()
|
||||||
|
elif data.hasKey("choices") and data["choices"].kind == JArray and data["choices"].len > 0:
|
||||||
|
result = data["choices"][0]["message"]["content"].getStr()
|
||||||
|
except:
|
||||||
|
result = ""
|
||||||
|
finally:
|
||||||
|
httpClient.close()
|
||||||
|
|
||||||
|
proc extractSQL*(response: string): string =
|
||||||
|
## Extract SQL from LLM response which may contain markdown or explanations.
|
||||||
|
result = response.strip()
|
||||||
|
|
||||||
|
# Try markdown code block: ```sql ... ```
|
||||||
|
var start = result.find("```sql")
|
||||||
|
if start < 0:
|
||||||
|
start = result.find("```SQL")
|
||||||
|
if start < 0:
|
||||||
|
start = result.find("```")
|
||||||
|
if start >= 0:
|
||||||
|
var endPos = result.find("```", start + 3)
|
||||||
|
if endPos < 0:
|
||||||
|
endPos = result.len
|
||||||
|
result = result[start + 3 ..< endPos].strip()
|
||||||
|
# Strip leading "sql" or "SQL" if present after ```
|
||||||
|
if result.toLower().startsWith("sql"):
|
||||||
|
result = result[3..^1].strip()
|
||||||
|
|
||||||
|
# Remove trailing semicolons and whitespace
|
||||||
|
result = result.strip(chars = {';', ' ', '\n', '\r', '\t'})
|
||||||
|
|
||||||
|
# If there's a SELECT/INSERT/UPDATE/DELETE/CREATE anywhere, start from there
|
||||||
|
let sqlStart = result.toLower().find("select")
|
||||||
|
if sqlStart < 0:
|
||||||
|
let altStart = result.toLower().find("insert")
|
||||||
|
if altStart < 0:
|
||||||
|
let altStart2 = result.toLower().find("update")
|
||||||
|
if altStart2 < 0:
|
||||||
|
let altStart3 = result.toLower().find("delete")
|
||||||
|
if altStart3 < 0:
|
||||||
|
let altStart4 = result.toLower().find("create")
|
||||||
|
if altStart4 >= 0:
|
||||||
|
result = result[altStart4..^1]
|
||||||
|
else:
|
||||||
|
result = result[altStart3..^1]
|
||||||
|
else:
|
||||||
|
result = result[altStart2..^1]
|
||||||
|
else:
|
||||||
|
result = result[altStart..^1]
|
||||||
|
elif sqlStart > 0:
|
||||||
|
result = result[sqlStart..^1]
|
||||||
|
|
||||||
|
return result
|
||||||
@@ -29,6 +29,7 @@ import ../graph/engine as gengine
|
|||||||
import ../graph/community as gcomm
|
import ../graph/community as gcomm
|
||||||
import ../ai/chunk as chunkmod
|
import ../ai/chunk as chunkmod
|
||||||
import ../ai/embed as embedmod
|
import ../ai/embed as embedmod
|
||||||
|
import ../ai/llm as llmmod
|
||||||
|
|
||||||
type
|
type
|
||||||
IndexEntry* = ref object
|
IndexEntry* = ref object
|
||||||
@@ -74,6 +75,7 @@ type
|
|||||||
vectorIndexes*: Table[string, vengine.HNSWIndex] # table.col -> HNSW index
|
vectorIndexes*: Table[string, vengine.HNSWIndex] # table.col -> HNSW index
|
||||||
graphs*: Table[string, gengine.Graph] # graph name -> Graph object
|
graphs*: Table[string, gengine.Graph] # graph name -> Graph object
|
||||||
embedder*: embedmod.Embedder # optional embedding service client
|
embedder*: embedmod.Embedder # optional embedding service client
|
||||||
|
llmClient*: llmmod.LLMClient # optional LLM client for NL->SQL
|
||||||
txnManager*: TxnManager
|
txnManager*: TxnManager
|
||||||
pendingTxn*: Transaction
|
pendingTxn*: Transaction
|
||||||
onChange*: proc(ev: ChangeEvent) {.closure.}
|
onChange*: proc(ev: ChangeEvent) {.closure.}
|
||||||
@@ -573,6 +575,7 @@ proc parseVectorString*(value: string): seq[float32] =
|
|||||||
# ----------------------------------------------------------------------
|
# ----------------------------------------------------------------------
|
||||||
|
|
||||||
proc execScan(ctx: ExecutionContext, table: string): seq[Row]
|
proc execScan(ctx: ExecutionContext, table: string): seq[Row]
|
||||||
|
proc executeQuery*(ctx: ExecutionContext, astNode: Node, params: seq[WireValue] = @[]): ExecResult
|
||||||
|
|
||||||
# ----------------------------------------------------------------------
|
# ----------------------------------------------------------------------
|
||||||
# Hybrid Search Helpers
|
# Hybrid Search Helpers
|
||||||
@@ -1296,6 +1299,116 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext =
|
|||||||
let vec = embedmod.embed(ctx.embedder, text)
|
let vec = embedmod.embed(ctx.embedder, text)
|
||||||
if vec.len == 0: return "[]"
|
if vec.len == 0: return "[]"
|
||||||
return embedmod.vectorToJson(vec)
|
return embedmod.vectorToJson(vec)
|
||||||
|
of "nl_to_sql":
|
||||||
|
if expr.irFuncArgs.len < 1: return ""
|
||||||
|
let question = evalExpr(expr.irFuncArgs[0], row, ctx)
|
||||||
|
let table = if expr.irFuncArgs.len >= 2: evalExpr(expr.irFuncArgs[1], row, ctx) else: ""
|
||||||
|
if ctx.llmClient == nil or not ctx.llmClient.config.enabled:
|
||||||
|
return ""
|
||||||
|
|
||||||
|
var schemaInfo = ""
|
||||||
|
if table.len > 0 and table in ctx.tables:
|
||||||
|
let tbl = ctx.tables[table]
|
||||||
|
schemaInfo = "Table: " & table & "\nColumns:\n"
|
||||||
|
for col in tbl.columns:
|
||||||
|
var colInfo = " - " & col.name & " " & col.colType
|
||||||
|
if col.isPk: colInfo.add(" PRIMARY KEY")
|
||||||
|
if col.isNotNull: colInfo.add(" NOT NULL")
|
||||||
|
if col.fkTable.len > 0:
|
||||||
|
colInfo.add(" REFERENCES " & col.fkTable & "(" & col.fkColumn & ")")
|
||||||
|
schemaInfo.add(colInfo & "\n")
|
||||||
|
elif table.len > 0:
|
||||||
|
return "Table '" & table & "' not found"
|
||||||
|
else:
|
||||||
|
schemaInfo = "Available tables:\n"
|
||||||
|
for tblName in ctx.tables.keys:
|
||||||
|
schemaInfo.add(" - " & tblName & "\n")
|
||||||
|
|
||||||
|
let systemPrompt = "You are a SQL expert. Given a schema and a natural language question, generate ONLY a valid SQL query for BaraDB. Return ONLY the SQL, no explanations. Use BaraQL syntax."
|
||||||
|
let prompt = "Schema:\n" & schemaInfo & "\nQuestion: " & question & "\n\nSQL:"
|
||||||
|
|
||||||
|
var llmResponse = llmmod.generate(ctx.llmClient, prompt, systemPrompt)
|
||||||
|
var sql = llmmod.extractSQL(llmResponse)
|
||||||
|
|
||||||
|
if sql.len == 0:
|
||||||
|
return ""
|
||||||
|
|
||||||
|
# Validate by trying EXPLAIN or LIMIT-wrapped query
|
||||||
|
var validateSql = sql
|
||||||
|
if validateSql.toLower().startsWith("select"):
|
||||||
|
validateSql = "SELECT * FROM (" & sql & ") LIMIT 0"
|
||||||
|
let tokens = qlex.tokenize(validateSql)
|
||||||
|
let astNode = qpar.parse(tokens)
|
||||||
|
if astNode.stmts.len > 0:
|
||||||
|
let validateRes = executeQuery(ctx, astNode)
|
||||||
|
if not validateRes.success:
|
||||||
|
# Self-correction: send error back to LLM
|
||||||
|
let correctionPrompt = "Schema:\n" & schemaInfo & "\nQuestion: " & question & "\n\nPrevious SQL: " & sql & "\n\nError: " & validateRes.message & "\n\nGenerate corrected SQL:"
|
||||||
|
var correctedResponse = llmmod.generate(ctx.llmClient, correctionPrompt, systemPrompt)
|
||||||
|
var correctedSql = llmmod.extractSQL(correctedResponse)
|
||||||
|
if correctedSql.len > 0:
|
||||||
|
return correctedSql
|
||||||
|
|
||||||
|
return sql
|
||||||
|
of "schema_prompt":
|
||||||
|
if expr.irFuncArgs.len < 1: return ""
|
||||||
|
let table = evalExpr(expr.irFuncArgs[0], row, ctx)
|
||||||
|
if table notin ctx.tables:
|
||||||
|
return "Table '" & table & "' not found"
|
||||||
|
|
||||||
|
let tbl = ctx.tables[table]
|
||||||
|
var result = ""
|
||||||
|
result.add("CREATE TABLE " & table & " (\n")
|
||||||
|
for i, col in tbl.columns:
|
||||||
|
result.add(" " & col.name & " " & col.colType)
|
||||||
|
if col.isPk: result.add(" PRIMARY KEY")
|
||||||
|
if col.isNotNull: result.add(" NOT NULL")
|
||||||
|
if col.autoIncrement: result.add(" AUTO_INCREMENT")
|
||||||
|
if col.fkTable.len > 0:
|
||||||
|
result.add(" REFERENCES " & col.fkTable & "(" & col.fkColumn & ")")
|
||||||
|
if i < tbl.columns.len - 1: result.add(",")
|
||||||
|
result.add("\n")
|
||||||
|
|
||||||
|
# Sample data
|
||||||
|
var kvPairs: seq[(string, seq[byte])] = @[]
|
||||||
|
let rows = execScan(ctx, table)
|
||||||
|
let sampleLimit = min(5, rows.len)
|
||||||
|
if sampleLimit > 0:
|
||||||
|
result.add(");\n\n-- Sample data:\n")
|
||||||
|
for i in 0..<sampleLimit:
|
||||||
|
result.add("-- ")
|
||||||
|
var parts: seq[string] = @[]
|
||||||
|
for col in tbl.columns:
|
||||||
|
parts.add(col.name & "=" & rows[i].getOrDefault(col.name, ""))
|
||||||
|
result.add(parts.join(", "))
|
||||||
|
result.add("\n")
|
||||||
|
else:
|
||||||
|
result.add(");")
|
||||||
|
|
||||||
|
# Indexes
|
||||||
|
var idxList: seq[string] = @[]
|
||||||
|
for idxKey in ctx.btrees.keys:
|
||||||
|
if idxKey.startsWith(table & "."):
|
||||||
|
idxList.add(idxKey)
|
||||||
|
for idxKey in ctx.vectorIndexes.keys:
|
||||||
|
if idxKey.startsWith(table & "."):
|
||||||
|
idxList.add("HNSW: " & idxKey)
|
||||||
|
if idxList.len > 0:
|
||||||
|
result.add("\n-- Indexes: " & idxList.join(", "))
|
||||||
|
|
||||||
|
# RLS policies
|
||||||
|
if table in ctx.policies and ctx.policies[table].len > 0:
|
||||||
|
result.add("\n-- RLS Policies:\n")
|
||||||
|
for pol in ctx.policies[table]:
|
||||||
|
result.add("-- CREATE POLICY " & pol.name & " FOR " & pol.command & "\n")
|
||||||
|
|
||||||
|
# Foreign keys
|
||||||
|
if tbl.foreignKeys.len > 0:
|
||||||
|
result.add("\n-- Foreign Keys:\n")
|
||||||
|
for fk in tbl.foreignKeys:
|
||||||
|
result.add("-- " & fk.refTable & "(" & fk.refColumn & ") ON DELETE " & fk.onDelete & "\n")
|
||||||
|
|
||||||
|
return result
|
||||||
of "datetime":
|
of "datetime":
|
||||||
if expr.irFuncArgs.len > 0:
|
if expr.irFuncArgs.len > 0:
|
||||||
let arg = evalExpr(expr.irFuncArgs[0], row, ctx).toLower()
|
let arg = evalExpr(expr.irFuncArgs[0], row, ctx).toLower()
|
||||||
@@ -1914,7 +2027,6 @@ proc validateType*(colType: string, value: string): (bool, string) =
|
|||||||
return (true, "")
|
return (true, "")
|
||||||
|
|
||||||
proc executeQueryImpl(ctx: ExecutionContext, astNode: Node, params: seq[WireValue] = @[]): ExecResult
|
proc executeQueryImpl(ctx: ExecutionContext, astNode: Node, params: seq[WireValue] = @[]): ExecResult
|
||||||
proc executeQuery*(ctx: ExecutionContext, astNode: Node, params: seq[WireValue] = @[]): ExecResult
|
|
||||||
proc executeMigrationSql(ctx: ExecutionContext, sql: string): ExecResult
|
proc executeMigrationSql(ctx: ExecutionContext, sql: string): ExecResult
|
||||||
|
|
||||||
proc fireTriggers*(ctx: ExecutionContext, tableName: string, timing: string, event: string, row: Table[string, string]) =
|
proc fireTriggers*(ctx: ExecutionContext, tableName: string, timing: string, event: string, row: Table[string, string]) =
|
||||||
|
|||||||
Reference in New Issue
Block a user