Struct PostgresVectorStore

pub struct PostgresVectorStore { /* private fields */ }

PostgresVectorStore

This is the implementation of EmbeddingStore for a Postgres database with the pgvector extension enabled. This store takes a table name and an embedding model. If a table already exists with the same name and does not have the expected columns any calls to PostgresVectorStore::store or PostgresVectorStore::store_batch will fail.

Required Environment Variables

• POSTGRES_USERNAME: The username to connect to the database with
• POSTGRES_PASSWORD: The password to connect to the database with
• POSTGRES_HOST: The host to connect to the database with
• POSTGRES_DATABASE: The database to connect with

Output table format

Columns: | id (int) | content (text) | embedding (vector) | metadata (jsonb) |

Examples

use rag_toolchain::stores::*;
use rag_toolchain::common::*;

async fn store(embeddings: Vec<Embedding>) {
    let embedding_model: OpenAIEmbeddingModel = OpenAIEmbeddingModel::TextEmbedding3Small;
    let table_name: &str = "table_name";
    let store: PostgresVectorStore = PostgresVectorStore::try_new(table_name, embedding_model)
        .await.unwrap();
    store.store_batch(embeddings).await.unwrap();
}

Implementations

impl PostgresVectorStore

pub async fn try_new( table_name: &str, embedding_model: impl EmbeddingModel, ) -> Result<Self, PostgresVectorStoreError>

PostgresVectorStore::try_new

This constructor is used to create a new PostgresVectorStore. It will read the required environment variables in. Try and connect to your postgres database and then create a table with the given name and the expected columns. If the table already exists with the same name it will not be re-created.

Arguments

• table_name: &str - The name of the table to store the embeddings in.
• embedding_model: impl EmbeddingModel - The embedding model to use to store the embeddings

Errors

• [PostgresVectorError::EnvVarError] if the required environment variables are not set.
• [PostgresVectorError::ConnectionError] if the connection to the database could not be established.
• [PostgresVectorError::TableCreationError] if the table could not be created.

Returns

• PostgresVectorStore if the connection and table creation is successful

pub async fn try_new_with_pool( pool: Pool<Postgres>, table_name: &str, embedding_model: impl EmbeddingModel, ) -> Result<Self, PostgresVectorStoreError>

PostgresVectorStore::try_new_with_pool

This is an alternative constructor that allows you to pass in a connection pool. This was added as it may be the case people want to establish one connection pool and then shared it across multiple PostgresVectorStores managing different tables.

Arguments

• pool: sqlx::Pool<Postgres> - a pre established connection pool.
• table_name: &str - The name of the table to store the embeddings in.
• embedding_model: implEmbeddingModel - The embedding model used for the genrated embeddings.

Errors

• [PostgresVectorError::TableCreationError] if the table could not be created

Returns

• PostgresVectorStore if the table creation is successful.

pub fn get_pool(&self) -> Pool<Postgres>

PostgresVectorStore::get_pool

Getter for the internal connection pool. This is useful if you want to do any further operations on the database such as enabling an index on the table.

Returns

• Pool - The connection pool

pub fn as_retriever<T: AsyncEmbeddingClient>( &self, embedding_client: T, distance_function: DistanceFunction, ) -> PostgresVectorRetriever<T>

PostgresVectorStore::as_retriever

This function allows us to convert the store into a retriever. Note that the returned retriever is bound to the same table as the store.

Arguments

• embedding_client: AsyncEmbeddingClient - The client we use to embed ` income text before the similarity search.
• distance_function: DistanceFunction - The distance function to use to compare the embeddings

Returns

PostgresVectorRetriever - The retriever that can be used to search for similar text.

Trait Implementations

impl Clone for PostgresVectorStore

fn clone(&self) -> PostgresVectorStore

Returns a copy of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for PostgresVectorStore

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl EmbeddingStore for PostgresVectorStore

async fn store( &self, embedding: Embedding, ) -> Result<(), PostgresVectorStoreError>

PostgresVectorStore::store

This is done as a single insert statement.

Arguments

• embedding: Embedding - to insert

Errors

• [PostgresVectorError::InsertError] if the insert fails

Returns

• [()] if the insert succeeds

async fn store_batch( &self, embeddings: Vec<Embedding>, ) -> Result<(), PostgresVectorStoreError>

PostgresVectorStore::store_batch

This is done as a single transaction with multiple insert statements.

Arguments

• embeddings: Vec<Embedding> - A vector of embeddings to insert

Errors

• [PostgresVectorError::TransactionError] if the transaction fails

Returns

• [()] if the transaction succeeds

type ErrorType = PostgresVectorStoreError

The custom error type for the store