[tab] evaluation/check_problem_difficulty.py (python)
#!/usr/bin/env python3
"""
Script to process JSONL files and print file names and entries.
"""

import json
import sys
import argparse
from google import genai
import os
from openai import OpenAI
gemini = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
openai = OpenAI(
    # This is the default and can be omitted
    api_key=os.environ.get("OPENAI_API_KEY"),
)

def get_problem_difficulty_openai(problem, response):
    """ask OpenAI GPT-5 (or latest available) to determine the difficulty of the problem"""
    prompt = f"""
    You are a helpful assistant that determines the difficulty of a math problem.
    The problem is: {problem}
    I tried asking an LLM and got the following response: {response}
    Rate the difficulty of the problem on a scale of 1 to 10, where 1 is the easiest and 10 is the hardest. Give me a one-liner why you think this is the case. 
    """
    try:
        completion = openai.chat.completions.create(
            model="gpt-5",  # Replace with actual model name if different
            messages=[
                {"role": "system", "content": "You are a grading assistant."},
                {"role": "user", "content": prompt}
            ],
            max_completion_tokens=4096
        )
        return completion.choices[0].message.content.strip()
    except Exception as e:
        print(f"OpenAI error: {e}")
        return "Error"

def get_problem_difficulty_gemini(problem, response):
    """ask gemini to determine the difficulty of the problem"""
    prompt = f"""
    You are a helpful assistant that determines the difficulty of a math problem.
    The problem is: {problem}
    I tried to solve the problem and got the following response: {response}
    Return the difficulty of the problem on a scale of 1 to 10, where 1 is the easiest and 10 is the hardest.
    """
    response = gemini.models.generate_content(model="gemini-2.5-flash", contents=[prompt])
    return response.text

def process_jsonl_file(file_path):
    """Process a single JSONL file and print its entries."""
    print(f"
{'='*60}")
    print(f"Processing file: {file_path}")
    print(f"{'='*60}
")
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, start=1):
                line = line.strip()
                if not line:  # Skip empty lines
                    continue
                
                try:
                    entry = json.loads(line)
                    print(f"Entry {line_num}:")
                    print(f"Problem:
 {entry['input'][0]['content']}")
                    difficulty = get_problem_difficulty_openai(entry['input'][0]['content'], entry['output'])
                    print(f"OpenAI GPT-5 Difficulty verdict: {difficulty}")
                    print("-" * 60)
                except json.JSONDecodeError as e:
                    print(f"Error parsing entry on line {line_num}: {e}")
                    print(f"Line content: {line}")
                    print("-" * 60)
    except FileNotFoundError:
        print(f"Error: File '{file_path}' not found.")
    except Exception as e:
        print(f"Error processing file '{file_path}': {e}")


def main():
    parser = argparse.ArgumentParser(
        description='Process JSONL files and print file names and entries.'
    )
    parser.add_argument(
        'files',
        nargs='+',
        help='One or more JSONL files to process'
    )
    
    args = parser.parse_args()
    
    for file_path in args.files:
        process_jsonl_file(file_path)


if __name__ == '__main__':
    main()