#!/usr/bin/env python3
"""Evaluate RAG system performance."""

from pygeai import Client
import json

def evaluate_rag(assistant_id: str, test_cases: list):
    """Evaluate RAG with test cases."""
    client = Client()
    results = []
    
    for test in test_cases:
        query = test["query"]
        expected_keywords = test.get("expected_keywords", [])
        
        response = client.rag.query(
            assistant_id=assistant_id,
            query=query,
            include_sources=True
        )
        
        # Check if expected keywords are in answer
        keywords_found = sum(1 for kw in expected_keywords if kw.lower() in response.answer.lower())
        keyword_score = keywords_found / len(expected_keywords) if expected_keywords else 1.0
        
        results.append({
            "query": query,
            "answer_length": len(response.answer),
            "num_sources": len(response.sources),
            "keyword_score": keyword_score,
            "avg_source_score": sum(s.score for s in response.sources) / len(response.sources) if response.sources else 0
        })
    
    # Calculate overall metrics
    avg_sources = sum(r["num_sources"] for r in results) / len(results)
    avg_keyword_score = sum(r["keyword_score"] for r in results) / len(results)
    
    print(f"Average sources per query: {avg_sources:.2f}")
    print(f"Average keyword score: {avg_keyword_score:.2%}")
    
    return results

if __name__ == "__main__":
    test_cases = [
        {"query": "How do I create an agent?", "expected_keywords": ["agent", "create"]},
        {"query": "What models are available?", "expected_keywords": ["model", "gpt", "claude"]}
    ]
    
    results = evaluate_rag("rag-123", test_cases)
    print(json.dumps(results, indent=2))
