app.py

import os
import random
import time
import logging
import asyncio
import threading
import uuid
import json

from datetime import datetime, timezone
from collections import defaultdict, Counter

from huggingface_hub import InferenceClient, HfApi, CommitOperationAdd, hf_hub_download, list_repo_files
from datasets import Dataset, load_dataset, concatenate_datasets
import gradio as gr
import pandas as pd

from model import ModelWrapper, get_model


HF_TOKEN = os.environ.get("HF_TOKEN")
HF_API = HfApi(token=HF_TOKEN)
REPO_ID = "aracape/LA-Arena"

MAX_REQUESTS_PER_HOUR = 20
MAX_TOKENS = 512
TEMPERATURE = 0.5
DEFAULT_SYSTEM = "You are a helpful learning assistant who supports students and helps them learn."
EXTENDED_SYSTEM = """
## Core Principles

**Guide, Don't Tell**: Your primary role is to facilitate learning through thoughtful questioning and scaffolded hints. Avoid giving direct answers unless absolutely necessary for the student's learning progression.

**Socratic Method with Flexibility**: Use questions to guide students toward insights, but remain adaptive. If a student is completely stuck on a prerequisite concept or needs a direct factual clarification to move forward, provide it concisely, then return to guided questioning.

**Build Understanding Incrementally**: Break complex problems into manageable steps. Start with what the student knows, then build toward the solution progressively.

## Your Approach

### 1. **Start with Diagnosis**
- Ask questions to understand what the student already knows
- Identify specific points of confusion
- Assess their current level of understanding

Example questions:
- "What have you tried so far?"
- "Which part of the problem feels most challenging?"
- "Can you explain what you understand about [concept] in your own words?"

### 2. **Provide Scaffolded Hints**
When offering hints, follow this progression:
- **First hint**: Point to a relevant concept or approach without revealing the solution
- **Second hint**: Break down the problem into smaller sub-problems
- **Third hint**: Provide a similar worked example or analogy
- **Only if needed**: Give more direct guidance while still leaving the final step to the student

### 3. **Ask Thoughtful, Purposeful Questions**
Your questions should:
- Direct attention to relevant concepts or relationships
- Prompt specific analytical thinking
- Help students recognize patterns or connections
- Encourage self-correction

Avoid vague questions like "Does that make sense?" Instead use:
- "What happens if you apply [concept] to this part?"
- "How does this connect to [related idea] we discussed?"
- "What do you notice about [specific element]?"

### 4. **Ensure Accuracy**
- Provide factually correct information in all hints and guidance
- If you're pointing toward a concept, ensure your description is precise
- Verify that your hints lead toward the correct solution path

### 5. **Adapt Your Approach**
- **For conceptual questions**: Use Socratic dialogue extensively
- **For factual clarifications**: Provide brief, direct answers then return to guided inquiry
- **For multi-step problems**: Break into phases with checkpoints
- **For completely stuck students**: Offer a more direct hint to unstick them, then step back

## Response Structure

1. **Acknowledge** what the student has shared
2. **Ask diagnostic questions** if needed to understand their thinking
3. **Provide a scaffolded hint or question** that moves them forward
4. **Encourage next steps** by indicating what they should think about or try next

## What to Avoid

- Giving complete solutions or final answers
- Asking too many questions at once (overwhelming)
- Being vague or unhelpfully abstract
- Using overly Socratic approaches when direct clarification is needed
- Providing hints that are too advanced for the student's current level

## Example Interaction Pattern

**Poor**: "The answer is X because of Y and Z."

**Good**: "I see you're working on [problem]. You mentioned [student's thought]. That's a good starting point. What do you think would happen if you [relevant prompt]? Consider how [related concept] might apply here."

Remember: Your success is measured by the student's learning journey, not by how quickly they reach the answer. Help them build confidence and genuine understanding through guided discovery.
"""

# Ensure consistent random responses
random.seed(time.time_ns())

logger = logging.getLogger("LA Arena")
logger.setLevel(logging.DEBUG)
if not logger.handlers:
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter('%(name)s - %(levelname)s - %(message)s'))
    logger.addHandler(handler)

_rl_lock = threading.Lock()
request_tracker = defaultdict(list)

MODEL_NAMES = {
    "baseline": "Llama 3.2 1B (Baseline)",
    "fine_tuned": "Llama 3.2 1B (Fine-tuned)",
    "prompted": "Llama 3.2 1B (Prompted)"
}

def exceeded_rate_limit(request: gr.Request):
    now = time.time()
    hour_ago = now - 3600
    client_id = getattr(request, "client", None)
    ip = getattr(client_id, "host", None) or request.headers.get("x-forwarded-for", "unknown")

    with _rl_lock:
        request_tracker[ip] = [t for t in request_tracker[ip] if t > hour_ago]
        if len(request_tracker[ip]) >= MAX_REQUESTS_PER_HOUR:
            return True
        request_tracker[ip].append(now)
 
    return False

def get_messages(prompt, variant, history):
    system_message = DEFAULT_SYSTEM
    if variant == "prompted":
        system_message += "\n" + EXTENDED_SYSTEM
    messages = [{"role": "system", "content": system_message}]
    messages.extend(history)
    messages.append({"role": "user", "content": prompt})
    return messages

def respond_single_model(
    message,
    history: list[dict[str, str]],
    model_choice: str,
    request: gr.Request
):
    """Chat with a single model"""
    if exceeded_rate_limit(request):
        yield "Sorry you exceeded the rate limit for this hour"

    model = get_model(model_choice)
     
    messages = get_messages(message, model_choice, history)
    logger.debug(f"{model_choice}: {messages}")
    
    yield  model.generate(messages, MAX_TOKENS, TEMPERATURE)


async def respond_two_models(prompt, request: gr.Request):
    if exceeded_rate_limit(request):
        msg = "Sorry you exceeded the rate limit for this hour"
        return msg, msg, "rate_limit", "rate_limit", ""

    model_keys = random.sample(["baseline", "fine_tuned", "prompted"], 2)
    model_a_key, model_b_key = model_keys[0], model_keys[1]
    model_a = get_model(model_a_key)
    model_b = get_model(model_b_key) 

    def run_model(model: ModelWrapper, variant):
        # TEMP: Mock responses for testing data saving
        # return f"Here is a response: {random.randint(0, 10)}"

        messages = get_messages(prompt, variant, [])
        logger.debug(f"{variant}: {messages}")
        return model.generate(messages, MAX_TOKENS, TEMPERATURE)
    
    response_a, response_b = await asyncio.gather(
        asyncio.to_thread(run_model, model_a, model_a_key),
        asyncio.to_thread(run_model, model_b, model_b_key)
    )

    # logger.debug(f"{model_a_key}: {response_a}")
    # logger.debug(f"{model_b_key}: {response_b}")
    
    return response_a, response_b, model_a_key, model_b_key, ""


def save_vote(prompt, response_a, response_b, model_a, model_b, choice):
    if not response_a or not response_b:
        logger.warning("No responses to vote on yet")
        return

    record = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "prompt": prompt,
        "response_a": response_a,
        "response_b": response_b,
        "model_a": model_a,
        "model_b": model_b,
        "choice": choice,
        "id": str(uuid.uuid4()),
    }

    # Write to a temp file
    tmp_path = f"/tmp/{record['id']}.json"
    with open(tmp_path, "w", encoding="utf-8") as f:
        json.dump(record, f, ensure_ascii=False)
    
    repo_path = f"votes/{record['id']}.json"

    reveal_msg = f"\n\n**Model A:** {MODEL_NAMES.get(model_a, model_a)} | **Model B:** {MODEL_NAMES.get(model_b, model_b)}"

    try:
        HF_API.create_commit(
            repo_id=REPO_ID,
            repo_type="space",
            operations=[CommitOperationAdd(path_in_repo=repo_path, path_or_fileobj=tmp_path)],
            commit_message=f"Add vote {record['id']}"
        )
        return "### Vote saved! 🗳️" + reveal_msg
    except Exception as e:
        logging.exception("Vote save failed")
        return f"### Vote save failed: {e}" + reveal_msg
        

def compute_win_rates(repo_id=REPO_ID):
    files = [p for p in list_repo_files(repo_id=repo_id, repo_type="space") if p.startswith("votes/") and p.endswith(".json")]
    wins, total = Counter(), 0
    for p in files:
        local = hf_hub_download(repo_id=repo_id, repo_type="space", filename=p)
        with open(local, "r", encoding="utf-8") as f:
            row = json.load(f)
        if row.get("choice") == "A":
            wins[row["model_a"]] += 1; total += 1
        elif row.get("choice") == "B":
            wins[row["model_b"]] += 1; total += 1
    return {k: (v / total if total else 0.0) for k, v in wins.items()}
    
def get_leaderboard_df(repo_id=REPO_ID):
    rates = compute_win_rates(repo_id)
    if not rates:
        return pd.DataFrame(columns=["Model", "Win Rate"]), "No votes yet — submit a prompt and cast the first vote!"
    df = pd.DataFrame(
        [(model, f"{rate*100:.1f}%") for model, rate in rates.items()],
        columns=["Model", "Win Rate"]
    ).sort_values("Win Rate", key=lambda s: s.str.rstrip("%").astype(float), ascending=False)
    return df, f"Updated leaderboard ({len(df)} models)"


def create_leaderboard_interface():
    gr.Markdown("Win rates computed from arena matchups and voting data")
    df, md = get_leaderboard_df(REPO_ID)
    status_md = gr.Markdown(md)
    table = gr.Dataframe(
        value=df,
        headers=["Model", "Win Rate"],
        datatype=["str", "str"],
        interactive=False,
        wrap=True,
        row_count=(0, "dynamic"),
        col_count=(2, "fixed")
    )
    refresh_btn = gr.Button("Refresh")

    # Wire the refresh button
    refresh_btn.click(
        fn=lambda: get_leaderboard_df(REPO_ID),
        inputs=None,
        outputs=[table, status_md],
    )


# Mode 1: Simple Chat with a Model
def create_chat_interface():
    """Single model chat interface"""
    with gr.Blocks() as chat_block:
        gr.Markdown("### Chat with a Model")
        gr.Markdown("Select a model and start testing!")
        
        model_dropdown = gr.Dropdown(
            choices=["baseline", "fine_tuned", "prompted"],
            value="fine_tuned",
            label="Choose Model",
            info="Select which model to chat with"
        )
        
        chatbot = gr.ChatInterface(
            fn=respond_single_model,
            additional_inputs=[model_dropdown],
            type="messages",
            title="",
            description="",
        )
    
    return chat_block

# Mode 2: Arena Comparison
def create_arena_interface():
    """Blind A/B testing interface - we'll implement this next"""
    with gr.Column() as arena:
        gr.Markdown("## Head to Head Battle")
        gr.Markdown("*What model do you think would help you learn the most?*")
        
        prompt_box = gr.Textbox(
            label="Enter your prompt",
            placeholder="Type a question or prompt here...",
            lines=3
        )
        submit_btn = gr.Button("Generate Responses", variant="primary")
        
        with gr.Row():
            with gr.Column():
                response_a = gr.Textbox(label="👤 Model A", lines=4, interactive=False)
            with gr.Column():
                response_b = gr.Textbox(label="👤 Model B", lines=4, interactive=False)
                
        # Hidden states
        model_a = gr.State()
        model_b = gr.State()
        
        gr.Markdown("### Which response is better?")
        with gr.Row():
            vote_a = gr.Button("👈 A is Better")
            vote_tie = gr.Button("🤝 Tie")
            vote_b = gr.Button("👉 B is Better")
        
        result_display = gr.Markdown("")
       
        submit_btn.click(
            respond_two_models, 
            inputs=[prompt_box], 
            outputs=[response_a, response_b, model_a, model_b, result_display],
            concurrency_limit=8,
        )
        for btn, choice in [(vote_a, "A"), (vote_tie, "Tie"), (vote_b, "B")]:
            btn.click(
                save_vote,
                inputs=[prompt_box, response_a, response_b, model_a, model_b, gr.State(choice)],
                outputs=[result_display],
                concurrency_id="voting_queue",
            )
    return arena

# Main app with tabs
with gr.Blocks(title="Model Evaluation Platform") as demo:
    gr.Markdown("# Learning Assistant Arena")
    gr.Markdown("Put different LLMs to the test")
    
    with gr.Tabs():
        with gr.Tab("🥊 Arena Mode"):
            create_arena_interface()

        with gr.Tab("💬 Chat Mode"):
            create_chat_interface()
         
        with gr.Tab("📊 Leaderboard"):
            create_leaderboard_interface()

if __name__ == "__main__":
    demo.launch()