82 lines
2.0 KiB
Rust
82 lines
2.0 KiB
Rust
use anyhow::Result;
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// Vector embedding record in pgvector.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct VectorRecord {
|
|
pub id: String,
|
|
pub chunk_id: String,
|
|
pub kind: String, // "text" | "symptom"
|
|
pub embedding: Vec<f32>, // 768-dimensional for nomic
|
|
pub tokens: u32,
|
|
}
|
|
|
|
/// pgvector client.
|
|
pub struct VectorStore {
|
|
// In production: PostgreSQL connection
|
|
// For now: in-memory vec
|
|
records: Vec<VectorRecord>,
|
|
}
|
|
|
|
impl VectorStore {
|
|
/// Create a new vector store.
|
|
pub fn new() -> Self {
|
|
Self {
|
|
records: Vec::new(),
|
|
}
|
|
}
|
|
|
|
/// Insert a vector record.
|
|
pub fn insert(&mut self, record: VectorRecord) -> Result<()> {
|
|
self.records.push(record);
|
|
Ok(())
|
|
}
|
|
|
|
/// Search by cosine similarity.
|
|
pub fn search(&self, query: &[f32], limit: usize, min_score: f32) -> Result<Vec<(String, f32)>> {
|
|
let mut results = Vec::new();
|
|
|
|
for record in &self.records {
|
|
if let Some(score) = cosine_similarity(query, &record.embedding) {
|
|
if score >= min_score {
|
|
results.push((record.id.clone(), score));
|
|
}
|
|
}
|
|
}
|
|
|
|
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
|
Ok(results.into_iter().take(limit).collect())
|
|
}
|
|
|
|
/// Get all records.
|
|
pub fn all(&self) -> Vec<&VectorRecord> {
|
|
self.records.iter().collect()
|
|
}
|
|
}
|
|
|
|
/// Compute cosine similarity between two vectors.
|
|
fn cosine_similarity(a: &[f32], b: &[f32]) -> Option<f32> {
|
|
if a.len() != b.len() {
|
|
return None;
|
|
}
|
|
|
|
let mut dot_product = 0.0;
|
|
let mut norm_a = 0.0;
|
|
let mut norm_b = 0.0;
|
|
|
|
for (x, y) in a.iter().zip(b.iter()) {
|
|
dot_product += x * y;
|
|
norm_a += x * x;
|
|
norm_b += y * y;
|
|
}
|
|
|
|
let norm_a = norm_a.sqrt();
|
|
let norm_b = norm_b.sqrt();
|
|
|
|
if norm_a == 0.0 || norm_b == 0.0 {
|
|
return None;
|
|
}
|
|
|
|
Some(dot_product / (norm_a * norm_b))
|
|
}
|