From e23b1d61d2d0be4b58aa703d9ebf8dba093d27f6 Mon Sep 17 00:00:00 2001 From: dimgigov Date: Sun, 17 May 2026 15:38:14 +0300 Subject: [PATCH] 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 --- src/barabadb/ai/llm.nim | 129 ++++++++++++++++++++++++++++++++ src/barabadb/query/executor.nim | 114 +++++++++++++++++++++++++++- 2 files changed, 242 insertions(+), 1 deletion(-) create mode 100644 src/barabadb/ai/llm.nim diff --git a/src/barabadb/ai/llm.nim b/src/barabadb/ai/llm.nim new file mode 100644 index 0000000..14a4618 --- /dev/null +++ b/src/barabadb/ai/llm.nim @@ -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 diff --git a/src/barabadb/query/executor.nim b/src/barabadb/query/executor.nim index 8e1f181..3109e0b 100644 --- a/src/barabadb/query/executor.nim +++ b/src/barabadb/query/executor.nim @@ -29,6 +29,7 @@ import ../graph/engine as gengine import ../graph/community as gcomm import ../ai/chunk as chunkmod import ../ai/embed as embedmod +import ../ai/llm as llmmod type IndexEntry* = ref object @@ -74,6 +75,7 @@ type vectorIndexes*: Table[string, vengine.HNSWIndex] # table.col -> HNSW index graphs*: Table[string, gengine.Graph] # graph name -> Graph object embedder*: embedmod.Embedder # optional embedding service client + llmClient*: llmmod.LLMClient # optional LLM client for NL->SQL txnManager*: TxnManager pendingTxn*: Transaction onChange*: proc(ev: ChangeEvent) {.closure.} @@ -573,6 +575,7 @@ proc parseVectorString*(value: string): seq[float32] = # ---------------------------------------------------------------------- proc execScan(ctx: ExecutionContext, table: string): seq[Row] +proc executeQuery*(ctx: ExecutionContext, astNode: Node, params: seq[WireValue] = @[]): ExecResult # ---------------------------------------------------------------------- # Hybrid Search Helpers @@ -1296,6 +1299,116 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext = let vec = embedmod.embed(ctx.embedder, text) if vec.len == 0: return "[]" 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.. 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": if expr.irFuncArgs.len > 0: let arg = evalExpr(expr.irFuncArgs[0], row, ctx).toLower() @@ -1914,7 +2027,6 @@ proc validateType*(colType: string, value: string): (bool, string) = return (true, "") 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 fireTriggers*(ctx: ExecutionContext, tableName: string, timing: string, event: string, row: Table[string, string]) =