Skip to main content

gonzalo_vector/
lib.rs

1//! Vector-search capability layer for gonzalo.
2//!
3//! Provides the [`VectorIndex`] trait (and its exact in-memory implementation
4//! [`MemoryVectorIndex`]) plus the [`Embedder`] seam so caliban-side model
5//! providers and gonzalo-hosted embedders share a single interface.
6//!
7//! All index operations are async so remote / approximate backends can be
8//! added in later milestones without breaking callers.
9
10pub mod index;
11pub use index::MemoryVectorIndex;
12
13pub mod shard;
14pub use shard::{DEFAULT_SHARDS, shard_of};
15
16pub mod record_index;
17pub use record_index::RecordVectorIndex;
18
19#[cfg(feature = "hnsw")]
20pub mod hnsw;
21#[cfg(feature = "hnsw")]
22pub use hnsw::HnswVectorIndex;
23
24use async_trait::async_trait;
25use gonzalo_core::{KeyPrefix, RecordKey, Result};
26
27// ---------------------------------------------------------------------------
28// Embedder
29// ---------------------------------------------------------------------------
30
31/// Turns text into an embedding vector.
32///
33/// The default deployment delegates this to the caller (caliban, which talks
34/// to model providers); gonzalo can also host its own embedder. This trait
35/// is the seam for both.
36#[async_trait]
37pub trait Embedder: Send + Sync {
38    async fn embed(&self, text: &str) -> Result<Vec<f32>>;
39}
40
41// ---------------------------------------------------------------------------
42// Match
43// ---------------------------------------------------------------------------
44
45/// One search hit: the record key and its similarity score (cosine, in
46/// `[-1.0, 1.0]`; higher is more similar).
47#[derive(Debug, Clone, PartialEq)]
48pub struct Match {
49    pub key: RecordKey,
50    pub score: f32,
51}
52
53// ---------------------------------------------------------------------------
54// VectorIndex
55// ---------------------------------------------------------------------------
56
57/// A vector index keyed by [`RecordKey`].
58///
59/// Async so remote / approximate backends can implement it later; the
60/// in-memory impl ([`MemoryVectorIndex`]) is exact brute-force cosine kNN.
61#[async_trait]
62pub trait VectorIndex: Send + Sync {
63    /// Insert or replace the vector for `key`.
64    async fn upsert(&self, key: RecordKey, vector: Vec<f32>) -> Result<()>;
65
66    /// Remove `key` if present (no error if absent).
67    async fn remove(&self, key: &RecordKey) -> Result<()>;
68
69    /// Return the top-`k` matches to `query`, restricted to keys matching
70    /// `filter`, ordered by descending score.
71    async fn query(&self, query: &[f32], k: usize, filter: &KeyPrefix) -> Result<Vec<Match>>;
72
73    /// Insert or replace many vectors at once.
74    ///
75    /// The default loops over [`upsert`](Self::upsert). A durable backend
76    /// overrides this: committing once per vector would mean one manifest
77    /// commit per vector on a bulk load.
78    async fn upsert_many(&self, items: Vec<(RecordKey, Vec<f32>)>) -> Result<()> {
79        for (key, vector) in items {
80            self.upsert(key, vector).await?;
81        }
82        Ok(())
83    }
84
85    /// Every key in the index matching `filter`. Order is unspecified.
86    ///
87    /// Used to rebuild derived state that would otherwise be lost across a
88    /// restart, such as a `KnowledgeStore`'s per-record chunk counts.
89    async fn keys(&self, filter: &KeyPrefix) -> Result<Vec<RecordKey>>;
90}
91
92// ---------------------------------------------------------------------------
93// Arc<T> delegation
94// ---------------------------------------------------------------------------
95
96#[async_trait]
97impl<T: VectorIndex + ?Sized> VectorIndex for std::sync::Arc<T> {
98    async fn upsert(&self, key: RecordKey, vector: Vec<f32>) -> Result<()> {
99        (**self).upsert(key, vector).await
100    }
101    async fn remove(&self, key: &RecordKey) -> Result<()> {
102        (**self).remove(key).await
103    }
104    async fn query(&self, query: &[f32], k: usize, filter: &KeyPrefix) -> Result<Vec<Match>> {
105        (**self).query(query, k, filter).await
106    }
107    async fn keys(&self, filter: &KeyPrefix) -> Result<Vec<RecordKey>> {
108        (**self).keys(filter).await
109    }
110    async fn upsert_many(&self, items: Vec<(RecordKey, Vec<f32>)>) -> Result<()> {
111        (**self).upsert_many(items).await
112    }
113}