rig_vectorize/

lib.rs

//! Cloudflare Vectorize integration for the Rig framework.
//!
//! This crate provides a vector store implementation using Cloudflare Vectorize,
//! a globally distributed vector database built for AI applications.
//!
//! # Example
//!
//! ```ignore
//! use rig_core::client::ProviderClient;
//! use rig_core::providers::openai;
//! use rig_vectorize::VectorizeVectorStore;
//!
//! let openai = openai::Client::from_env()?;
//! let embedding_model = openai.embedding_model(openai::TEXT_EMBEDDING_3_SMALL);
//!
//! let vector_store = VectorizeVectorStore::new(
//!     embedding_model,
//!     "your-account-id",
//!     "your-index-name",
//!     std::env::var("CLOUDFLARE_API_TOKEN").unwrap(),
//! );
//! ```

mod client;

// Re-export client types
pub use client::{
    DeleteByIdsRequest, DeleteResult, ListVectorsResult, QueryRequest, QueryResult, ReturnMetadata,
    UpsertRequest, UpsertResult, VectorIdEntry, VectorInput, VectorMatch, VectorizeClient,
    VectorizeError, VectorizeFilter,
};

use client::{QueryRequest as ApiQueryRequest, VectorInput as ApiVectorInput};
use rig_core::embeddings::EmbeddingModel;
use rig_core::vector_store::request::VectorSearchRequest;
use rig_core::vector_store::{InsertDocuments, VectorStoreError, VectorStoreIndex};
use rig_core::{Embed, embeddings::Embedding};
use serde::{Deserialize, Serialize};
use uuid::Uuid;

impl From<VectorizeError> for VectorStoreError {
    fn from(err: VectorizeError) -> Self {
        VectorStoreError::datastore(err)
    }
}

/// A vector store backed by Cloudflare Vectorize.
///
/// This struct implements [`VectorStoreIndex`] to provide vector similarity search
/// using Cloudflare's globally distributed Vectorize service.
#[derive(Debug, Clone)]
pub struct VectorizeVectorStore<M> {
    /// The embedding model used to generate query embeddings.
    model: M,
    /// The HTTP client for Vectorize API.
    client: VectorizeClient,
}

impl<M> VectorizeVectorStore<M> {
    /// Creates a new Vectorize vector store.
    ///
    /// # Arguments
    /// * `model` - The embedding model to use for query embedding
    /// * `account_id` - Cloudflare account ID
    /// * `index_name` - Name of the Vectorize index
    /// * `api_token` - Cloudflare API token with Vectorize read permissions
    pub fn new(
        model: M,
        account_id: impl Into<String>,
        index_name: impl Into<String>,
        api_token: impl Into<String>,
    ) -> Self {
        Self {
            model,
            client: VectorizeClient::new(account_id, index_name, api_token),
        }
    }
}

impl<M> VectorizeVectorStore<M>
where
    M: EmbeddingModel + Sync + Send,
{
    /// Validates the filter, embeds the query, and returns the threshold-filtered matches.
    async fn query_matches(
        &self,
        req: &VectorSearchRequest<VectorizeFilter>,
        return_metadata: ReturnMetadata,
    ) -> Result<Vec<VectorMatch>, VectorStoreError> {
        if let Some(filter) = req.filter() {
            filter.validate()?;
        }

        let embedding = self.model.embed_text(req.query()).await?;

        let query_request = ApiQueryRequest {
            vector: embedding.vec,
            top_k: req.samples(),
            return_values: Some(false),
            return_metadata: Some(return_metadata),
            filter: req.filter().as_ref().map(|f| f.clone().into_inner()),
        };

        let result = self.client.query(query_request).await?;

        Ok(result
            .matches
            .into_iter()
            .filter(|m| req.threshold().is_none_or(|t| m.score >= t))
            .collect())
    }
}

impl<M> VectorStoreIndex for VectorizeVectorStore<M>
where
    M: EmbeddingModel + Sync + Send,
{
    type Filter = VectorizeFilter;

    async fn top_n<T: for<'a> Deserialize<'a> + Send>(
        &self,
        req: VectorSearchRequest<Self::Filter>,
    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
        let matches = self.query_matches(&req, ReturnMetadata::All).await?;

        // Convert results to the expected format
        let results = matches
            .into_iter()
            .map(|m| {
                let metadata = m.metadata.unwrap_or(serde_json::Value::Null);
                let doc: T = serde_json::from_value(metadata)?;
                Ok((m.score, m.id, doc))
            })
            .collect::<Result<Vec<_>, serde_json::Error>>()?;

        Ok(results)
    }

    async fn top_n_ids(
        &self,
        req: VectorSearchRequest<Self::Filter>,
    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
        let matches = self.query_matches(&req, ReturnMetadata::None).await?;

        // Convert results to (score, id) tuples
        Ok(matches.into_iter().map(|m| (m.score, m.id)).collect())
    }
}

impl<M> InsertDocuments for VectorizeVectorStore<M>
where
    M: EmbeddingModel + Sync + Send,
{
    async fn insert_documents<Doc: Serialize + Embed + Send>(
        &self,
        documents: Vec<(Doc, Vec<Embedding>)>,
    ) -> Result<(), VectorStoreError> {
        let vectors =
            rig_core::vector_store::flatten_embedded(documents, |metadata, embedding| {
                Ok(ApiVectorInput {
                    id: Uuid::new_v4().to_string(),
                    values: embedding.vec,
                    metadata: Some(metadata.clone()),
                    namespace: None,
                })
            })?;

        if vectors.is_empty() {
            return Ok(());
        }

        tracing::debug!("Upserting {} vectors to Vectorize", vectors.len());

        const BATCH_SIZE: usize = 1000;

        for batch in vectors.chunks(BATCH_SIZE) {
            let request = UpsertRequest {
                vectors: batch.to_vec(),
            };

            self.client.upsert(request).await?;
        }

        Ok(())
    }
}