daimon_plugin_pgvector/
lib.rs

1//! # daimon-plugin-pgvector
2//!
3//! A pgvector-backed [`VectorStore`](daimon_core::VectorStore) plugin for
4//! the [Daimon](https://docs.rs/daimon) AI agent framework.
5//!
6//! This crate provides [`PgVectorStore`], which stores document embeddings
7//! in PostgreSQL using the [pgvector](https://github.com/pgvector/pgvector)
8//! extension. It supports cosine, L2, and inner-product distance metrics
9//! with HNSW indexing.
10//!
11//! ## Quick Start
12//!
13//! ```ignore
14//! use daimon_plugin_pgvector::{PgVectorStoreBuilder, DistanceMetric};
15//! use daimon::retriever::SimpleKnowledgeBase;
16//! use std::sync::Arc;
17//!
18//! let store = PgVectorStoreBuilder::new("postgresql://user:pass@localhost/db", 1536)
19//!     .table("my_docs")
20//!     .distance_metric(DistanceMetric::Cosine)
21//!     .build()
22//!     .await?;
23//!
24//! // Compose with an embedding model for a full RAG pipeline:
25//! let kb = SimpleKnowledgeBase::new(embedding_model, store);
26//! ```
27//!
28//! ## Manual Schema Setup
29//!
30//! If you prefer to manage migrations yourself, disable auto-migration
31//! and use the SQL from [`migrations`]:
32//!
33//! ```ignore
34//! let store = PgVectorStoreBuilder::new(conn_str, 1536)
35//!     .auto_migrate(false)
36//!     .build()
37//!     .await?;
38//! ```
39//!
40//! Then run the SQL from [`migrations::CREATE_EXTENSION`],
41//! [`migrations::create_table_sql`], and [`migrations::create_hnsw_index_sql`]
42//! against your database.
43
44mod builder;
45pub mod migrations;
46mod store;
47
48pub use builder::PgVectorStoreBuilder;
49pub use store::PgVectorStore;
50
51/// Distance metric used for vector similarity search.
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub enum DistanceMetric {
54    /// Cosine similarity (1 - cosine distance). Best for normalized embeddings.
55    Cosine,
56    /// Euclidean (L2) distance. Best for absolute spatial similarity.
57    L2,
58    /// Inner product. Best for maximum inner product search (MIPS).
59    InnerProduct,
60}
61
62impl DistanceMetric {
63    /// Returns the pgvector operator class for HNSW index creation.
64    pub fn ops_class(&self) -> &'static str {
65        match self {
66            Self::Cosine => "vector_cosine_ops",
67            Self::L2 => "vector_l2_ops",
68            Self::InnerProduct => "vector_ip_ops",
69        }
70    }
71}
72
73impl Default for DistanceMetric {
74    fn default() -> Self {
75        Self::Cosine
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn test_distance_metric_ops_class() {
85        assert_eq!(DistanceMetric::Cosine.ops_class(), "vector_cosine_ops");
86        assert_eq!(DistanceMetric::L2.ops_class(), "vector_l2_ops");
87        assert_eq!(DistanceMetric::InnerProduct.ops_class(), "vector_ip_ops");
88    }
89
90    #[test]
91    fn test_distance_metric_default() {
92        assert_eq!(DistanceMetric::default(), DistanceMetric::Cosine);
93    }
94}
daimon_plugin_pgvector/lib.rs

daimon_plugin_pgvector/
lib.rs
