From 13bc17cfa851b4e252ae6be125889290268cf714 Mon Sep 17 00:00:00 2001 From: dimgigov Date: Sun, 17 May 2026 15:26:24 +0300 Subject: [PATCH] feat: 10.1.4 Chunking + embedding pipeline - New modules: src/barabadb/ai/chunk.nim (text chunking) and embed.nim (HTTP embedding client) - chunk() SQL function: returns JSON array of chunks with configurable size/overlap - embed_text() SQL function: calls external embedding API (OpenAI/Ollama compatible) - Auto-embedding on INSERT: when VECTOR column is null but TEXT column is populated, generates embeddings via configured embedder - Configurable via env vars: BARADB_EMBED_ENDPOINT, BARADB_EMBED_MODEL, BARADB_EMBED_API_KEY - All 340+ existing tests pass --- src/barabadb/ai/chunk.nim | 142 ++++++++++++++++++++++++++++++++ src/barabadb/ai/embed.nim | 87 +++++++++++++++++++ src/barabadb/query/executor.nim | 74 +++++++++++++++++ 3 files changed, 303 insertions(+) create mode 100644 src/barabadb/ai/chunk.nim create mode 100644 src/barabadb/ai/embed.nim diff --git a/src/barabadb/ai/chunk.nim b/src/barabadb/ai/chunk.nim new file mode 100644 index 0000000..bb5fbcd --- /dev/null +++ b/src/barabadb/ai/chunk.nim @@ -0,0 +1,142 @@ +## Chunking — Text splitting for RAG pipelines +## +## Splits long text into overlapping chunks suitable for embedding. +## Strategies: paragraph, sentence, fixed-size with overlap. + +import std/strutils +import std/sequtils +import std/json + +type + ChunkStrategy* = enum + csParagraph = "paragraph" # Split by double newlines + csSentence = "sentence" # Split by sentence boundaries + csFixed = "fixed" # Fixed-size with overlap + csRecursive = "recursive" # Try paragraph, then sentence, then fixed + + ChunkConfig* = object + maxChunkSize*: int # Max characters per chunk (default 1024) + chunkOverlap*: int # Character overlap between chunks (default 128) + strategy*: ChunkStrategy # Chunking strategy (default recursive) + minChunkSize*: int # Minimum chunk size (default 64) + separators*: seq[string] # Custom separators for recursive splitting + +proc defaultChunkConfig*(): ChunkConfig = + ChunkConfig( + maxChunkSize: 1024, + chunkOverlap: 128, + strategy: csRecursive, + minChunkSize: 64, + separators: @["\n\n", "\n", ". ", "? ", "! ", "; ", ", ", " "], + ) + +proc splitByParagraphs(text: string): seq[string] = + result = @[] + for para in text.split("\n\n"): + let trimmed = para.strip() + if trimmed.len > 0: + result.add(trimmed) + +proc splitBySentences(text: string): seq[string] = + result = @[] + var current = "" + var i = 0 + while i < text.len: + current.add(text[i]) + if text[i] in {'.', '?', '!'}: + if i + 1 < text.len and text[i + 1] == ' ': + inc i + current.add(' ') + let trimmed = current.strip() + if trimmed.len > 0: + result.add(trimmed) + current = "" + inc i + let remaining = current.strip() + if remaining.len > 0: + result.add(remaining) + +proc splitFixed(text: string, chunkSize: int, overlap: int): seq[string] = + result = @[] + if text.len <= chunkSize: + if text.strip().len > 0: + result.add(text.strip()) + return + + var pos = 0 + while pos < text.len: + let endPos = min(pos + chunkSize, text.len) + var chunk = text[pos ..< endPos] + + if endPos < text.len: + var breakPos = chunk.rfind(". ") + if breakPos < 0: + breakPos = chunk.rfind("? ") + if breakPos < 0: + breakPos = chunk.rfind("! ") + if breakPos < 0: + breakPos = chunk.rfind("\n\n") + if breakPos < 0: + breakPos = chunk.rfind("\n") + if breakPos < 0: + breakPos = chunk.rfind(" ") + if breakPos > chunkSize div 4: + chunk = chunk[0 .. breakPos] + pos += breakPos + 1 + else: + pos += chunkSize - overlap + else: + pos = text.len + + let trimmed = chunk.strip() + if trimmed.len > 0: + result.add(trimmed) + +proc chunk*(text: string, config: ChunkConfig = defaultChunkConfig()): seq[string] = + if text.len <= config.minChunkSize: + let trimmed = text.strip() + if trimmed.len > 0: + return @[trimmed] + return @[] + + case config.strategy + of csParagraph: + result = splitByParagraphs(text) + of csSentence: + result = splitBySentences(text) + of csFixed: + result = splitFixed(text, config.maxChunkSize, config.chunkOverlap) + of csRecursive: + # Try paragraph first + var paragraphs = splitByParagraphs(text) + if paragraphs.len > 1: + for para in paragraphs: + if para.len > config.maxChunkSize: + for sentence in splitBySentences(para): + if sentence.len > config.maxChunkSize: + result.add(splitFixed(sentence, config.maxChunkSize, config.chunkOverlap)) + else: + result.add(sentence) + else: + result.add(para) + else: + var sentences = splitBySentences(text) + if sentences.len > 1: + for sentence in sentences: + if sentence.len > config.maxChunkSize: + result.add(splitFixed(sentence, config.maxChunkSize, config.chunkOverlap)) + else: + result.add(sentence) + else: + result = splitFixed(text, config.maxChunkSize, config.chunkOverlap) + + result = result.filterIt(it.len >= config.minChunkSize) + +proc chunkToJson*(text: string, config: ChunkConfig = defaultChunkConfig()): JsonNode = + let chunks = chunk(text, config) + var arr = newJArray() + var idx = 0 + for c in chunks: + arr.add(%*{"index": idx, "text": c, "size": c.len}) + inc idx + return arr diff --git a/src/barabadb/ai/embed.nim b/src/barabadb/ai/embed.nim new file mode 100644 index 0000000..4f85427 --- /dev/null +++ b/src/barabadb/ai/embed.nim @@ -0,0 +1,87 @@ +## Embedding client — calls external embedding APIs +## +## Configurable HTTP client for generating vector embeddings from text. +## Supports OpenAI-compatible and Ollama APIs. + +import std/httpclient +import std/json +import std/strutils +import std/os + +type + EmbedderConfig* = object + endpoint*: string # e.g. "http://localhost:11434/api/embeddings" + model*: string # e.g. "nomic-embed-text" + apiKey*: string # API key (for OpenAI-compatible APIs) + dimensions*: int # Expected embedding dimensions + timeoutMs*: int # Request timeout in ms + enabled*: bool # Whether auto-embedding is enabled + + Embedder* = ref object + config*: EmbedderConfig + +proc defaultEmbedderConfig*(): EmbedderConfig = + EmbedderConfig( + endpoint: getEnv("BARADB_EMBED_ENDPOINT", ""), + model: getEnv("BARADB_EMBED_MODEL", "nomic-embed-text"), + apiKey: getEnv("BARADB_EMBED_API_KEY", ""), + dimensions: 768, + timeoutMs: 30000, + enabled: false, + ) + +proc newEmbedder*(config: EmbedderConfig = defaultEmbedderConfig()): Embedder = + result = Embedder(config: config) + result.config.enabled = config.endpoint.len > 0 + +proc embed*(e: Embedder, text: string): seq[float32] = + result = @[] + if not e.config.enabled: + return + + var client = newHttpClient(timeout = e.config.timeoutMs) + try: + var body = %*{"model": e.config.model, "prompt": text} + if e.config.apiKey.len > 0: + client.headers["Authorization"] = "Bearer " & e.config.apiKey + client.headers["Content-Type"] = "application/json" + + let resp = client.request(e.config.endpoint, httpMethod = HttpPost, body = $body) + let data = parseJson(resp.body) + + if data.hasKey("embedding"): + for val in data["embedding"]: + result.add(float32(val.getFloat())) + elif data.hasKey("data") and data["data"].kind == JArray and data["data"].len > 0: + for val in data["data"][0]["embedding"]: + result.add(float32(val.getFloat())) + except: + discard + finally: + client.close() + +proc embedBatch*(e: Embedder, texts: seq[string]): seq[seq[float32]] = + result = newSeq[seq[float32]](texts.len) + for i, text in texts: + result[i] = e.embed(text) + +proc vectorToJson*(vec: seq[float32]): string = + var parts: seq[string] = @[] + for v in vec: + parts.add($v) + return "[" & parts.join(",") & "]" + +proc jsonToVector*(s: string): seq[float32] = + result = @[] + var cleaned = s.strip() + if cleaned.startsWith("[") and cleaned.endsWith("]"): + cleaned = cleaned[1..^2] + elif cleaned.startsWith("(") and cleaned.endsWith(")"): + cleaned = cleaned[1..^2] + for part in cleaned.split(","): + let p = part.strip() + if p.len > 0: + try: + result.add(parseFloat(p)) + except: + discard diff --git a/src/barabadb/query/executor.nim b/src/barabadb/query/executor.nim index d181d3c..8e1f181 100644 --- a/src/barabadb/query/executor.nim +++ b/src/barabadb/query/executor.nim @@ -27,6 +27,8 @@ import ../fts/engine as fts import ../vector/engine as vengine import ../graph/engine as gengine import ../graph/community as gcomm +import ../ai/chunk as chunkmod +import ../ai/embed as embedmod type IndexEntry* = ref object @@ -71,6 +73,7 @@ type ftsIndexes*: Table[string, fts.InvertedIndex] # table.col -> FTS index 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 txnManager*: TxnManager pendingTxn*: Transaction onChange*: proc(ev: ChangeEvent) {.closure.} @@ -1269,6 +1272,30 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext = return $(%* outArr) except: return resultsJson + of "chunk": + if expr.irFuncArgs.len < 1: return "[]" + let text = evalExpr(expr.irFuncArgs[0], row, ctx) + let maxSize = if expr.irFuncArgs.len >= 2: + try: parseInt(evalExpr(expr.irFuncArgs[1], row, ctx)) except: 1024 + else: 1024 + let overlap = if expr.irFuncArgs.len >= 3: + try: parseInt(evalExpr(expr.irFuncArgs[2], row, ctx)) except: 128 + else: 128 + let cfg = chunkmod.ChunkConfig(maxChunkSize: maxSize, chunkOverlap: overlap, + strategy: chunkmod.csRecursive, minChunkSize: 64) + let chunks = chunkmod.chunk(text, cfg) + var jsonChunks = newJArray() + for i, c in chunks: + jsonChunks.add(%*{"index": i, "text": c, "size": c.len}) + return $(jsonChunks) + of "embed_text": + if expr.irFuncArgs.len < 1: return "[]" + let text = evalExpr(expr.irFuncArgs[0], row, ctx) + if ctx.embedder == nil or not ctx.embedder.config.enabled: + return "[]" + let vec = embedmod.embed(ctx.embedder, text) + if vec.len == 0: return "[]" + return embedmod.vectorToJson(vec) of "datetime": if expr.irFuncArgs.len > 0: let arg = evalExpr(expr.irFuncArgs[0], row, ctx).toLower() @@ -1542,6 +1569,53 @@ proc execInsert*(ctx: ExecutionContext, table: string, fields: seq[string], valu meta[col] = val vengine.insert(vecIdx, docId, vec, meta) + # Auto-embed: if table has VECTOR column with null value but TEXT column + # with content, and embedder is configured, generate embedding + if ctx.embedder != nil and ctx.embedder.config.enabled: + for vecKey in ctx.vectorIndexes.keys: + if not vecKey.startsWith(table & "."): continue + let vecCol = vecKey[table.len + 1..^1] + let vecStr = getValue(rowVals, fields, vecCol) + if vecStr.len == 0 or vecStr == "null" or vecStr == "[]": + var sourceText = "" + for i, f in fields: + if i < rowVals.len and (f == "text" or f == "content" or f == "body"): + sourceText = rowVals[i] + break + if sourceText.len > 0: + let vec = embedmod.embed(ctx.embedder, sourceText) + if vec.len > 0: + let vecStr2 = "[" & vec.mapIt($it).join(",") & "]" + var updateKey = "" + var updateVals: seq[string] = @[] + for i, f in fields: + if i < rowVals.len: + if f == vecCol: + updateVals.add(f & "=" & escapeRowVal(vecStr2)) + elif updateKey.len == 0: + updateKey = f & "=" & escapeRowVal(rowVals[i]) + else: + updateVals.add(f & "=" & escapeRowVal(rowVals[i])) + elif f == vecCol: + updateVals.add(f & "=" & escapeRowVal(vecStr2)) + if updateVals.len > 0: + let fullKey = table & "." & updateKey + let valStr = updateVals.join(",") + if ctx.pendingTxn != nil and ctx.pendingTxn.state == tsActive: + discard ctx.txnManager.write(ctx.pendingTxn, fullKey, cast[seq[byte]](valStr)) + else: + ctx.db.put(fullKey, cast[seq[byte]](valStr)) + var docId: uint64 = 0 + for ch in fullKey: + docId = docId * 31 + uint64(ord(ch)) + var meta = initTable[string, string]() + meta["key"] = fullKey + for col, val in row: + if col.len > 0 and col != "$key" and col != "$value": + meta[col] = val + meta[vecCol] = vecStr2 + vengine.insert(ctx.vectorIndexes[vecKey], docId, vec, meta) + # Update Graph objects for graph node/edge tables for graphName, graph in ctx.graphs: if table == graphName & "_nodes":