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