Skip to content

feat(rag): add hybrid dense + sparse retrieval pipeline#1283

Open
Avi-47 wants to merge 1 commit intomofa-org:mainfrom
Avi-47:feat/rag-hybrid-retrieval
Open

feat(rag): add hybrid dense + sparse retrieval pipeline#1283
Avi-47 wants to merge 1 commit intomofa-org:mainfrom
Avi-47:feat/rag-hybrid-retrieval

Conversation

@Avi-47
Copy link
Contributor

@Avi-47 Avi-47 commented Mar 16, 2026

Summary

This PR introduces a hybrid retrieval pipeline for MoFA’s RAG system that combines:

  • Dense semantic retrieval (vector similarity)
  • Sparse keyword retrieval (BM25)

The results from both retrievers are fused using Reciprocal Rank Fusion (RRF) to produce a unified ranking.

Hybrid retrieval significantly improves recall and robustness in RAG pipelines by leveraging both semantic similarity and exact keyword matching.

This work builds on the previously introduced BM25 sparse retrieval module and extends the MoFA RAG system toward production-grade hybrid search.

Relates to #305
Depends on: #1266 (BM25 sparse retrieval)

TL;DR

Adds hybrid dense + sparse retrieval to the MoFA RAG pipeline.

Key additions:
• HybridRetriever trait (kernel)
• HybridSearchPipeline implementation (foundation)
• Reciprocal Rank Fusion (RRF)
• Parallel retrieval using tokio::join!


Motivation

The current MoFA RAG pipeline primarily relies on dense vector similarity. While effective for semantic search, dense retrieval may miss:

  • exact keyword matches
  • domain-specific terminology
  • rare tokens poorly represented in embeddings

Sparse retrieval (BM25) complements dense retrieval by providing:

  • deterministic keyword search
  • interpretable scoring
  • strong performance on rare terms

Hybrid retrieval combines both approaches to achieve higher recall and better ranking quality.


Architecture

The implementation follows MoFA’s microkernel architecture:

query
 ↓
dense retriever (VectorStore search)
 ↓
BM25 retriever
 ↓
Reciprocal Rank Fusion (RRF)
 ↓
top-k fused results

Hybrid Retrieval Architecture Diagram

┌─────────────────────────────────────────────────────────────────────────────┐
│                           Hybrid Retrieval Pipeline                         │
└─────────────────────────────────────────────────────────────────────────────┘

                                    ┌─────────────┐
                                    │   Query     │
                                    └──────┬──────┘
                                           │
                    ┌──────────────────────┼──────────────────────┐
                    │                      │                      │
                    ▼                      ▼                      ▼
         ┌──────────────────┐  ┌──────────────────┐  ┌──────────────────┐
         │   Dense Retriever│  │  Sparse Retriever│  │   Embedding      │
         │  (Vector Store)  │  │     (BM25)       │  │   Adapter        │
         └────────┬─────────┘  └────────┬─────────┘  └──────────────────┘
                  │                     │
                  │                     │
                  ▼                     ▼
         ┌──────────────────┐  ┌──────────────────┐
         │   Ranked Docs    │  │   Ranked Docs    │
         │   [d1, d2, d3]   │  │   [s1, s2, s3]   │
         └────────┬─────────┘  └────────┬─────────┘
                  │                     │
                  └──────────┬──────────┘
                             │
                             ▼
              ┌───────────────────────────────┐
              │   Reciprocal Rank Fusion      │
              │   Formula: score = 1/(k+rank) │
              │   Default k = 60            │
              └───────────────┬───────────────┘
                              │
                              ▼
                   ┌─────────────────────┐
                   │  Fused & Re-ranked  │
                   │     Documents       │
                   └──────────┬──────────┘
                              │
                              ▼
                   ┌─────────────────────┐
                   │      top_k          │
                   │    Results          │
                   └─────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│                            Module Structure                                 │
└─────────────────────────────────────────────────────────────────────────────┘

mofa-kernel (Traits/Interfaces)
├── rag/
│   ├── hybrid.rs          → HybridRetriever trait
│   │                         - retrieve(query, top_k)
│   │                         - retrieve_with_rrf(query, top_k, rrf_k)
│   ├── types.rs           → ScoredDocument, Document
│   └── vector_store.rs    → VectorStore trait

mofa-foundation (Implementations)
└── rag/
    └── hybrid/
        ├── mod.rs             → Exports
        ├── rrf.rs             → reciprocal_rank_fusion()
        │                         - DEFAULT_RRF_K = 60
        │                         - Merges multiple ranked lists
        └── hybrid_retriever.rs → HybridSearchPipeline
                                      - dense_retriever: Arc<dyn VectorStore>
                                      - sparse_retriever: Arc<dyn Retriever>
                                      - embedding_adapter: LlmEmbeddingAdapter
                                  → GenericHybridRetriever
                                      - Works with any two retrievers

┌─────────────────────────────────────────────────────────────────────────────┐
│                         RRF Score Calculation                               │
└─────────────────────────────────────────────────────────────────────────────┘

Given two ranked lists:
  Dense results: [doc_A, doc_B, doc_C]
  Sparse results: [doc_B, doc_A, doc_D]

RRF Score for each document:
  doc_A: 1/(0+60) + 1/(1+60) = 0.0164 + 0.0164 = 0.0328
  doc_B: 1/(1+60) + 1/(0+60) = 0.0164 + 0.0164 = 0.0328  
  doc_C: 1/(2+60) = 0.0161
  doc_D: 1/(2+60) = 0.0161

Final ranking (descending by RRF score):
  doc_A, doc_B > doc_C, doc_D

Kernel Layer

crates/mofa-kernel/src/rag/hybrid.rs

Adds a new trait:

trait HybridRetriever {
    async fn retrieve(&self, query: &str, top_k: usize)
        -> AgentResult<Vec<ScoredDocument>>;
}

This defines the abstraction for hybrid retrieval pipelines.


Foundation Layer

New module:

crates/mofa-foundation/src/rag/hybrid/

Components:

HybridSearchPipeline

Combines:

  • dense vector search via VectorStore
  • sparse retrieval via Retriever

Runs both retrievers in parallel using tokio::join!.

GenericHybridRetriever

A simpler hybrid retriever that combines any two Retriever implementations.

Reciprocal Rank Fusion (RRF)

Implemented in:

rrf.rs

Formula:

score = 1 / (rank + k)

Where:

  • rank = document position in each result list
  • k = smoothing parameter (default 60)

RRF is widely used in production hybrid search systems.


Key Features

  • hybrid dense + sparse retrieval
  • parallel retrieval execution
  • Reciprocal Rank Fusion ranking
  • configurable fusion parameters
  • plug-and-play with existing RAG pipeline
  • no breaking API changes

Files Added

crates/mofa-kernel/src/rag/hybrid.rs

crates/mofa-foundation/src/rag/hybrid/
  hybrid_retriever.rs
  rrf.rs
  mod.rs

Updated exports:

crates/mofa-kernel/src/rag/mod.rs
crates/mofa-foundation/src/rag/mod.rs

Example Usage

use mofa_foundation::rag::{
    hybrid::HybridSearchPipeline,
    InMemoryVectorStore,
    Bm25Retriever
};

let hybrid = HybridSearchPipeline::new(
    dense_store,
    sparse_retriever,
    embedder
);

let results = hybrid.retrieve("distributed systems architecture", 5).await?;

Testing

Verification steps:

cargo check
cargo test
cargo clippy --workspace --all-features

Results:

  • All 1114 workspace tests pass
  • No new clippy warnings
  • Hybrid module unit tests included

Impact

This PR enables hybrid retrieval in MoFA, which is a key capability for production RAG systems.

Benefits:

  • improved recall
  • better ranking quality
  • robust search for mixed semantic + keyword queries

Future Work

Possible follow-ups:

  • pgvector vector store backend
  • cross-encoder reranking
  • query expansion
  • hybrid retrieval configuration in CLI

Breaking changes: None
Existing APIs remain unchanged.

@Avi-47 Avi-47 marked this pull request as ready for review March 16, 2026 03:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant