[tab] gdr/refining.py (python)
import os
from google import genai
from dotenv import load_dotenv
import judgement_prompts
import refinement_prompts

load_dotenv()

key = os.getenv("GEMINI_API_KEY_1")
client = genai.Client(api_key=key)
model = client.models

JUDGEMENT_PROMPT = judgement_prompts.judgement_prompt_v1

REFINEMENT_PROMPT = refinement_prompts.generative_data_refinement_v3   


def judge_problem_difficulty(problem: str) -> str:
    """Judges the difficulty of a problem."""
    print("--- Step 1: Judging problem difficulty... ---")
    try:
        prompt = JUDGEMENT_PROMPT.format(problem=problem)
        response = model.generate_content(
            model="gemini-2.5-flash", contents=prompt
        )
        judgement = response.text.strip()
        if judgement in ["Easy", "Medium", "Hard"]:
            return judgement
        else:
            print(f"Warning: Unexpected difficulty '{judgement}'. Defaulting to 'Medium'.")
            return "Medium"
    except Exception as e:
        print(f"Error during judgment: {e}")
        return "Medium"

def refine_thought_trace_advanced(problem: str, difficulty: str, original_thought_trace: str) -> str:
    """Refines a thought trace using the advanced rubric-based prompt."""
    print(f"--- Step 2: Refining trace for '{difficulty}' difficulty using advanced prompt... ---")
    try:
        prompt = REFINEMENT_PROMPT.format(
            problem=problem,
            difficulty=difficulty,
            original_thought_trace=original_thought_trace
        )
        response = model.generate_content(
            model="gemini-2.5-flash", contents=prompt
        )
        return response.text.strip()
    except Exception as e:
        print(f"Error during refinement: {e}")
        return "Error: Could not refine trace."


def print_header(title: str, width: int = 80):
    """Print a nicely formatted header."""
    print("
" + "=" * width)
    print(f"  {title}")
    print("=" * width)

def print_section(title: str, content: str, width: int = 80):
    """Print a section with title and content."""
    print(f"
{title}")
    print("─" * width)
    print(content.strip())

def print_comparison(original: str, refined: str, width: int = 100):
    """Print original and refined traces in a nice comparison format."""
    print("
" + "TRACE COMPARISON".center(width))
    print("=" * width)
    
    # Split into left and right columns
    col_width = (width - 5) // 2  # 5 chars for " │ " separator
    
    print(f"{'ORIGINAL TRACE':<{col_width}} │ {'REFINED TRACE'}")
    print("─" * col_width + "─┼─" + "─" * (width - col_width - 3))
    
    # Clean and split traces
    orig_lines = [line.strip() for line in original.strip().split('
') if line.strip()]
    ref_lines = [line.strip() for line in refined.strip().split('
') if line.strip()]
    
    # Print side by side
    max_lines = max(len(orig_lines), len(ref_lines))
    for i in range(max_lines):
        orig_line = orig_lines[i] if i < len(orig_lines) else ""
        ref_line = ref_lines[i] if i < len(ref_lines) else ""
        
        # Truncate if too long and add ellipsis
        if len(orig_line) > col_width - 2:
            orig_line = orig_line[:col_width-5] + "..."
        if len(ref_line) > col_width - 2:
            ref_line = ref_line[:col_width-5] + "..."
            
        print(f"{orig_line:<{col_width}} │ {ref_line}")
    
    print("=" * width)

if __name__ == "__main__":
    # Problem setup
    hard_problem = "The sum of the ages of a father and son is 45 years. Five years ago, the product of their ages was 4 times the father's age at that time. What are the current ages of the father and son?"
    hard_trace = """
1. Let F be father's age and S be son's age.
2. F + S = 45.
3. Five years ago, father was F-5 and son was S-5.
4. (F-5)(S-5) = 4(F-5).
5. S-5 = 4, so S=9.
6. Sub S into first eq: F + 9 = 45.
7. F = 36.
8. So, father is 36, son is 9.
    """
    
    # Main execution with improved formatting
    print_header("ADAPTIVE THOUGHT TRACE REFINEMENT")
    
    print_section("Problem Statement", hard_problem)
    
    print("
Analyzing problem difficulty...")
    judged_difficulty = judge_problem_difficulty(hard_problem)
    
    print(f"
Assessment Complete: {judged_difficulty} Difficulty")
    
    print(f"
Refining trace for {judged_difficulty.lower()} difficulty...")
    refined_trace = refine_thought_trace_advanced(hard_problem, judged_difficulty, hard_trace)
    
    # Display results
    print_comparison(hard_trace, refined_trace)
    
    print(f"
Refinement Complete!")
    print(f"Difficulty Level: {judged_difficulty}")
    print(f"Optimization: Adapted for {judged_difficulty.lower()}-level reasoning")