swiftide_pgvector/
lib.rs

1mod persist;
2mod retrieve;
3
4use anyhow::Result;
5use derive_builder::Builder;
6use sqlx::PgPool;
7use swiftide_core::Persist;
8
9#[derive(Builder, Debug, Clone)]
10pub struct PgVector {
11    /// The database connection pool.
12    pool: PgPool,
13    /// The name of the table to store the vectors in.
14    #[builder(default = "String::from(\"swiftide_rag\")")]
15    table_name: String,
16    /// The size of the vectors to store in the database.
17    vector_size: i32,
18    /// The batch size to use when storing nodes.
19    #[builder(default = "128")]
20    batch_size: usize,
21}
22
23impl PgVector {
24    pub async fn try_new(pool: PgPool, vector_size: i32) -> Result<Self> {
25        let vector = PgVectorBuilder::default()
26            .pool(pool)
27            .vector_size(vector_size)
28            .build()?;
29
30        vector.setup().await?;
31        Ok(vector)
32    }
33
34    pub fn get_pool(&self) -> &PgPool {
35        &self.pool
36    }
37}