feat: full-stack search engine — FTS5, vector search, hybrid merge, embedding cache, chunker

The Full Stack (All Custom):
- Vector DB: embeddings stored as BLOB, cosine similarity in pure Rust
- Keyword Search: FTS5 virtual tables with BM25 scoring + auto-sync triggers
- Hybrid Merge: weighted fusion of vector + keyword results (configurable weights)
- Embeddings: provider abstraction (OpenAI, custom URL, noop fallback)
- Chunking: line-based markdown chunker with heading preservation
- Caching: embedding_cache table with LRU eviction
- Safe Reindex: rebuild FTS5 + re-embed missing vectors

New modules:
- src/memory/embeddings.rs — EmbeddingProvider trait + OpenAI + Noop + factory
- src/memory/vector.rs — cosine similarity, vec↔bytes, ScoredResult, hybrid_merge
- src/memory/chunker.rs — markdown-aware document splitting

Upgraded:
- src/memory/sqlite.rs — FTS5 schema, embedding column, hybrid recall, cache, reindex
- src/config/schema.rs — MemoryConfig expanded with embedding/search settings
- All callers updated to pass api_key for embedding provider

739 tests passing, 0 clippy warnings (Rust 1.93.1), cargo-deny clean
This commit is contained in:
argenis de la rosa 2026-02-14 00:00:23 -05:00
parent 4fceba0740
commit 0e7f501fd6
10 changed files with 1423 additions and 96 deletions

190
src/memory/embeddings.rs Normal file
View file

@ -0,0 +1,190 @@
use async_trait::async_trait;
/// Trait for embedding providers — convert text to vectors
#[async_trait]
pub trait EmbeddingProvider: Send + Sync {
/// Provider name
fn name(&self) -> &str;
/// Embedding dimensions
fn dimensions(&self) -> usize;
/// Embed a batch of texts into vectors
async fn embed(&self, texts: &[&str]) -> anyhow::Result<Vec<Vec<f32>>>;
/// Embed a single text
async fn embed_one(&self, text: &str) -> anyhow::Result<Vec<f32>> {
let mut results = self.embed(&[text]).await?;
results
.pop()
.ok_or_else(|| anyhow::anyhow!("Empty embedding result"))
}
}
// ── Noop provider (keyword-only fallback) ────────────────────
pub struct NoopEmbedding;
#[async_trait]
impl EmbeddingProvider for NoopEmbedding {
fn name(&self) -> &str {
"none"
}
fn dimensions(&self) -> usize {
0
}
async fn embed(&self, _texts: &[&str]) -> anyhow::Result<Vec<Vec<f32>>> {
Ok(Vec::new())
}
}
// ── OpenAI-compatible embedding provider ─────────────────────
pub struct OpenAiEmbedding {
client: reqwest::Client,
base_url: String,
api_key: String,
model: String,
dims: usize,
}
impl OpenAiEmbedding {
pub fn new(base_url: &str, api_key: &str, model: &str, dims: usize) -> Self {
Self {
client: reqwest::Client::new(),
base_url: base_url.trim_end_matches('/').to_string(),
api_key: api_key.to_string(),
model: model.to_string(),
dims,
}
}
}
#[async_trait]
impl EmbeddingProvider for OpenAiEmbedding {
fn name(&self) -> &str {
"openai"
}
fn dimensions(&self) -> usize {
self.dims
}
async fn embed(&self, texts: &[&str]) -> anyhow::Result<Vec<Vec<f32>>> {
if texts.is_empty() {
return Ok(Vec::new());
}
let body = serde_json::json!({
"model": self.model,
"input": texts,
});
let resp = self
.client
.post(format!("{}/v1/embeddings", self.base_url))
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&body)
.send()
.await?;
if !resp.status().is_success() {
let status = resp.status();
let text = resp.text().await.unwrap_or_default();
anyhow::bail!("Embedding API error {status}: {text}");
}
let json: serde_json::Value = resp.json().await?;
let data = json
.get("data")
.and_then(|d| d.as_array())
.ok_or_else(|| anyhow::anyhow!("Invalid embedding response: missing 'data'"))?;
let mut embeddings = Vec::with_capacity(data.len());
for item in data {
let embedding = item
.get("embedding")
.and_then(|e| e.as_array())
.ok_or_else(|| anyhow::anyhow!("Invalid embedding item"))?;
#[allow(clippy::cast_possible_truncation)]
let vec: Vec<f32> = embedding
.iter()
.filter_map(|v| v.as_f64().map(|f| f as f32))
.collect();
embeddings.push(vec);
}
Ok(embeddings)
}
}
// ── Factory ──────────────────────────────────────────────────
pub fn create_embedding_provider(
provider: &str,
api_key: Option<&str>,
model: &str,
dims: usize,
) -> Box<dyn EmbeddingProvider> {
match provider {
"openai" => {
let key = api_key.unwrap_or("");
Box::new(OpenAiEmbedding::new(
"https://api.openai.com",
key,
model,
dims,
))
}
name if name.starts_with("custom:") => {
let base_url = name.strip_prefix("custom:").unwrap_or("");
let key = api_key.unwrap_or("");
Box::new(OpenAiEmbedding::new(base_url, key, model, dims))
}
_ => Box::new(NoopEmbedding),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn noop_name() {
let p = NoopEmbedding;
assert_eq!(p.name(), "none");
assert_eq!(p.dimensions(), 0);
}
#[tokio::test]
async fn noop_embed_returns_empty() {
let p = NoopEmbedding;
let result = p.embed(&["hello"]).await.unwrap();
assert!(result.is_empty());
}
#[test]
fn factory_none() {
let p = create_embedding_provider("none", None, "model", 1536);
assert_eq!(p.name(), "none");
}
#[test]
fn factory_openai() {
let p = create_embedding_provider("openai", Some("key"), "text-embedding-3-small", 1536);
assert_eq!(p.name(), "openai");
assert_eq!(p.dimensions(), 1536);
}
#[test]
fn factory_custom_url() {
let p = create_embedding_provider("custom:http://localhost:1234", None, "model", 768);
assert_eq!(p.name(), "openai"); // uses OpenAiEmbedding internally
assert_eq!(p.dimensions(), 768);
}
}