Story Crater Bot
33eaf1b4f8
- Add database schema with pgvector extension (L0/L1/L2 memories) - Implement pgvector-backed vector store with similarity search - Add Ollama embeddings client for 768-dim nomic embeddings - Implement ingest worker to process records into L0/L1 memory - Implement query worker with semantic search across memory tiers - Rewrite HTTP server with database connection pooling - Wire all endpoints to actual backend (ingest, query, projects, skills) - Update main.rs to use DATABASE_URL from environment - All code compiles, ready for Docker build and deployment
112 lines
3.8 KiB
Rust
112 lines
3.8 KiB
Rust
use anyhow::Result;
|
|
use mem_llm::{EmbeddingsClient, RerankClient};
|
|
use mem_store::VectorStore;
|
|
use pgvector::Vector;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Query result with provenance
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct QueryResult {
|
|
pub level: String, // "L0", "L1", "L2", "corpus"
|
|
pub score: f32,
|
|
pub text: String,
|
|
pub source: Option<String>,
|
|
pub provenance: Vec<String>, // parent IDs
|
|
}
|
|
|
|
/// Query worker — semantic search + reranking
|
|
pub struct QueryWorker {
|
|
vector_store: std::sync::Arc<VectorStore>,
|
|
embeddings: std::sync::Arc<EmbeddingsClient>,
|
|
reranker: std::sync::Arc<RerankClient>,
|
|
}
|
|
|
|
impl QueryWorker {
|
|
/// Create query worker
|
|
pub fn new(
|
|
vector_store: VectorStore,
|
|
embeddings: EmbeddingsClient,
|
|
reranker: RerankClient,
|
|
) -> Self {
|
|
Self {
|
|
vector_store: std::sync::Arc::new(vector_store),
|
|
embeddings: std::sync::Arc::new(embeddings),
|
|
reranker: std::sync::Arc::new(reranker),
|
|
}
|
|
}
|
|
|
|
/// Execute semantic query: embed -> search vector -> rerank -> result
|
|
pub async fn query(
|
|
&self,
|
|
project: &str,
|
|
question: &str,
|
|
limit: Option<i64>,
|
|
) -> Result<Vec<QueryResult>> {
|
|
let limit = limit.unwrap_or(5);
|
|
|
|
// Embed the question
|
|
let question_embedding = self.embeddings.embed(question).await?;
|
|
|
|
// Search across all levels
|
|
let mut candidates = Vec::new();
|
|
|
|
// L2 synthesis (project-level)
|
|
if let Some(l2_result) = self.vector_store.search_l2(project, &question_embedding).await? {
|
|
candidates.push(QueryResult {
|
|
level: "L2".to_string(),
|
|
score: l2_result.score,
|
|
text: l2_result.item.content.clone(),
|
|
source: Some(format!("project:{}", project)),
|
|
provenance: vec![l2_result.item.id.to_string()],
|
|
});
|
|
}
|
|
|
|
// L1 per-query memories
|
|
let l1_results = self.vector_store.search_l1(project, &question_embedding, limit).await?;
|
|
for l1_result in l1_results {
|
|
candidates.push(QueryResult {
|
|
level: "L1".to_string(),
|
|
score: l1_result.score,
|
|
text: l1_result.item.content.clone(),
|
|
source: Some(format!("query:{}", l1_result.item.query_id)),
|
|
provenance: vec![l1_result.item.id.to_string()],
|
|
});
|
|
}
|
|
|
|
// Reference corpus
|
|
let corpus_results = self.vector_store.search_corpus(project, &question_embedding, limit).await?;
|
|
for corpus_result in corpus_results {
|
|
candidates.push(QueryResult {
|
|
level: "corpus".to_string(),
|
|
score: corpus_result.score,
|
|
text: corpus_result.item.content.clone(),
|
|
source: Some(format!("doc:{}", corpus_result.item.name)),
|
|
provenance: vec![corpus_result.item.id.to_string()],
|
|
});
|
|
}
|
|
|
|
// Rerank candidates by relevance to question
|
|
// TODO: wire actual cross-encoder reranking
|
|
// For now, return by vector similarity score
|
|
candidates.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
|
|
candidates.truncate(limit as usize);
|
|
|
|
Ok(candidates)
|
|
}
|
|
|
|
/// Get project synthesis (L2) directly
|
|
pub async fn get_synthesis(&self, project: &str) -> Result<Option<QueryResult>> {
|
|
if let Some(l2) = self.vector_store.get_l2(project).await? {
|
|
Ok(Some(QueryResult {
|
|
level: "L2".to_string(),
|
|
score: 1.0,
|
|
text: l2.content,
|
|
source: Some(format!("project:{}", project)),
|
|
provenance: vec![l2.id.to_string()],
|
|
}))
|
|
} else {
|
|
Ok(None)
|
|
}
|
|
}
|
|
}
|