test.py

"""
test.py

End-to-end testing script for the refactored RAG pipeline.
"""

import logging
from pathlib import Path

from config import COLLECTION_NAME, DB_PATH
from embedding.embedder import Embedder
from embedding.model import EmbeddingModel
from loader import DocumentLoader
from builder import PromptBuilder
from pipeline.indexing import IndexingPipeline
from pipeline.querying import QueryPipeline
from pipeline.rag_pipeline import RAGPipeline
from retrival.context_builder import ContextBuilder
from retrival.router import IntentRouter
from retrival.search_service import HeuristicReranker, SearchService
from src.chunker import Chunker
from vectordb.croma_store import ChromaVectorStore
from vectordb.retriever import Retriever

# Mock LLM for local testing without external API requirements
class MockLLM:
    def generate(self, prompt: str) -> str:
        return f"[Mock LLM Response generated successfully for prompt length: {len(prompt)} chars]"

def main() -> None:
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger("RAG_Test")

    logger.info("Initializing components...")
    loader = DocumentLoader(data_dir="data")
    chunker = Chunker()
    embedder = Embedder()
    embedding_model = EmbeddingModel()
    vector_store = ChromaVectorStore(db_path=DB_PATH, collection_name=COLLECTION_NAME)

    # 1. Test Indexing Pipeline
    indexing_pipeline = IndexingPipeline(
        loader=loader,
        chunker=chunker,
        embedder=embedder,
        vector_store=vector_store,
    )
    
    indexed_count = indexing_pipeline.index()
    logger.info("Total indexed chunks: %d", indexed_count)

    # 2. Test Retrieval & Query Pipeline
    retriever = Retriever(vector_store=vector_store, embedding_model=embedding_model)
    context_builder = ContextBuilder(max_chunks=5)
    reranker = HeuristicReranker()
    
    search_service = SearchService(
        retriever=retriever,
        reranker=reranker,
        context_builder=context_builder,
    )
    
    prompt_builder = PromptBuilder()
    mock_llm = MockLLM()

    querying_pipeline = QueryPipeline(
        search_service=search_service,
        prompt_builder=prompt_builder,
        llm=mock_llm,
    )

    rag = RAGPipeline(
        indexing=indexing_pipeline,
        querying=querying_pipeline,
    )

    # 3. Execution Test
    test_query = "What coping mechanisms help with anxiety?"
    logger.info("Running query: '%s'", test_query)
    
    response = rag.answer(question=test_query)
    print("\n" + "=" * 50)
    print("PIPELINE OUTPUT:")
    print("=" * 50)
    print(response)
    print("=" * 50)

if __name__ == "__main__":
    main()