Building a Self-Reflective RAG Pipeline: From Smart Chunking to Insightful Retrieval — Part 2

Dual-Store Retrieval, Cross-Encoder Re-ranking & Context Compression

This section explores how to retrieve relevant chunks during a query. It describes how Hybrid search (dense + sparse) handles synonyms and exact keywords. A cross-encoder then re-ranks the merged results, and a lightweight LLM compresses the final context sent to the answer-generation step. It shows the benefits of the smart chunking and storage strategies introduced in Part 1.

Generated using DALL-E

Retrieval-Augmented Generation (RAG) pipelines are becoming essential in building intelligent applications that can retrieve and reason over large volumes of unstructured data. To make such pipelines efficient and scalable, it is important to design them with clarity, modularity, and long-term maintainability in mind.

This blog series is structured into four parts, offering a practical, step-by-step guide to building a self-reflective RAG pipeline — from preparing high-quality chunks to designing intelligent retrieval strategies and scaling the system for real-world applications:

1. **Part 1: Smart Chunking & Storage Fundamentals**
2. Part 2: Dual-Store Retrieval: Child, QA and Re-ranking (This article)
3. **Part 3: Self-Reflective RAG Techniques**
4. **Part 4: End-to-End Walk-through, Evaluation & Scaling**

Each part focuses on a critical aspect of the pipeline, helping developers build a robust system that goes beyond basic retrieval to deliver insightful and context-aware results.

Introduction

Child chunks are smaller segments of text optimized for precise semantic retrieval, while QA chunks are pre-generated pairs of questions and answers designed to handle queries phrased as natural questions. In this stage of the pipeline, we describe how an incoming user query is answered by retrieving relevant information from two specialized stores and then refining the ranked results. Initially, the query is converted into a semantic embedding (e.g., using a transformer model) and used in a hybrid search (dense + sparse) that combines both dense and sparse techniques. The hybrid similarity search (*alpha=0.5*) evenly balances dense semantic similarity and sparse keyword matching, ensuring comprehensive retrieval coverage. Practically, this approach performs a vector similarity search (dense search) in parallel with a traditional keyword search (such as BM25, sparse search), and merging these results into a combined candidate set. The dense vector search captures semantic meaning, while the sparse keyword search captures precise textual matches. Finally, A cross-encoder then re-ranks the merged results, and a lightweight LLM compresses the final context sent to the answer-generation step.

Workflow Diagram

Parallel Retrieval from Child-Chunk and QA-Chunk Stores

The pipeline retrieves relevant content from two specialized Weaviate vector stores simultaneously: one indexing child chunks and another indexing pre-generated question-answer (QA) pairs. This dual retrieval ensures comprehensive coverage, enabling the retrieval of content that matches either direct textual fragments or more naturally phrased queries.

Child Chunk Retrieval:

The child chunk retrieval leverages Weaviate’s vector store, instantiated via *WeaviateVectorStore*. The retrieval process starts by embedding the user’s query using the HuggingFace embedding model (*EMBEDDING_MODEL*) running on GPU if available (device = “cuda” if cuda.is_available() else “cpu”). These embeddings are normalized to enhance search consistency.

The child store is initialized with specific attributes, including the file and parent identifiers:

child_store = WeaviateVectorStore(
    client=weaviate_client,
    index_name=CHILD_INDEX,
    text_key="text",
    embedding=emb,
    attributes=[file_id_key, parent_id_key],
)

A hybrid similarity search (*similarity_search*) is executed using the user query to retrieve the top five matching child chunks:

docs = child_store.similarity_search(state["question"], k=5, alpha=0.5)

These results, containing child chunks along with their associated parent identifiers, are stored for subsequent re-ranking.

QA Chunk Retrieval:

In parallel, the pipeline performs a similar retrieval from the QA chunk store. The QA chunks are retrieved through a separate Weaviate store specifically indexing generated QA pairs, enhancing the pipeline’s ability to handle queries phrased naturally as questions:

qa_store = WeaviateVectorStore(
    client=weaviate_client,
    index_name=QA_INDEX,
    text_key="text",
    embedding=emb,
    attributes=[file_id_key, parent_id_key],
)

The QA retrieval also uses hybrid similarity search with the same embedding and parameters as the child chunk retrieval:

docs = qa_store.similarity_search(state["question"], k=5, alpha=0.5)

The retrieved QA chunks similarly include metadata linking back to their parent content.

Both retrieval processes use robust resource management (try…finally) to ensure the Weaviate client connection is properly closed after the retrieval, thereby avoiding resource leaks.

The parallel execution of these retrieval paths significantly boosts the pipeline’s effectiveness, ensuring relevant content is captured comprehensively before proceeding to the cross-encoder re-ranking stage.

Tip –Tune ALPHA on a held-out question set (0 → keyword-only, 1 → dense-only).

Cross-Encoder Re-ranking of Top Candidates

After retrieving child and QA chunks, the pipeline applies a precise cross-encoder re-ranking process to refine the results. The cross-encoder utilized here is the pre-trained model *cross-encoder/ms-marco-MiniLM-L-6-v2*, which is optimized for evaluating the relevance of query-passage pairs.

The re-ranking process involves the following steps:

1. Pair Formation: Query-passage pairs are formed by combining the user query with each retrieved chunk (both child and QA chunks):

pairs = [[query, doc.page_content] for doc in all_dox]

2. Relevance Scoring: Each pair is then evaluated by the cross-encoder to generate relevance scores:

scores = cross_encoder.predict(pairs)

3. Ranking: The scores are sorted in descending order to identify the top candidates. This is done efficiently using NumPy sorting:

sorted_indices = np.argsort(-scores)

4. Top Selection: The top 5 highest-scoring chunks are selected as the most relevant:

top_ranked_dox = []
for i in sorted_indices[:5]:
    top_ranked_dox.append(all_dox[i])

5. Parent ID Extraction: Finally, unique parent document identifiers (*doc_id*) from the metadata of these top-ranked chunks are extracted and deduplicated to ensure each parent is considered only once:

parent_ids = [str((result.metadata["doc_id"])) for result in top_ranked_dox]
parent_ids = list(set(parent_ids))

These refined parent IDs are then passed forward for the final parent chunk retrieval, ensuring the pipeline’s efficiency and accuracy in delivering highly relevant and context-rich information to the user.

Parent Chunk Aggregation and Retrieval

MongoDB provides efficient querying capabilities and scalability for retrieving large parent chunks, enabling fast and reliable access during pipeline execution. After determining the most relevant chunks through cross-encoder re-ranking, the pipeline proceeds to fetch the comprehensive parent chunks from a MongoDB collection. This ensures the final output provided to the user contains rich, contextual information derived from larger text segments.

The retrieval process utilizes a MongoDB query leveraging the unique parent document identifiers obtained during the re-ranking stage. The process is executed as follows:

1. Establish Database Connection: The pipeline securely connects to MongoDB using *MongoClient* within a context manager to ensure proper handling and closing of the database connection:

with MongoClient(MONGO_URI) as client:

2. Query Execution: It retrieves the relevant parent chunks by matching their unique identifiers (*parent_id_key*) against the retrieved list of parent IDs:

coll = client[MONGO_DB_NAME][MONGO_PARENT_COLLECTION]
cursor = coll.find(
         { parent_id_key: { "$in": state["parent_ids"] } },
         {
             "_id": False,
             parent_id_key: True,
             "content": True,
             "source": True
         }
     )

3. Content Aggregation: The content from the retrieved parent chunks is concatenated into a single continuous string to ensure easy handling for subsequent pipeline processing:

documents = [
         Document(
             page_content=doc["content"],
             metadata={
                 parent_id_key: doc[parent_id_key],
                 "source": doc.get("source"),
             }
         )
         for doc in cursor
     ]

4. Pipeline Continuation: The aggregated content is then prepared for the next step, typically content compression:

return Command(
    update={
        "parent_docs": documents,
    },
    goto="compress"
)

This structured approach efficiently retrieves and aggregates parent chunks, providing extensive context necessary for generating accurate, comprehensive answers.

Context Compression

After retrieving the relevant parent chunks, the pipeline applies context compression using a Large Language Model (LLM) to extract and summarize the most pertinent information for the user’s query. This step significantly improves efficiency by converting extensive content into concise and focused snippets relevant to the query.

The pipeline leverages LangChain’s *LLMChainExtractor*, utilizing OpenAI’s Chat model (*ChatOpenAI*) configured with a specified GPT model (*GPT_MODEL*) and a deterministic temperature setting (*0*).

The updated compression process includes parallel processing for enhanced performance:

1. Initialization of LLM:

llm = ChatOpenAI(model_name=GPT_MODEL, temperature=0, openai_api_key=key, cache=False)

2. Compressor Setup:

Instantiate the *LLMChainExtractor*:

compressor = LLMChainExtractor.from_llm(llm=llm)

*LLMChainExtractor* uses an LLM to analyze document content and the query to generate concise summaries, significantly reducing context while retaining the most relevant information.

3. Parallel Document Compression:

Parent documents are compressed concurrently using Python’s *ThreadPoolExecutor*, significantly speeding up the compression process:

def compress_context(state: RetrieverState) -> Command:
    query = state["question"]
    parent_docs = state["parent_docs"]

    def compress_document(doc):
        return compressor.compress_documents([doc], query)

    compressed_docs = []

    with ThreadPoolExecutor() as executor:
        future_to_doc = {executor.submit(compress_document, doc): doc for doc in parent_docs}

        for future in as_completed(future_to_doc):
            actual = future.result()
            if len(actual) > 0:
                compressed_docs.extend(actual)

    return Command(
        update={
            "compressed_docs": compressed_docs,
        },
    )

4. Updating Pipeline State:

The compressed documents are stored in the pipeline state, ensuring subsequent processes have streamlined and relevant data:

return Command(
    update={
        "compressed_docs": compressed_docs,
    },
)

This compression approach optimizes the pipeline’s efficiency and accuracy, enabling delivery of succinct, context-rich answers and a more effective user experience.

Building the Retrieval Workflow Graph with LangGraph

from langgraph.graph import StateGraph
from rag.retrieval.retriever_state import RetrieverState
from rag.retrieval.search_child_node import get_child_chunks
from rag.retrieval.search_qa_node import get_qa_chunks
from rag.retrieval.re_ranking_node import cross_encoder_re_rank
from rag.retrieval.search_parent_node import get_contents_by_parent_id
from rag.retrieval.contextual_compressor_node import compress_context

def generate_graph():
    workflow = StateGraph(RetrieverState)

    workflow.add_node("child", get_child_chunks)
    workflow.add_node("qa", get_qa_chunks)
    workflow.add_node("re_ranking", cross_encoder_re_rank)
    workflow.add_node("parent", get_contents_by_parent_id)
    workflow.add_node("compress", compress_context)

    workflow.set_entry_point("child")
    workflow.set_entry_point("qa")

    workflow.set_finish_point("compress")
    # memory = SqliteSaver.from_conn_string(":memory:")
    chain = workflow.compile()
    # chain = workflow.compile(checkpointer=memory, interrupt_before=["save"])
    return chain

Executing the Retrieval Workflow Graph

# execute_graph.py
from rag.retrieval.graph_generator import generate_graph

query = "Explain Strategic Design of DDD"

inputs = {"question": query}
config = {"recursion_limit": 50}

graph = generate_graph()

output = graph.invoke(inputs, config=config)

print(f'Context: {output["compressed_docs"]}')

# Benefits of the Dual-Store Architecture

This **“dual-store**” design — with separate indexes for child chunks and QA chunks, plus a distinct store for parent chunks — offers multiple advantages. First, by splitting documents into small child chunks for embedding, we enable very precise matching: smaller text pieces yield more focused semantic vectors. At the same time, by linking back to larger parents, we recover full context when needed. As Dify’s parent-child retrieval paper notes[1], using child chunks for query matching and then retrieving their enclosing parent sections **balances the precision–context tradeoff**. Searching a QA chunk store adds another layer of robustness: queries phrased in natural question form can match pre-generated QA pairs, improving recall when the same information might not be retrieved by direct text search.

Second, this approach leverages hybrid search to handle diverse query types. The sparse (keyword) index catches exact terms, and the dense (embedding) index captures synonyms and general meaning. By fusing both, we avoid missing answers that would slip through if we used only one method. Finally, the cross-encoder re-ranking ensures that the final top hits are of highest relevance. In summary, the dual-store pipeline allows us to efficiently search vast knowledge bases while providing the LLM with both pinpoint relevance and rich context for answer generation — improving accuracy and completeness in the RAG system.

# Source Code

**Clone the repository:**

git clone https://github.com/mail2mhossain/self_reflective_rag.git cd self_reflective_rag

**Create a Conda environment (Assuming Anaconda is installed):**

conda create -n self_reflective_rag_env python=3.11

**Activate the environment:**

conda activate self_reflective_rag_env

4. **Install the required packages:**

pip install torch==2.5.1 torchvision torchaudio - index-url https://download.pytorch.org/whl/cu121 pip install git+https://github.com/huggingface/transformers pip install git+https://github.com/huggingface/accelerate pip install -r requirements.txt

5. **Run the Retrieval Workflow Graph:**

python -m rag.retrieval.execute_graph [from root directory]

**To remove the environment after use:**

conda remove - name self_reflective_rag_env - all

**Reference [1]** [Dify-AI, “Enhancing Retrieval with Parent–Child Chunking,” 2024](https://dify.ai/blog/introducing-parent-child-retrieval-for-enhanced-knowledge).