Upload 7 files

chunking_parent.py

@@ -0,0 +1,79 @@
+# file: chunking.py
+import uuid
+from typing import List, Tuple, Dict, Any
+from langchain_core.documents import Document
+from langchain.storage import InMemoryStore
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+# --- Configuration for Parent-Child Splitting ---
+# Parent chunks are the larger documents passed to the LLM for context.
+PARENT_CHUNK_SIZE = 2000
+PARENT_CHUNK_OVERLAP = 200
+
+# Child chunks are the smaller, more granular documents used for retrieval.
+CHILD_CHUNK_SIZE = 400
+CHILD_CHUNK_OVERLAP = 100
+
+def create_parent_child_chunks(
+    full_text: str
+) -> Tuple[List[Document], InMemoryStore, Dict[str, str]]:
+    """
+    Implements the Parent Document strategy for chunking.
+
+    1. Splits the document into larger "parent" chunks.
+    2. Splits the parent chunks into smaller "child" chunks.
+    3. The child chunks are used for retrieval, while the parent chunks
+       are used to provide context to the LLM.
+
+    Args:
+        full_text: The entire text content of the document.
+
+    Returns:
+        A tuple containing:
+        - A list of the small "child" documents for the vector store.
+        - An in-memory store mapping parent document IDs to the parent documents.
+        - A dictionary mapping child document IDs to their parent's ID.
+    """
+    if not full_text:
+        print("Warning: Input text for chunking is empty.")
+        return [], InMemoryStore(), {}
+
+    print("Creating parent and child chunks...")
+    
+    # This splitter creates the large documents that will be stored.
+    parent_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=PARENT_CHUNK_SIZE,
+        chunk_overlap=PARENT_CHUNK_OVERLAP,
+    )
+
+    # This splitter creates the small, granular chunks for retrieval.
+    child_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=CHILD_CHUNK_SIZE,
+        chunk_overlap=CHILD_CHUNK_OVERLAP,
+    )
+
+    parent_documents = parent_splitter.create_documents([full_text])
+    
+    docstore = InMemoryStore()
+    child_documents = []
+    child_to_parent_id_map = {}
+
+    # Generate unique IDs for each parent document and add them to the store
+    parent_ids = [str(uuid.uuid4()) for _ in parent_documents]
+    docstore.mset(list(zip(parent_ids, parent_documents)))
+
+    # Split each parent document into smaller child documents
+    for i, p_doc in enumerate(parent_documents):
+        parent_id = parent_ids[i]
+        _child_docs = child_splitter.split_documents([p_doc])
+        
+        for _child_doc in _child_docs:
+            child_id = str(uuid.uuid4())
+            _child_doc.metadata["parent_id"] = parent_id
+            _child_doc.metadata["child_id"] = child_id
+            child_to_parent_id_map[child_id] = parent_id
+        
+        child_documents.extend(_child_docs)
+
+    print(f"Created {len(parent_documents)} parent chunks and {len(child_documents)} child chunks.")
+    return child_documents, docstore, child_to_parent_id_map

document_processor.py

@@ -0,0 +1,88 @@
+# file: document_processing.py
+
+import os
+import time
+import httpx
+from pathlib import Path
+from urllib.parse import urlparse, unquote
+from llama_index.readers.file import PyMuPDFReader
+from llama_index.core import Document as LlamaDocument
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pydantic import HttpUrl
+from typing import List
+
+# Define the batch size for parallel processing
+BATCH_SIZE = 25
+
+def _process_page_batch(documents_batch: List[LlamaDocument]) -> str:
+    """
+    Helper function to extract content from a batch of LlamaIndex Document objects
+    and join them into a single string.
+    """
+    return "\n\n".join([d.get_content() for d in documents_batch])
+
+async def ingest_and_parse_document(doc_url: HttpUrl) -> str:
+    """
+    Asynchronously downloads a document, saves it locally, and parses it to
+    Markdown text using PyMuPDFReader with parallel processing.
+
+    Args:
+        doc_url: The Pydantic-validated URL of the document to process.
+
+    Returns:
+        A single string containing the document's extracted text.
+    """
+    print(f"Initiating download from: {doc_url}")
+    LOCAL_STORAGE_DIR = "data/"
+    os.makedirs(LOCAL_STORAGE_DIR, exist_ok=True)
+
+    try:
+        # Asynchronously download the document
+        async with httpx.AsyncClient() as client:
+            response = await client.get(str(doc_url), timeout=30.0, follow_redirects=True)
+            response.raise_for_status()
+            doc_bytes = response.content
+        print("Download successful.")
+
+        # Determine a valid local filename
+        parsed_path = urlparse(str(doc_url)).path
+        filename = unquote(os.path.basename(parsed_path)) or "downloaded_document.pdf"
+        local_file_path = Path(os.path.join(LOCAL_STORAGE_DIR, filename))
+
+        # Save the document locally
+        with open(local_file_path, "wb") as f:
+            f.write(doc_bytes)
+        print(f"Document saved locally at: {local_file_path}")
+
+        # Parse the document using LlamaIndex's PyMuPDFReader
+        print("Parsing document with PyMuPDFReader...")
+        loader = PyMuPDFReader()
+        docs_from_loader = loader.load_data(file_path=local_file_path)
+
+        # Parallelize the extraction of text from loaded pages
+        start_time = time.perf_counter()
+        all_extracted_texts = []
+        with ThreadPoolExecutor(max_workers=os.cpu_count() or 4) as executor:
+            futures = [
+                executor.submit(_process_page_batch, docs_from_loader[i:i + BATCH_SIZE])
+                for i in range(0, len(docs_from_loader), BATCH_SIZE)
+            ]
+            for future in as_completed(futures):
+                all_extracted_texts.append(future.result())
+
+        doc_text = "\n\n".join(all_extracted_texts)
+        elapsed_time = time.perf_counter() - start_time
+        print(f"Time taken for parallel text extraction: {elapsed_time:.4f} seconds.")
+
+        if not doc_text:
+            raise ValueError("Document parsing yielded no content.")
+
+        print(f"Parsing complete. Extracted {len(doc_text)} characters.")
+        return doc_text
+
+    except httpx.HTTPStatusError as e:
+        print(f"Error downloading document: {e}")
+        raise
+    except Exception as e:
+        print(f"An unexpected error occurred during document processing: {e}")
+        raise

embedding.py

@@ -0,0 +1,40 @@
+# file: embedding.py
+
+import torch
+from sentence_transformers import SentenceTransformer
+from typing import List
+
+# --- Configuration ---
+EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
+
+class EmbeddingClient:
+    """A client for generating text embeddings using a local sentence transformer model."""
+
+    def __init__(self, model_name: str = EMBEDDING_MODEL_NAME):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = SentenceTransformer(model_name, device=self.device)
+        print(f"EmbeddingClient initialized with model '{model_name}' on device '{self.device}'.")
+
+    def create_embeddings(self, texts: List[str]) -> torch.Tensor:
+        """
+        Generates embeddings for a list of text chunks.
+
+        Args:
+            texts: A list of strings to be embedded.
+
+        Returns:
+            A torch.Tensor containing the generated embeddings.
+        """
+        if not texts:
+            return torch.tensor([])
+            
+        print(f"Generating embeddings for {len(texts)} text chunks on {self.device}...")
+        try:
+            embeddings = self.model.encode(
+                texts, convert_to_tensor=True, show_progress_bar=False
+            )
+            print("Embeddings generated successfully.")
+            return embeddings
+        except Exception as e:
+            print(f"An error occurred during embedding generation: {e}")
+            raise

generation.py

@@ -0,0 +1,57 @@
+# file: generation.py
+from groq import AsyncGroq
+from typing import List, Dict
+
+# --- Configuration ---
+GROQ_MODEL_NAME = "llama3-8b-8192"
+
+async def generate_answer(query: str, context_chunks: List[Dict], groq_api_key: str) -> str:
+    """
+    Generates a final answer using the Groq API based on the query and retrieved context.
+
+    Args:
+        query: The user's original question.
+        context_chunks: A list of the most relevant, reranked document chunks.
+        groq_api_key: The API key for the Groq service.
+
+    Returns:
+        A string containing the generated answer.
+    """
+    if not groq_api_key:
+        return "Error: Groq API key is not set."
+    if not context_chunks:
+        return "I do not have enough information to answer this question based on the provided document."
+
+    print("Generating final answer with Groq...")
+    client = AsyncGroq(api_key=groq_api_key)
+
+    # Format the context for the prompt
+    context_str = "\n\n---\n\n".join(
+        [f"Context Chunk:\n{chunk['content']}" for chunk in context_chunks]
+    )
+
+    prompt = (
+        "You are an expert Q&A system. Your task is to extract information with 100% accuracy from the provided text. Provide a brief and direct answer."
+        "Do not mention the context in your response. Answer *only* using the information from the provided document."
+        "Do not infer, add, or assume any information that is not explicitly written in the source text. If the answer is not in the document, state that the information is not available."
+        "When the question involves numbers, percentages, or monetary values, extract the exact figures from the text."
+        "Double-check that the value corresponds to the correct plan or condition mentioned in the question."
+        "\n\n"
+        f"CONTEXT:\n{context_str}\n\n"
+        f"QUESTION:\n{query}\n\n"
+        "ANSWER:"
+    )
+
+    try:
+        chat_completion = await client.chat.completions.create(
+            messages=[{"role": "user", "content": prompt}],
+            model=GROQ_MODEL_NAME,
+            temperature=0.2, # Lower temperature for more factual answers
+            max_tokens=500,
+        )
+        answer = chat_completion.choices[0].message.content
+        print("Answer generated successfully.")
+        return answer
+    except Exception as e:
+        print(f"An error occurred during Groq API call: {e}")
+        return "Could not generate an answer due to an API error."

main3.py

@@ -0,0 +1,123 @@
+# file: main.py
+import time
+import os
+import asyncio
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, HttpUrl
+from typing import List, Dict, Any
+from dotenv import load_dotenv
+
+from document_processor import ingest_and_parse_document
+from chunking_parent import create_parent_child_chunks # <-- Using the parent-child function
+from embedding import EmbeddingClient
+from retrieval_parent import Retriever, generate_hypothetical_document
+from generation import generate_answer
+
+load_dotenv()
+
+app = FastAPI(
+    title="Modular RAG API",
+    description="A modular API for Retrieval-Augmented Generation with Parent-Child Retrieval.",
+    version="2.2.1", # Updated version
+)
+
+GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
+embedding_client = EmbeddingClient()
+retriever = Retriever(embedding_client=embedding_client)
+
+# --- Pydantic Models ---
+class RunRequest(BaseModel):
+    document_url: HttpUrl
+    questions: List[str]
+
+class RunResponse(BaseModel):
+    answers: List[str]
+
+class TestRequest(BaseModel): # <-- Model for test endpoints
+    document_url: HttpUrl
+
+# --- NEW: Test Endpoint for Parent-Child Chunking ---
+@app.post("/test/chunk", response_model=Dict[str, Any], tags=["Testing"])
+async def test_chunking_endpoint(request: TestRequest):
+    """
+    Tests the parent-child chunking strategy.
+    Returns parent chunks, child chunks, and the time taken.
+    """
+    print("--- Running Parent-Child Chunking Test ---")
+    start_time = time.perf_counter()
+
+    try:
+        # Step 1: Parse the document to get raw text
+        markdown_content = await ingest_and_parse_document(request.document_url)
+        
+        # Step 2: Create parent and child chunks
+        child_documents, docstore, _ = create_parent_child_chunks(markdown_content)
+        
+        end_time = time.perf_counter()
+        duration = end_time - start_time
+        print(f"--- Parsing and Chunking took {duration:.2f} seconds ---")
+
+        # Convert Document objects to a JSON-serializable list for the response
+        child_chunk_results = [
+            {"page_content": doc.page_content, "metadata": doc.metadata}
+            for doc in child_documents
+        ]
+        
+        # Retrieve parent documents from the in-memory store
+        parent_docs = docstore.mget(list(docstore.store.keys()))
+        parent_chunk_results = [
+            {"page_content": doc.page_content, "metadata": doc.metadata}
+            for doc in parent_docs if doc
+        ]
+        
+        return {
+            "total_time_seconds": duration,
+            "parent_chunk_count": len(parent_chunk_results),
+            "child_chunk_count": len(child_chunk_results),
+            "parent_chunks": parent_chunk_results,
+            "child_chunks": child_chunk_results,
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"An error occurred during chunking test: {str(e)}")
+
+
+@app.post("/hackrx/run", response_model=RunResponse)
+async def run_rag_pipeline(request: RunRequest):
+    try:
+        print("--- Kicking off RAG Pipeline with Parent-Child Strategy ---")
+        
+        # --- STAGE 1: DOCUMENT INGESTION ---
+        markdown_content = await ingest_and_parse_document(request.document_url)
+        
+        # --- STAGE 2: PARENT-CHILD CHUNKING ---
+        child_documents, docstore, _ = create_parent_child_chunks(markdown_content)
+
+        if not child_documents:
+            raise HTTPException(status_code=400, detail="Document could not be processed into chunks.")
+
+        # --- STAGE 3: INDEXING ---
+        retriever.index(child_documents, docstore)
+
+        # --- CONCURRENT WORKFLOW ---
+        hyde_tasks = [generate_hypothetical_document(q, GROQ_API_KEY) for q in request.questions]
+        all_hyde_docs = await asyncio.gather(*hyde_tasks)
+
+        retrieval_tasks = [
+            retriever.retrieve(q, hyde_doc)
+            for q, hyde_doc in zip(request.questions, all_hyde_docs)
+        ]
+        all_retrieved_chunks = await asyncio.gather(*retrieval_tasks)
+
+        answer_tasks = [
+            generate_answer(q, chunks, GROQ_API_KEY)
+            for q, chunks in zip(request.questions, all_retrieved_chunks)
+        ]
+        final_answers = await asyncio.gather(*answer_tasks)
+
+        print("--- RAG Pipeline Completed Successfully ---")
+        return RunResponse(answers=final_answers)
+
+    except Exception as e:
+        raise HTTPException(
+            status_code=500, detail=f"An internal server error occurred: {str(e)}"
+        )

requirements.txt

@@ -0,0 +1,217 @@
+accelerate==1.10.0
+aiofiles==24.1.0
+aiohappyeyeballs==2.6.1
+aiohttp==3.12.15
+aiosignal==1.4.0
+aiosqlite==0.21.0
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.9.0
+asyncio==3.4.3
+attrs==25.3.0
+backoff==2.2.1
+banks==2.2.0
+beautifulsoup4==4.13.4
+bitarray==3.6.0
+blinker==1.9.0
+cachetools==5.5.2
+catalogue==2.0.10
+certifi==2025.8.3
+cffi==1.17.1
+charset-normalizer==3.4.2
+click==8.2.1
+colbert-ai==0.2.21
+colorama==0.4.6
+coloredlogs==15.0.1
+contourpy==1.3.3
+cryptography==45.0.6
+cycler==0.12.1
+dataclasses-json==0.6.7
+datasets==4.0.0
+defusedxml==0.7.1
+deprecated==1.2.18
+dill==0.3.8
+dirtyjson==1.0.8
+distro==1.9.0
+dotenv==0.9.9
+effdet==0.4.1
+emoji==2.14.1
+et-xmlfile==2.0.0
+faiss-cpu==1.11.0.post1
+fast-pytorch-kmeans==0.2.2
+fastapi==0.116.1
+filelock==3.18.0
+filetype==1.2.0
+flashrank==0.2.10
+flask==3.1.1
+flatbuffers==25.2.10
+fonttools==4.59.0
+frozenlist==1.7.0
+fsspec==2025.3.0
+git-python==1.0.3
+gitdb==4.0.12
+gitpython==3.1.45
+google-api-core==2.25.1
+google-auth==2.40.3
+google-cloud-vision==3.10.2
+googleapis-common-protos==1.70.0
+greenlet==3.2.3
+griffe==1.9.0
+groq==0.30.0
+grpcio==1.74.0
+grpcio-status==1.74.0
+h11==0.16.0
+hf-xet==1.1.5
+html5lib==1.1
+httpcore==1.0.9
+httptools==0.6.4
+httpx==0.28.1
+httpx-sse==0.4.1
+huggingface-hub==0.34.3
+humanfriendly==10.0
+idna==3.10
+itsdangerous==2.2.0
+jinja2==3.1.6
+jiter==0.10.0
+joblib==1.5.1
+jsonpatch==1.33
+jsonpointer==3.0.0
+kiwisolver==1.4.8
+langchain==0.3.27
+langchain-community==0.3.27
+langchain-core==0.3.72
+langchain-pymupdf4llm==0.4.1
+langchain-text-splitters==0.3.9
+langdetect==1.0.9
+langsmith==0.4.10
+llama-cloud==0.1.35
+llama-cloud-services==0.6.53
+llama-index==0.13.0
+llama-index-cli==0.5.0
+llama-index-core==0.13.0
+llama-index-embeddings-openai==0.5.0
+llama-index-indices-managed-llama-cloud==0.9.0
+llama-index-instrumentation==0.4.0
+llama-index-llms-openai==0.5.1
+llama-index-readers-file==0.5.0
+llama-index-readers-llama-parse==0.5.0
+llama-index-workflows==1.2.0
+llama-parse==0.6.53
+lxml==6.0.0
+markdown==3.8.2
+markupsafe==3.0.2
+marshmallow==3.26.1
+matplotlib==3.10.5
+mpmath==1.3.0
+msoffcrypto-tool==5.4.2
+multidict==6.6.3
+multiprocess==0.70.16
+mypy-extensions==1.1.0
+nest-asyncio==1.6.0
+networkx==3.5
+ninja==1.11.1.4
+nltk==3.9.1
+numpy==2.3.2
+olefile==0.47
+omegaconf==2.3.0
+onnx==1.18.0
+onnxruntime==1.22.1
+openai==1.99.3
+opencv-python==4.11.0.86
+openpyxl==3.1.5
+orjson==3.11.1
+packaging==25.0
+pandas==2.2.3
+pdf2image==1.17.0
+pdfminer-six==20250506
+pi-heif==1.1.0
+pikepdf==9.10.2
+pillow==11.3.0
+pinecone-client==6.0.0
+pinecone-plugin-interface==0.0.7
+platformdirs==4.3.8
+propcache==0.3.2
+proto-plus==1.26.1
+protobuf==6.31.1
+psutil==7.0.0
+pyarrow==21.0.0
+pyasn1==0.6.1
+pyasn1-modules==0.4.2
+pycocotools==2.0.10
+pycparser==2.22
+pydantic==2.11.7
+pydantic-core==2.33.2
+pydantic-settings==2.10.1
+pymupdf==1.26.3
+pymupdf4llm==0.0.27
+pypandoc==1.15
+pyparsing==3.2.3
+pypdf==5.9.0
+pypdf2==3.0.1
+pypdfium2==4.30.0
+pytesseract==0.3.13
+python-dateutil==2.9.0.post0
+python-docx==1.2.0
+python-dotenv==1.1.1
+python-iso639==2025.2.18
+python-magic==0.4.27
+python-multipart==0.0.20
+python-oxmsg==0.0.2
+python-pptx==1.0.2
+pytz==2025.2
+pyyaml==6.0.2
+ragatouille==0.0.9.post2
+rank-bm25==0.2.2
+rapidfuzz==3.13.0
+regex==2025.7.34
+requests==2.32.4
+requests-toolbelt==1.0.0
+rsa==4.9.1
+safetensors==0.5.3
+scikit-learn==1.7.1
+scipy==1.16.1
+sentence-transformers==5.0.0
+setuptools==80.9.0
+six==1.17.0
+smmap==5.0.2
+sniffio==1.3.1
+soupsieve==2.7
+sqlalchemy==2.0.42
+srsly==2.5.1
+starlette==0.47.2
+striprtf==0.0.26
+sympy==1.14.0
+tenacity==9.1.2
+tesseract==0.1.3
+threadpoolctl==3.6.0
+tiktoken==0.9.0
+timm==1.0.19
+tokenizers==0.21.4
+torch==2.7.1
+torchvision==0.22.1
+tqdm==4.67.1
+transformers==4.49.0
+typing==3.7.4.3
+typing-extensions==4.14.1
+typing-inspect==0.9.0
+typing-inspection==0.4.1
+tzdata==2025.2
+ujson==5.10.0
+unstructured==0.18.11
+unstructured-client==0.42.2
+unstructured-inference==1.0.5
+unstructured-pytesseract==0.3.15
+urllib3==2.5.0
+uvicorn==0.35.0
+uvloop==0.21.0
+voyager==2.1.0
+watchfiles==1.1.0
+webencodings==0.5.1
+websockets==15.0.1
+werkzeug==3.1.3
+wrapt==1.17.2
+xlrd==2.0.2
+xlsxwriter==3.2.5
+xxhash==3.5.0
+yarl==1.20.1
+zstandard==0.23.0

retrieval_parent.py

@@ -0,0 +1,155 @@
+# file: retrieval.py
+
+import time
+import asyncio
+import numpy as np
+import torch
+from groq import AsyncGroq
+from rank_bm25 import BM25Okapi
+from sentence_transformers import CrossEncoder
+from sklearn.preprocessing import MinMaxScaler
+from torch.nn.functional import cosine_similarity
+from typing import List, Dict, Tuple
+from langchain.storage import InMemoryStore
+
+from embedding import EmbeddingClient
+from langchain_core.documents import Document
+
+# --- Configuration ---
+HYDE_MODEL = "llama3-8b-8192"
+RERANKER_MODEL = 'cross-encoder/ms-marco-MiniLM-L6-v2'
+INITIAL_K_CANDIDATES = 20
+TOP_K_CHUNKS = 10 
+
+async def generate_hypothetical_document(query: str, groq_api_key: str) -> str:
+    # ... (this function remains unchanged) ...
+    if not groq_api_key:
+        print("Groq API key not set. Skipping HyDE generation.")
+        return ""
+
+    print(f"Starting HyDE generation for query: '{query}'...")
+    client = AsyncGroq(api_key=groq_api_key)
+    prompt = (
+        f"Write a brief, formal passage that answers the following question. "
+        f"Use specific terminology as if it were from a larger document. "
+        f"Do not include the question or conversational text.\n\n"
+        f"Question: {query}\n\n"
+        f"Hypothetical Passage:"
+    )
+
+    try:
+        chat_completion = await client.chat.completions.create(
+            messages=[{"role": "user", "content": prompt}],
+            model=HYDE_MODEL,
+            temperature=0.7,
+            max_tokens=500,
+        )
+        return chat_completion.choices[0].message.content
+    except Exception as e:
+        print(f"An error occurred during HyDE generation: {e}")
+        return ""
+
+
+class Retriever:
+    """Manages hybrid search with parent-child retrieval."""
+
+    def __init__(self, embedding_client: EmbeddingClient):
+        self.embedding_client = embedding_client
+        self.reranker = CrossEncoder(RERANKER_MODEL, device=self.embedding_client.device)
+        self.bm25 = None
+        self.document_chunks = []
+        self.chunk_embeddings = None
+        self.docstore = InMemoryStore() # <-- ADD THIS
+        print(f"Retriever initialized with reranker '{RERANKER_MODEL}'.")
+
+    def index(self, child_documents: List[Document], docstore: InMemoryStore): # <-- MODIFY THIS
+        """Builds the search index from child documents and stores parent documents."""
+        self.document_chunks = child_documents # Store child docs for mapping
+        self.docstore = docstore # Store the parent documents
+        
+        corpus = [doc.page_content for doc in child_documents]
+        if not corpus:
+            print("No documents to index.")
+            return
+
+        print("Indexing child documents for retrieval...")
+        tokenized_corpus = [doc.split(" ") for doc in corpus]
+        self.bm25 = BM25Okapi(tokenized_corpus)
+        self.chunk_embeddings = self.embedding_client.create_embeddings(corpus)
+        print("Indexing complete.")
+
+    def _hybrid_search(self, query: str, hyde_doc: str) -> List[Tuple[int, float]]:
+        # ... (this function remains unchanged) ...
+        if self.bm25 is None or self.chunk_embeddings is None:
+            raise ValueError("Retriever has not been indexed. Call index() first.")
+
+        enhanced_query = f"{query}\n\n{hyde_doc}" if hyde_doc else query
+        tokenized_query = query.split(" ")
+        bm25_scores = self.bm25.get_scores(tokenized_query)
+        query_embedding = self.embedding_client.create_embeddings([enhanced_query])
+        dense_scores = cosine_similarity(query_embedding, self.chunk_embeddings).cpu().numpy().flatten()
+
+        scaler = MinMaxScaler()
+        norm_bm25 = scaler.fit_transform(bm25_scores.reshape(-1, 1)).flatten()
+        norm_dense = scaler.fit_transform(dense_scores.reshape(-1, 1)).flatten()
+        combined_scores = 0.5 * norm_bm25 + 0.5 * norm_dense
+        
+        top_indices = np.argsort(combined_scores)[::-1][:INITIAL_K_CANDIDATES]
+        return [(idx, combined_scores[idx]) for idx in top_indices]
+
+
+    async def _rerank(self, query: str, candidates: List[dict]) -> List[dict]:
+        # ... (this function remains unchanged) ...
+        if not candidates:
+            return []
+
+        print(f"Reranking {len(candidates)} candidates...")
+        rerank_input = [[query, chunk["content"]] for chunk in candidates]
+        
+        rerank_scores = await asyncio.to_thread(
+            self.reranker.predict, rerank_input, show_progress_bar=False
+        )
+
+        for candidate, score in zip(candidates, rerank_scores):
+            candidate['rerank_score'] = score
+        
+        candidates.sort(key=lambda x: x['rerank_score'], reverse=True)
+        return candidates[:TOP_K_CHUNKS]
+
+
+    async def retrieve(self, query: str, hyde_doc: str) -> List[Dict]: # <-- MODIFY THIS
+        """Executes the full retrieval pipeline and returns parent documents."""
+        print(f"Retrieving documents for query: '{query}'")
+        
+        # 1. Hybrid search returns indices of the best CHILD documents
+        initial_candidates_info = self._hybrid_search(query, hyde_doc)
+        
+        retrieved_child_docs = [{
+            "content": self.document_chunks[idx].page_content,
+            "metadata": self.document_chunks[idx].metadata,
+        } for idx, score in initial_candidates_info]
+
+        # 2. Rerank the CHILD documents
+        reranked_child_docs = await self._rerank(query, retrieved_child_docs)
+
+        # 3. Get the unique parent IDs from the reranked child documents
+        parent_ids = []
+        for doc in reranked_child_docs:
+            parent_id = doc["metadata"]["parent_id"]
+            if parent_id not in parent_ids:
+                parent_ids.append(parent_id)
+
+        # 4. Retrieve the full PARENT documents from the docstore
+        retrieved_parents = self.docstore.mget(parent_ids)
+        
+        # Filter out any None results in case of a miss
+        final_parent_docs = [doc for doc in retrieved_parents if doc is not None]
+
+        # 5. Format for the generation step
+        final_context = [{
+            "content": doc.page_content,
+            "metadata": doc.metadata
+        } for doc in final_parent_docs]
+
+        print(f"Retrieved {len(final_context)} final parent chunks for context.")
+        return final_context