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
This commit is contained in:
2026-05-17 15:26:24 +03:00
parent 8a395225c0
commit 13bc17cfa8
3 changed files with 303 additions and 0 deletions
+142
View File
@@ -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
+87
View File
@@ -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
+74
View File
@@ -27,6 +27,8 @@ import ../fts/engine as fts
import ../vector/engine as vengine import ../vector/engine as vengine
import ../graph/engine as gengine import ../graph/engine as gengine
import ../graph/community as gcomm import ../graph/community as gcomm
import ../ai/chunk as chunkmod
import ../ai/embed as embedmod
type type
IndexEntry* = ref object IndexEntry* = ref object
@@ -71,6 +73,7 @@ type
ftsIndexes*: Table[string, fts.InvertedIndex] # table.col -> FTS index ftsIndexes*: Table[string, fts.InvertedIndex] # table.col -> FTS index
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
txnManager*: TxnManager txnManager*: TxnManager
pendingTxn*: Transaction pendingTxn*: Transaction
onChange*: proc(ev: ChangeEvent) {.closure.} onChange*: proc(ev: ChangeEvent) {.closure.}
@@ -1269,6 +1272,30 @@ proc evalExpr*(expr: IRExpr, row: Table[string, string], ctx: ExecutionContext =
return $(%* outArr) return $(%* outArr)
except: except:
return resultsJson 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": 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()
@@ -1542,6 +1569,53 @@ proc execInsert*(ctx: ExecutionContext, table: string, fields: seq[string], valu
meta[col] = val meta[col] = val
vengine.insert(vecIdx, docId, vec, meta) 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 # Update Graph objects for graph node/edge tables
for graphName, graph in ctx.graphs: for graphName, graph in ctx.graphs:
if table == graphName & "_nodes": if table == graphName & "_nodes":