rig
rig copied to clipboard
refactor: Generalize and simplify vector store interface
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:
- The interface forces developers to use the
DocumentEmbeddingstype, which is somewhat opinionated and a little over-engineered. - 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. - The process of integrating new vector stores is convoluted for non-document databases (e.g.: Postgres, LanceDB) since
DocumentEmbeddingswas designed for document vector stores. - The interface assumes that user's would use Rig constructs (e.g.:
DocumentEmbeddings) to populate their vector store.
Proposal
- Remove the
VectorStoretrait and simplify theVectorStoreIndextrait 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
DocumentEmbeddingstype entirely. - Update the
Agenttype accordingly (we could enforce that the typeTwhich is stored in the vector store also implementsToStringso that the we can easily insert the dynamic context in the agent's prompt) - Update the
EmbeddingsBuildertype accordingly - Update the existing vector stores integration
Alternatives
Open to alternatives
Implementation Checklist
- [x] #40
- [x] #41
@0xMochan @ThierryBleau what do you guys thinks of this?
@cvauclair This should close since #52 was merged right? Same for sub-issue #40
Is this issue resolved? (just asking for the sake of housekeeping - we shouldn't keep issues open if there have already been merged PRs that fix the issue)