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:
parent
4fceba0740
commit
0e7f501fd6
10 changed files with 1423 additions and 96 deletions
190
src/memory/embeddings.rs
Normal file
190
src/memory/embeddings.rs
Normal 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);
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue