Retrieval evaluation using LLM-as-a-judge via Pydantic AI ↗

%pip install -U cohere pydantic-ai

import requests
import cohere
import pandas as pd
from pydantic_ai import Agent
from pydantic_ai.models import KnownModelName
from collections import Counter

import os
co = cohere.ClientV2(os.getenv("COHERE_API_KEY"))

import nest_asyncio
nest_asyncio.apply()

import requests

def search_wikipedia(query, limit=10):
    url = "https://en.wikipedia.org/w/api.php"
    params = {
        'action': 'query',
        'list': 'search',
        'srsearch': query,
        'format': 'json',
        'srlimit': limit
    }

    response = requests.get(url, params=params)
    data = response.json()
    
    # Format the results
    results = []
    for item in data['query']['search']:
        results.append({
            "title": item["title"],
            "snippet": item["snippet"].replace("<span class=\"searchmatch\">", "").replace("</span>", ""),
        })
            
    return results

geography_questions = [
    "What is the capital of France?",
    "What is the longest river in the world?",
    "What is the largest desert in the world?",
    "What is the highest mountain peak on Earth?",
    "What are the major tectonic plates?",
    "What is the Ring of Fire?",
    "What is the largest ocean on Earth?",
    "What are the Seven Wonders of the Natural World?",
    "What causes the Northern Lights?",
    "What is the Great Barrier Reef?"
]

# Run search_wikipedia for each question
results = []

for question in geography_questions:
    question_results = search_wikipedia(question, limit=10)
    
    # Format the results as requested
    formatted_results = []
    for item in question_results:
        formatted_result = f"{item['title']}\n{item['snippet']}"
        formatted_results.append(formatted_result)
    
    # Add to the results list
    results.append({
        "question": question,
        "search_results": formatted_results
    })

# Rerank the search results for each question
top_n = 3
results_reranked_top_n = []

for item in results:
    question = item["question"]
    documents = item["search_results"]
    
    # Rerank the documents using Cohere
    reranked = co.rerank(
        model="rerank-v3.5",
        query=question,
        documents=documents,
        top_n=top_n  # Get top 3 results
    )
    
    # Format the reranked results
    top_results = []
    for result in reranked.results:
        top_results.append(documents[result.index])
    
    # Add to the reranked results list
    results_reranked_top_n.append({
        "question": question,
        "search_results": top_results
    })

# Print a sample of the reranked results
print(f"Original question: {results_reranked_top_n[0]['question']}")
print(f"Top 3 reranked results:")
for i, result in enumerate(results_reranked_top_n[0]['search_results']):
    print(f"\n{i+1}. {result}")

results_top_n = []

for item in results:
    results_top_n.append({
        "question": item["question"],
        "search_results": item["search_results"][:top_n]
    })
    
# Print a sample of the top_n results (without reranking)
print(f"Original question: {results_top_n[0]['question']}")
print(f"Top {top_n} results (without reranking):")
for i, result in enumerate(results_top_n[0]['search_results']):
    print(f"\n{i+1}. {result}")

# System prompt for the AI evaluator
SYSTEM_PROMPT = """
You are an AI search evaluator. You will compare search results from two engines and
determine which set provides more relevant and diverse information. You will only
answer with the verdict rather than explaining your reasoning; simply say "Engine A" or
"Engine B".
"""

# Prompt template for each evaluation
PROMPT_TEMPLATE = """
For the following question, which search engine provides more relevant results?

## Question:
{query}

## Engine A:
{engine_a_results}

## Engine B:
{engine_b_results}
"""

def format_results(results):
    """Format search results in a readable way"""
    formatted = []
    for i, result in enumerate(results):
        formatted.append(f"Result {i+1}: {result[:200]}...")
    return "\n\n".join(formatted)

def judge_query(query, engine_a_results, engine_b_results, model_name):
    """Use a single model to judge which engine has better results"""
    agent = Agent(model_name, system_prompt=SYSTEM_PROMPT)
    
    # Format the results
    engine_a_formatted = format_results(engine_a_results)
    engine_b_formatted = format_results(engine_b_results)
    
    # Create the prompt
    prompt = PROMPT_TEMPLATE.format(
        query=query,
        engine_a_results=engine_a_formatted,
        engine_b_results=engine_b_formatted
    )
    
    # Get the model's judgment
    response = agent.run_sync(prompt)
    return response.data

def evaluate_search_results(reranked_results, regular_results, models):
    """
    Evaluate both sets of search results using multiple models.
    
    Args:
        reranked_results: List of dictionaries with 'question' and 'search_results'
        regular_results: List of dictionaries with 'question' and 'search_results'
        models: List of model names to use as judges
    
    Returns:
        DataFrame with evaluation results
    """
    # Prepare data structure for results
    evaluation_results = []
    
    # Evaluate each query
    for i in range(len(reranked_results)):
        query = reranked_results[i]['question']
        engine_a_results = reranked_results[i]['search_results']  # Reranked results
        engine_b_results = regular_results[i]['search_results']   # Regular results
        
        # Get judgments from each model
        judgments = []
        for model in models:
            judgment = judge_query(query, engine_a_results, engine_b_results, model)
            judgments.append(judgment)
        
        # Determine winner by majority vote
        votes = Counter(judgments)
        if votes["Engine A"] > votes["Engine B"]:
            winner = "Engine A"
        elif votes["Engine B"] > votes["Engine A"]:
            winner = "Engine B"
        else:
            winner = "Tie"
        
        # Add results for this query
        row = [query] + judgments + [winner]
        evaluation_results.append(row)
    
    # Create DataFrame
    column_names = ["question"] + [f"judge_{i+1} ({model})" for i, model in enumerate(models)] + ["winner"]
    df = pd.DataFrame(evaluation_results, columns=column_names)
    
    return df

# Define the search engines
engine_a = results_reranked_top_n
engine_b = results_top_n

# Define the models to use as judges
models = [
    "cohere:command-a-03-2025",
    "cohere:command-r-plus-08-2024",
    "cohere:command-r-08-2024",
    "cohere:c4ai-aya-expanse-32b",
    "cohere:c4ai-aya-expanse-8b",
]

# Get evaluation results
results_df = evaluate_search_results(engine_a, engine_b, models)

# Calculate overall statistics
winner_counts = Counter(results_df["winner"])
total_queries = len(results_df)

# Display summary of results
print("\nPercentage of questions won by each engine:")
for engine, count in winner_counts.items():
    percentage = (count / total_queries) * 100
    print(f"{engine}: {percentage:.2f}% ({count}/{total_queries})")
    
# Display dataframe
results_df.head()

# Save to CSV
results_csv = results_df.to_csv("search_results_evaluation.csv", index=False)