rig icon indicating copy to clipboard operation
rig copied to clipboard

refactor: Generalize and simplify vector store interface

Open cvauclair opened this issue 5 months ago • 2 comments

Feature Request

Generalize and simplify the vector store interface (and the integration of third party vector stores).

Motivation

Rig's current approach to vector stores and vector search is lacking in multiple ways:

  1. The interface forces developers to use the DocumentEmbeddings type, which is somewhat opinionated and a little over-engineered.
  2. The interface doesn't lend itself well to use cases where a developer already has a populated vector store since the interface expects the vector store to be modeled after DocumentEmbeddings.
  3. The process of integrating new vector stores is convoluted for non-document databases (e.g.: Postgres, LanceDB) since DocumentEmbeddings was designed for document vector stores.
  4. The interface assumes that user's would use Rig constructs (e.g.: DocumentEmbeddings) to populate their vector store.

Proposal

  • Remove the VectorStore trait and simplify the VectorStoreIndex trait to the following methods only:
pub trait VectorStoreIndex: Send + Sync {
    /// Get the top n documents based on the distance to the given embedding.
    /// The documents are deserialized into the given type.
    async fn top_n_from_query<T: for<'a> Deserialize<'a>>(
        &self,
        query: &str,
        n: usize,
    ) -> Result<Vec<(f64, T)>, VectorStoreError>;

    /// Same as `top_n_from_query` but returns the document ids only.
    async fn top_n_ids_from_query(
        &self,
        query: &str,
        n: usize,
    ) -> Result<Vec<(f64, String)>, VectorStoreError>;

    /// Get the top n documents based on the distance to the given embedding.
    /// The documents are deserialized into the given type.
    async fn top_n_from_embedding<T: for<'a> Deserialize<'a>>(
        &self,
        embedding: &Embedding,
        n: usize,
    ) -> Result<Vec<(f64, T)>, VectorStoreError>;

    /// Same as `top_n_from_embedding` but returns the document ids only.
    async fn top_n_ids_from_embedding(
        &self,
        embedding: &Embedding,
        n: usize,
    ) -> Result<Vec<(f64, String)>, VectorStoreError>;
}
  • Remove the DocumentEmbeddings type entirely.
  • Update the Agent type accordingly (we could enforce that the type T which is stored in the vector store also implements ToString so that the we can easily insert the dynamic context in the agent's prompt)
  • Update the EmbeddingsBuilder type accordingly
  • Update the existing vector stores integration

Alternatives

Open to alternatives

Implementation Checklist

  • [ ] #40
  • [x] #41

cvauclair avatar Sep 24 '24 21:09 cvauclair