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, Default)]
53pub enum DistanceMetric {
54    /// Cosine similarity (1 - cosine distance). Best for normalized embeddings.
55    #[default]
56    Cosine,
57    /// Euclidean (L2) distance. Best for absolute spatial similarity.
58    L2,
59    /// Inner product. Best for maximum inner product search (MIPS).
60    InnerProduct,
61}
62
63impl DistanceMetric {
64    /// Returns the pgvector operator class for HNSW index creation.
65    pub fn ops_class(&self) -> &'static str {
66        match self {
67            Self::Cosine => "vector_cosine_ops",
68            Self::L2 => "vector_l2_ops",
69            Self::InnerProduct => "vector_ip_ops",
70        }
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn test_distance_metric_ops_class() {
80        assert_eq!(DistanceMetric::Cosine.ops_class(), "vector_cosine_ops");
81        assert_eq!(DistanceMetric::L2.ops_class(), "vector_l2_ops");
82        assert_eq!(DistanceMetric::InnerProduct.ops_class(), "vector_ip_ops");
83    }
84
85    #[test]
86    fn test_distance_metric_default() {
87        assert_eq!(DistanceMetric::default(), DistanceMetric::Cosine);
88    }
89}
daimon_plugin_pgvector/lib.rs

daimon_plugin_pgvector/
lib.rs
