scratch/time_test.py

import asyncio
import time
import os
from dotenv import load_dotenv

# Load local environment if present
load_dotenv()

from app.core.config import get_settings
from app.services.embeddings import EmbeddingService
from app.services.llm import LLMService
from app.services.vector_store import FaissVectorStore
from app.services.reranker import RerankerService
from app.services.rag_pipeline import RAGPipeline

async def run_timing_test():
    print("Initializing services...")
    settings = get_settings()
    
    # Check key
    if not settings.openai_api_key:
        print("Warning: OPENAI_API_KEY is not set. Please set it in your environment or .env file.")
        # Try to read from env var just in case
        settings.openai_api_key = os.getenv("OPENAI_API_KEY", "")
    
    t0 = time.time()
    embedding_service = EmbeddingService(settings.embedding_model)
    print(f"Embedding service loaded in {time.time() - t0:.2f}s")
    
    t0 = time.time()
    vector_store = FaissVectorStore(
        embedding_service=embedding_service,
        docs_dir=settings.docs_dir,
        index_dir=settings.index_dir,
        chunk_size_tokens=settings.chunk_size_tokens,
        chunk_overlap_tokens=settings.chunk_overlap_tokens,
    )
    vector_store.build_or_load()
    print(f"FAISS index loaded in {time.time() - t0:.2f}s")
    
    t0 = time.time()
    reranker_service = RerankerService(settings.reranker_model)
    print(f"Reranker service loaded in {time.time() - t0:.2f}s")
    
    llm_service = LLMService(
        provider=settings.llm_provider,
        openai_api_key=settings.openai_api_key,
        openai_model=settings.openai_model,
        openai_rewrite_model=settings.openai_rewrite_model,
    )
    
    pipeline = RAGPipeline(
        vector_store=vector_store,
        llm_service=llm_service,
        reranker=reranker_service,
        top_k=settings.top_k,
        max_context_chunks=settings.max_context_chunks
    )
    
    # Run test chat
    print("\n--- Running pipeline.chat Timing Test ---")
    question = "How many casual leaves do I have?"
    history = []
    
    # Time rewrite query
    t_start = time.time()
    rewritten = await llm_service.rewrite_query(question, history)
    t_rewrite = time.time() - t_start
    print(f"1. Query Rewritten: '{question}' -> '{rewritten}' in {t_rewrite:.2f}s")
    
    # Time vector search
    t_start = time.time()
    queries = [rewritten, question]
    all_retrieved = vector_store.multi_search(queries, top_k=settings.top_k)
    t_search = time.time() - t_start
    print(f"2. Multi-search (FAISS) completed in {t_search:.2f}s (found {len(all_retrieved)} chunks)")
    
    # Time reranking
    t_start = time.time()
    reranked = reranker_service.rerank(rewritten, all_retrieved, top_n=settings.max_context_chunks)
    t_rerank = time.time() - t_start
    print(f"3. Reranking completed in {t_rerank:.2f}s")
    
    # Time answer generation
    t_start = time.time()
    answer = await llm_service.answer(question, reranked, history)
    t_answer = time.time() - t_start
    print(f"4. LLM Answer generated in {t_answer:.2f}s")
    print(f"Response: '{answer[:60]}...'")
    
    print(f"\nTotal Pipeline time: {t_rewrite + t_search + t_rerank + t_answer:.2f}s")
    
    await llm_service.close()

if __name__ == "__main__":
    asyncio.run(run_timing_test())