benchmark.py

import os
import json
import time
import pandas as pd
from typing import List, Dict
from datasets import Dataset
from ragas import evaluate
from ragas.metrics import faithfulness, answer_relevancy, context_precision

# Import logic from server (simplified)
import server

def run_benchmark(questions: List[str], ground_truth: List[str] = None):
    results_naive = []
    results_advanced = []

    print(f"🚀 Starting benchmark with {len(questions)} questions...")

    for q in questions:
        # Naive RAG
        start_time = time.time()
        retrieved_naive = server.retrieve(q, top_k=4)
        context_naive = "\n".join([c["text"] for c in retrieved_naive])
        
        # Simple query for naive (ignoring LLM call for speed if just measuring retrieval, 
        # but RAGAS needs answer, so we might need to mock or call)
        # For this benchmark, we'll call the real server logic
        resp_naive = server.query(server.QueryRequest(query=q, advanced=False))
        results_naive.append({
            "question": q,
            "answer": resp_naive["answer"],
            "contexts": [c["text"] for c in resp_naive["sources"]],
            "time": time.time() - start_time
        })

        # Advanced RAG
        start_time = time.time()
        resp_adv = server.query(server.QueryRequest(query=q, advanced=True))
        results_advanced.append({
            "question": q,
            "answer": resp_adv["answer"],
            "contexts": [c["text"] for c in resp_adv["sources"]],
            "time": time.time() - start_time
        })

    # Convert to Ragas format
    ds_naive = Dataset.from_list(results_naive)
    ds_adv = Dataset.from_list(results_advanced)

    print("📊 Evaluating Naive RAG...")
    score_naive = evaluate(ds_naive, metrics=[faithfulness, answer_relevancy, context_precision])
    
    print("📊 Evaluating Advanced RAG...")
    score_adv = evaluate(ds_adv, metrics=[faithfulness, answer_relevancy, context_precision])

    return score_naive, score_adv

if __name__ == "__main__":
    # Example questions
    test_questions = [
        "What is the main topic of the document?",
        "Can you summarize the key findings?",
        "Explain the methodology used in the text."
    ]
    
    # Ensure some documents are loaded in server.DOCUMENTS for testing
    # In a real scenario, you'd upload a specific test file first
    if not server.DOCUMENTS:
        print("⚠️ No documents found. Please upload a document to the server first.")
    else:
        sn, sa = run_benchmark(test_questions)
        
        df = pd.DataFrame([
            {"Method": "Naive RAG", **sn},
            {"Method": "Advanced RAG", **sa}
        ])
        print("\n🏆 Benchmark Results:")
        print(df.to_markdown())
        
        with open("benchmark_results.json", "w") as f:
            json.dump({"naive": sn, "advanced": sa}, f, indent=2)