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

embedding.py

@@ -0,0 +1,40 @@
+# file: embedding.py
+
+import torch
+from sentence_transformers import SentenceTransformer
+from typing import List
+
+# --- Configuration ---
+EMBEDDING_MODEL_NAME = "sentence-transformers/stsb-xlm-r-multilingual"
+
+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."

ingestion_router.py

@@ -0,0 +1,129 @@
+# file: ingestion_router.py
+import os
+import time
+import httpx
+import zipfile
+import io
+import asyncio
+from PIL import Image
+from pathlib import Path
+from urllib.parse import urlparse, unquote
+from pydantic import HttpUrl
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+# --- Import our custom parsers ---
+from pdf_parallel_parser import process_pdf_with_hybrid_parallel_sync
+from complex_parser import process_image_element
+
+# --- A simple parser for generic files (DOCX, etc.) using unstructured ---
+from unstructured.partition.auto import partition
+
+# --- Configuration ---
+LOCAL_STORAGE_DIR = "data/"
+
+# --- Synchronous, CPU-Bound Parsing Functions ---
+
+def _process_generic_file_sync(file_content: bytes, filename: str) -> str:
+    """Fallback parser for standard files like DOCX, PPTX, etc., using unstructured."""
+    print(f"Processing '{filename}' with unstructured (standard)...")
+    try:
+        elements = partition(file=io.BytesIO(file_content), file_filename=filename)
+        return "\n\n".join([el.text for el in elements])
+    except Exception as e:
+        print(f"Unstructured failed for {filename}: {e}")
+        return ""
+
+def _process_zip_file_in_parallel(zip_content: bytes, temp_dir: Path) -> str:
+    """Extracts and processes files from a ZIP archive in parallel."""
+    print("Initiating parallel processing of ZIP archive...")
+    all_extracted_texts = []
+
+    def process_single_zipped_file(zf: zipfile.ZipFile, file_info: zipfile.ZipInfo) -> str:
+        file_content = zf.read(file_info.filename)
+        file_extension = Path(file_info.filename).suffix.lower()
+        
+        # Route to the appropriate synchronous parser
+        if file_extension == '.pdf':
+            temp_file_path = temp_dir / Path(file_info.filename).name
+            temp_file_path.parent.mkdir(parents=True, exist_ok=True)
+            temp_file_path.write_bytes(file_content)
+            return process_pdf_with_hybrid_parallel_sync(temp_file_path)
+        elif file_extension in ['.png', '.jpg', '.jpeg']:
+            return process_image_element(Image.open(io.BytesIO(file_content)))
+        elif file_extension in ['.docx', '.pptx', '.html']:
+            return _process_generic_file_sync(file_content, file_info.filename)
+        else:
+            print(f"Skipping unsupported file in ZIP: {file_info.filename}")
+            return ""
+
+    with zipfile.ZipFile(io.BytesIO(zip_content)) as zf:
+        file_list = [info for info in zf.infolist() if not info.is_dir()]
+        with ThreadPoolExecutor(max_workers=os.cpu_count() or 4) as executor:
+            future_to_file = {executor.submit(process_single_zipped_file, zf, file_info): file_info for file_info in file_list}
+            for future in as_completed(future_to_file):
+                try:
+                    text = future.result()
+                    if text:
+                        all_extracted_texts.append(f"--- Content from: {future_to_file[future].filename} ---\n{text}")
+                except Exception as e:
+                    print(f"Error processing file '{future_to_file[future].filename}' from ZIP: {e}")
+    
+    return "\n\n".join(all_extracted_texts)
+
+# --- Main Asynchronous Ingestion and Routing Function ---
+
+async def ingest_and_parse_document(doc_url: HttpUrl) -> str:
+    """Asynchronously downloads and parses a document using the optimal strategy."""
+    print(f"Initiating processing for URL: {doc_url}")
+    os.makedirs(LOCAL_STORAGE_DIR, exist_ok=True)
+    start_time = time.perf_counter()
+
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.get(str(doc_url), timeout=120.0, follow_redirects=True)
+            response.raise_for_status()
+            doc_bytes = response.content
+        print("Download successful.")
+
+        filename = unquote(os.path.basename(urlparse(str(doc_url)).path)) or "downloaded_file"
+        local_file_path = Path(LOCAL_STORAGE_DIR) / filename
+        file_extension = local_file_path.suffix.lower()
+
+        # Run the appropriate CPU-bound parsing function in a separate thread
+        if file_extension == '.pdf':
+            local_file_path.write_bytes(doc_bytes)
+            doc_text = await asyncio.to_thread(process_pdf_with_hybrid_parallel_sync, local_file_path)
+        elif file_extension == '.zip':
+            doc_text = await asyncio.to_thread(_process_zip_file_in_parallel, doc_bytes, Path(LOCAL_STORAGE_DIR))
+        elif file_extension in ['.png', '.jpg', '.jpeg']:
+            image = Image.open(io.BytesIO(doc_bytes))
+            doc_text = await asyncio.to_thread(process_image_element, image)
+        elif file_extension in ['.docx', '.pptx', '.html']:
+            doc_text = await asyncio.to_thread(_process_generic_file_sync, doc_bytes, filename)
+        else:
+            raise ValueError(f"Unsupported file type: {file_extension}")
+
+        elapsed_time = time.perf_counter() - start_time
+        print(f"Total processing time: {elapsed_time:.4f} seconds.")
+        if not doc_text.strip():
+            raise ValueError("Document parsing yielded no content.")
+
+        return doc_text
+
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+        raise
+
+# # Example of how to run the pipeline
+# async def main():
+#     # Example URL pointing to a PDF with tables
+#     pdf_url = HttpUrl("https://www.w3.org/WAI/WCAG21/working-examples/pdf-table-linearized/table.pdf")
+#     try:
+#         content = await ingest_and_parse_document(pdf_url)
+#         print("\n--- FINAL EXTRACTED CONTENT ---")
+#         print(content)
+#     except Exception as e:
+#         print(f"Pipeline failed: {e}")
+
+# if __name__ == "__main__":
+#     asyncio.run(main())

main.py

@@ -0,0 +1,158 @@
+# 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
+
+# Assuming 'ingestion_router.py' is in the same directory and contains the function
+from ingestion_router import ingest_and_parse_document
+from chunking_parent import create_parent_child_chunks
+from embedding import EmbeddingClient
+from retrieval_parent import Retriever
+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.3.0", # 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):
+    documents: HttpUrl
+    questions: List[str]
+
+class RunResponse(BaseModel):
+    answers: List[str]
+
+class TestRequest(BaseModel):
+    documents: HttpUrl
+
+# --- NEW: Test Endpoint for Ingestion and Parsing ---
+@app.post("/test/ingestion", response_model=Dict[str, Any], tags=["Testing"])
+async def test_ingestion_endpoint(request: TestRequest):
+    """
+    Tests the complete ingestion and parsing pipeline.
+    Downloads a document from a URL, processes it using the modular
+    parsing strategy (e.g., parallel for PDF, standard for DOCX),
+    and returns the extracted Markdown content and time taken.
+    """
+    print("--- Running Document Ingestion & Parsing Test ---")
+    start_time = time.perf_counter()
+    try:
+        # Step 1: Call the main ingestion function from your router
+        markdown_content = await ingest_and_parse_document(request.documents)
+
+        end_time = time.perf_counter()
+        duration = end_time - start_time
+        print(f"--- Ingestion and Parsing took {duration:.2f} seconds ---")
+
+        if not markdown_content:
+            raise HTTPException(
+                status_code=404,
+                detail="Document processed, but no content was extracted."
+            )
+
+        return {
+            "total_time_seconds": duration,
+            "character_count": len(markdown_content),
+            "extracted_content": markdown_content,
+        }
+    except Exception as e:
+        # Catch potential download errors, parsing errors, or unsupported file types
+        raise HTTPException(status_code=500, detail=f"An error occurred during ingestion test: {str(e)}")
+
+
+# --- 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.documents)
+        
+        # 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.documents)
+        
+        # --- 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)
+
+        # --- STAGE 4: CONCURRENT RETRIEVAL & GENERATION ---
+        print("Starting retrieval for all questions...")
+        retrieval_tasks = [
+            retriever.retrieve(q, GROQ_API_KEY)
+            for q in request.questions
+        ]
+        all_retrieved_chunks = await asyncio.gather(*retrieval_tasks)
+        print("Retrieval complete. Starting final answer generation...")
+
+        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)}"
+        )

pdf_parallel_parser.py

@@ -0,0 +1,78 @@
+# file: pdf_parallel_parser.py
+
+import fitz  # PyMuPDF
+from PIL import Image
+import io
+import os
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+
+# Import the specialized parsers from our other module
+from complex_parser import process_table_element, process_image_element
+
+def _is_bbox_contained(inner_bbox, outer_bbox):
+    """Check if inner_bbox is fully inside outer_bbox."""
+    return (inner_bbox[0] >= outer_bbox[0] and
+            inner_bbox[1] >= outer_bbox[1] and
+            inner_bbox[2] <= outer_bbox[2] and
+            inner_bbox[3] <= outer_bbox[3])
+
+def _process_page(page: fitz.Page) -> str:
+    """
+    Processes a single PDF page to extract text, tables, and images.
+    - Tables are found and processed with the complex_parser.
+    - Plain text is extracted, excluding any text already inside a processed table.
+    """
+    page_content = []
+    
+    # 1. Find and process tables first
+    table_bboxes = []
+    try:
+        tables = page.find_tables()
+        pix = page.get_pixmap(dpi=200)
+        page_image = Image.open(io.BytesIO(pix.tobytes("png")))
+        
+        print(f"Page {page.number}: Found {len(tables.tables)} potential tables.")
+        for i, table in enumerate(tables):
+            table_bboxes.append(table.bbox)
+            table_image = page_image.crop(table.bbox)
+            markdown_table = process_table_element(table_image)
+            page_content.append(markdown_table)
+    except Exception as e:
+        print(f"Could not process tables on page {page.number}: {e}")
+
+    # 2. Extract text blocks, excluding those within table bounding boxes
+    text_blocks = page.get_text("blocks")
+    for block in text_blocks:
+        block_bbox = block[:4]
+        # Check if this text block is inside any of the tables we just processed
+        is_in_table = any(_is_bbox_contained(block_bbox, table_bbox) for table_bbox in table_bboxes)
+        if not is_in_table:
+            page_content.append(block[4].strip())
+            
+    # Note: Image extraction can be added here if needed, similar to table extraction.
+
+    return "\n".join(page_content)
+
+def process_pdf_with_hybrid_parallel_sync(file_path: Path) -> str:
+    """
+    Processes a PDF file in parallel using PyMuPDF and the complex_parser.
+    """
+    print(f"Processing PDF '{file_path.name}' with parallel page-by-page strategy...")
+    all_page_texts = []
+    doc = fitz.open(file_path)
+
+    with ThreadPoolExecutor(max_workers=os.cpu_count() or 4) as executor:
+        futures = {executor.submit(_process_page, page): page.number for page in doc}
+        
+        # Collect results in page order
+        results = ["" for _ in range(len(doc))]
+        for future in as_completed(futures):
+            page_num = futures[future]
+            try:
+                results[page_num] = future.result()
+            except Exception as e:
+                print(f"Error processing page {page_num}: {e}")
+        all_page_texts = results
+
+    return f"\n\n--- Page Break ---\n\n".join(all_page_texts)

retrieval_parent.py

@@ -0,0 +1,193 @@
+# file: retrieval_parent.py
+
+import time # <-- ADD THIS IMPORT
+import asyncio
+import numpy as np
+import torch
+import json
+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 ---
+GENERATION_MODEL = "llama3-8b-8192"
+RERANKER_MODEL = 'cross-encoder/stsb-distilroberta-base'
+INITIAL_K_CANDIDATES = 20
+TOP_K_CHUNKS = 10 
+
+async def generate_hypothetical_document(query: str, groq_api_key: str) -> str:
+    """Generates a hypothetical document to answer the query (HyDE)."""
+    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 directly answers the following question. "
+        f"This passage will be used to find similar documents. "
+        f"Do not include the question or any conversational text.\n\n"
+        f"Question: {query}\n\n"
+        f"Hypothetical Passage:"
+    )
+    
+    start_time = time.perf_counter() # <-- START TIMER
+    try:
+        chat_completion = await client.chat.completions.create(
+            messages=[{"role": "user", "content": prompt}],
+            model=GENERATION_MODEL,
+            temperature=0.7,
+            max_tokens=500,
+        )
+        end_time = time.perf_counter() # <-- END TIMER
+        print(f"--- HyDE generation took {end_time - start_time:.4f} seconds ---") # <-- PRINT DURATION
+        return chat_completion.choices[0].message.content
+    except Exception as e:
+        print(f"An error occurred during HyDE generation: {e}")
+        return ""
+
+async def generate_expanded_terms(query: str, groq_api_key: str) -> List[str]:
+    """Generates semantically related search terms for a query."""
+    if not groq_api_key:
+        print("Groq API key not set. Skipping Semantic Expansion.")
+        return [query] 
+
+    print(f"Starting Semantic Expansion for query: '{query}'...")
+    client = AsyncGroq(api_key=groq_api_key)
+    prompt = (
+        f"You are a search query expansion expert. Based on the following query, generate up to 4 additional, "
+        f"semantically related search terms. The terms should be relevant for finding information in technical documents. "
+        f"Return the original query plus the new terms as a single JSON list of strings.\n\n"
+        f"Query: \"{query}\"\n\n"
+        f"JSON List:"
+    )
+
+    start_time = time.perf_counter() # <-- START TIMER
+    try:
+        chat_completion = await client.chat.completions.create(
+            messages=[{"role": "user", "content": prompt}],
+            model=GENERATION_MODEL,
+            temperature=0.4,
+            max_tokens=200,
+            response_format={"type": "json_object"},
+        )
+        end_time = time.perf_counter() # <-- END TIMER
+        print(f"--- Semantic Expansion took {end_time - start_time:.4f} seconds ---") # <-- PRINT DURATION
+        
+        result_text = chat_completion.choices[0].message.content
+        terms = json.loads(result_text)
+        
+        if isinstance(terms, dict) and 'terms' in terms:
+            return terms['terms']
+        return terms
+    except Exception as e:
+        print(f"An error occurred during Semantic Expansion: {e}")
+        return [query]
+
+
+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()
+        print(f"Retriever initialized with reranker '{RERANKER_MODEL}'.")
+
+    def index(self, child_documents: List[Document], docstore: InMemoryStore):
+        """Builds the search index from child documents and stores parent documents."""
+        self.document_chunks = child_documents
+        self.docstore = docstore
+        
+        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, expanded_terms: List[str]) -> List[Tuple[int, float]]:
+        """Performs a hybrid search using expanded terms for BM25 and a HyDE doc for dense search."""
+        if self.bm25 is None or self.chunk_embeddings is None:
+            raise ValueError("Retriever has not been indexed. Call index() first.")
+
+        print(f"Running BM25 with expanded terms: {expanded_terms}")
+        bm25_scores = self.bm25.get_scores(expanded_terms)
+
+        enhanced_query = f"{query}\n\n{hyde_doc}" if hyde_doc else 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]:
+        """Reranks candidates using a CrossEncoder model."""
+        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, groq_api_key: str) -> List[Dict]:
+        """Executes the full retrieval pipeline: expansion, HyDE, hybrid search, and reranking."""
+        print(f"\n--- Retrieving documents for query: '{query}' ---")
+
+        hyde_task = generate_hypothetical_document(query, groq_api_key)
+        expansion_task = generate_expanded_terms(query, groq_api_key)
+        hyde_doc, expanded_terms = await asyncio.gather(hyde_task, expansion_task)
+
+        initial_candidates_info = self._hybrid_search(query, hyde_doc, expanded_terms)
+        
+        retrieved_child_docs = [{
+            "content": self.document_chunks[idx].page_content,
+            "metadata": self.document_chunks[idx].metadata,
+        } for idx, score in initial_candidates_info]
+
+        reranked_child_docs = await self._rerank(query, retrieved_child_docs)
+
+        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)
+
+        retrieved_parents = self.docstore.mget(parent_ids)
+        final_parent_docs = [doc for doc in retrieved_parents if doc is not None]
+
+        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