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:
@@ -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
|
||||
@@ -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
|
||||
@@ -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":
|
||||
|
||||
Reference in New Issue
Block a user