feat(query): Add support for single embedding retrieval with PGVector

[shamb0]

Part 2 of PR: swiftide/pull/392

• Added retrieval functionality to pgvector.
• Verified indexing and querying through unit tests and example demo applications. Integration works as expected.

[shamb0]

Hi @timonv

I’m reaching out regarding some unit test failures related to the retrieval process. Here are the specific scenarios that are failing:

failures:  
    pgvector::tests::test_persist_nodes::both_mode_with_metadata  
    pgvector::tests::test_persist_nodes::both_mode_without_metadata  
    pgvector::tests::test_persist_nodes::perfield_mode_with_metadata  
    pgvector::tests::test_persist_nodes::perfield_mode_without_metadata  

The failures seem to stem from the complexity of managing multiple columns of vector embeddings. The current implementation of Query<T> has limitations when querying individual vector embedding columns.

Key Issue

To ensure proper functionality, both the Embedding value and its corresponding EmbeddedField type must be specified in the query. However, it seems some field variables are missing these options:

• Query<T>::embedding
• Query<T>::sparse_embedding

To reproduce the issue for the above scenarios, set the use_adv_embedding_query flag to false (it is set to true by default). This should replicate the failures.

Temporary Workaround

As a temporary measure, I introduced a new type, Query<T>::adv_embedding, to highlight the limitation. However, this is not intended as a final solution.

Request for Guidance

I noticed your comment in the code:

// TODO: How would this work when doing a rollup query?  

It seems you’ve already been considering ways to consolidate query options. Could you please share your thoughts on how we might address this scenario effectively.

Thank you for your time and assistance! I’m looking forward to your insights.

[timonv]

Hey @shamb0,

Interesting, I'll get back to this more in depth asap. The query pipeline currently has a limitation, as you pointed out, when retrieving from multiple sources. Originally, the idea was for search strategies to provide a sane default to get started quickly, perhaps that can still hold, i.e. SimilaritySingleEmbedding is meant for when there is a single embedding. Storage layers can implement multiple strategies, so perhaps a MultipleNamedEmbeddings could be a good introduction here. I have a hunch adv_embedding isn't needed, as you pointed out, but I'd like to get back to that more in depth.

Zooming out and looking at the bigger picture and future architecture, I'm tempted to tackle the issue with strategies as much as possible. For instance, this was some mock coding I did recently on fusion search (multiple retrievers):

let fusion_retriever = FusionRetriever::builder()
    .add_retriever(SimpleSimilaritySearch::default(), Qdrant::default())
    .add_retriever(FullTextSearch::default(), Lancedb::default())
    .build();
    
impl<S: SearchStrategy> Retrieve<S> for FusionRetriever {
    fn retrieve(_strat, query) {
        let documents = tokio::join!(self.retrievers.iter().map(|strat, retriever| retriever.retrieve(strat, query));
        
        query.retrieved_documents(documents)
    }
}

// And then in a pipeline
query::Pipeline::default()
    ...
    .then_retrieve(fusion_retriever)
    .then(Reranker::builder().top_n(10).bm25().build()?)
    .query("How can I use the query pipeline in Swiftide?")
    .await?;

From my perspective I think it's relatively easy to implement, and having an architecture like this could really nicely enable other methods of retrieval as well (i.e. parent lookup, routed retrieval, other wild combinations). However, I do feel it introduces some odd abstractions; then again, if it's only for advanced usage, and there's sane, easier defaults to start with, perhaps it's not that bad. What are your thoughts on this?

My vision on the query pipeline (discussion greatly appreciated):

• Search strategies offer a sane default to get started quickly
• Retrieved documents should be first class citizens in the pipeline; you get a big pile, and they can be modified, mangled, and filtered by follow up steps, and thus not stored in the state, but in the query itself.
• Pre-retrieval transformation transform the query to something that suits the search strategy and retrieval methods, post retrieval works towards answer generation
• In general; I prefer to offer nice defaults, and not abstract away too much from the underlying api. If a user wants to do something more complicated, the underlying query api should be available.

// TODO: How would this work when doing a rollup query?

So one use case, you might want to generate subquestions on a query, embed each individually, then retrieve for each individually (perhaps even routed to different retrievers), then consolidate the documents, rerank, and put it back together in a single query. I think there is a nice, low level architectural solution to this, but I haven't given it that much thought yet.

I'm mostly working on agents right now, I don't really have the time to make large architectural changes.

Perhaps it's an idea to focus on one search strategy at a time? What are your thoughts on this?

[timonv]

@shamb0 You're not on our Discord right? Might be a bit faster to discuss things there! Love the effort you're putting into this.

[shamb0]

@shamb0 You're not on our Discord right? Might be a bit faster to discuss things there! Love the effort you're putting into this.

Sure @timonv, I'm joining now 😄 👍🏾 .

[shamb0]

Hey @shamb0,

Interesting, I'll get back to this more in depth asap. The query pipeline currently has a limitation, as you pointed out, when retrieving from multiple sources. Originally, the idea was for search strategies to provide a sane default to get started quickly, perhaps that can still hold, i.e. SimilaritySingleEmbedding is meant for when there is a single embedding. Storage layers can implement multiple strategies, so perhaps a MultipleNamedEmbeddings could be a good introduction here. I have a hunch adv_embedding isn't needed, as you pointed out, but I'd like to get back to that more in depth.

Zooming out and looking at the bigger picture and future architecture, I'm tempted to tackle the issue with strategies as much as possible. For instance, this was some mock coding I did recently on fusion search (multiple retrievers):

let fusion_retriever = FusionRetriever::builder()
    .add_retriever(SimpleSimilaritySearch::default(), Qdrant::default())
    .add_retriever(FullTextSearch::default(), Lancedb::default())
    .build();
    
impl<S: SearchStrategy> Retrieve<S> for FusionRetriever {
    fn retrieve(_strat, query) {
        let documents = tokio::join!(self.retrievers.iter().map(|strat, retriever| retriever.retrieve(strat, query));
        
        query.retrieved_documents(documents)
    }
}

// And then in a pipeline
query::Pipeline::default()
    ...
    .then_retrieve(fusion_retriever)
    .then(Reranker::builder().top_n(10).bm25().build()?)
    .query("How can I use the query pipeline in Swiftide?")
    .await?;

From my perspective I think it's relatively easy to implement, and having an architecture like this could really nicely enable other methods of retrieval as well (i.e. parent lookup, routed retrieval, other wild combinations). However, I do feel it introduces some odd abstractions; then again, if it's only for advanced usage, and there's sane, easier defaults to start with, perhaps it's not that bad. What are your thoughts on this?

My vision on the query pipeline (discussion greatly appreciated):

• Search strategies offer a sane default to get started quickly
• Retrieved documents should be first class citizens in the pipeline; you get a big pile, and they can be modified, mangled, and filtered by follow up steps, and thus not stored in the state, but in the query itself.
• Pre-retrieval transformation transform the query to something that suits the search strategy and retrieval methods, post retrieval works towards answer generation
• In general; I prefer to offer nice defaults, and not abstract away too much from the underlying api. If a user wants to do something more complicated, the underlying query api should be available.

// TODO: How would this work when doing a rollup query?

So one use case, you might want to generate subquestions on a query, embed each individually, then retrieve for each individually (perhaps even routed to different retrievers), then consolidate the documents, rerank, and put it back together in a single query. I think there is a nice, low level architectural solution to this, but I haven't given it that much thought yet.

I'm mostly working on agents right now, I don't really have the time to make large architectural changes.

Perhaps it's an idea to focus on one search strategy at a time? What are your thoughts on this?

Hi @timonv,

Thank you for sharing the architectural overview and your thoughts on FusionRetriever. The concept is fascinating and shows great potential for enhancing the query pipeline. The ability to integrate multiple search strategies from diverse sources and seamlessly feed the retrieved documents into post-retrieval stages is both flexible and robust. I also appreciate your emphasis on maintaining simplicity through sane defaults while enabling advanced use cases.

I’m interested in the use case you shared:

"One use case might involve generating subquestions from a query, embedding each individually, retrieving for each (potentially routed to different retrievers), consolidating the results, reranking them, and assembling everything back into a single query."

If your schedule is currently occupied with higher-priority tasks, would you be open to allowing me to prototype this concept? I’d love to start exploring the FusionRetriever architecture, using your suggested use case as the foundation. While I may need further guidance on the rerank process, I feel confident in tackling the rest.

Looking forward to your thoughts and feedback!

[shamb0]

Hi @timonv,

I've resolved the issue using the HybridSearch strategy.
Now unit test and integration looks OK and available for first round of review.

Thankyou

[timonv]

Nice, this works!

Hybrid Search is a specific search method in RAG that combines keyword search with vector search, it's very effective, but this is not it. See i.e. https://towardsdatascience.com/how-to-use-hybrid-search-for-better-llm-rag-retrieval-032f66810ebe

I'd suggest to just stick to the single embedding, or create a SimilarityMultipleNamedEmbeddings strategy, or we could go lazy (not so lazy, it's not trival), and add a Custom strategy that takes a sqlx query, or a closure that returns a sqlx query with the nodes yieled? That would enable a lot of doors in fully controlling the search process. I'd suggest to start with the small thing first (single embedding), and if you're still interested, I'm excited to see what you can think of.

[timonv]

There is also a query_all, should be no need for a loop :-)

[shamb0]

Closing this in favor of implementing the suggested alternative search strategies.

[timonv]

We could still merge the single similarity, that looks great @shamb0

[shamb0]

We could still merge the single similarity, that looks great @shamb0

Thanks, @timonv! Apologies for overlooking your message earlier. I'll proceed with keeping only the 'Single Similarity' search integration for the intake merger.

[shamb0]

Hi @timonv, I've completed the integration of the retrieval feature using the Single Similarity search strategy.

I plan to integrated below test cases as part of SimilarityMultipleNamedEmbeddings search strategy implementation.

pgvector::tests::test_persist_nodes::both_mode_with_metadata
pgvector::tests::test_persist_nodes::both_mode_without_metadata
pgvector::tests::test_persist_nodes::perfield_mode_with_metadata
pgvector::tests::test_persist_nodes::perfield_mode_without_metadata

Please review and share any feedback you may have on the PR.

[timonv]

@shamb0 Awesome, code looks great!

Since it's full circle now, could you add a very simple index and query integration test in swiftide/tests? After that, looks good to merge. Nice job!

[shamb0]

Thanks @timonv, for the reminder about the integration test. I’ve implemented two test scenarios inspired by qdrant and lanchedb. Could you review them to ensure they align with the integration test scope and requirements? Let me know if any enhancements are needed.

[timonv]

Awesome! Looks great and thanks again for the contribution and hard work. I'll try to get the pipelines fixed and merge it asap 🎉

[shamb0]

Awesome! Looks great and thanks again for the contribution and hard work. I'll try to get the pipelines fixed and merge it asap 🎉

Thanks @timonv