baseline.py

#!/usr/bin/env python3
"""
Baseline inference script for the Data Cleaning Environment.

Uses the OpenAI API client to run an LLM agent against the environment
for all 3 tasks (easy, medium, hard) and prints reproducible scores.

Usage:
    # Set your API key
    export OPENAI_API_KEY=sk-...

    # Run against local server (default)
    python baseline.py

    # Run against a deployed HF Space
    python baseline.py --base-url https://your-username-data-cleaning-env.hf.space

Requirements:
    pip install openai requests
"""

import argparse
import json
import os
import sys

import requests
try:
    from dotenv import load_dotenv
    load_dotenv()  
except ImportError:
    pass 

from huggingface_hub import InferenceClient

try:
    from openai import OpenAI
except ImportError:
    print("openai package not found. Install with: pip install openai")
    sys.exit(1)
try:
    from openai import OpenAI
except ImportError:
    print("openai package not found. Install with: pip install openai")
    sys.exit(1)


# ---------------------------------------------------------------------------
# Deterministic rule-based agent (no LLM needed for baseline)
# ---------------------------------------------------------------------------

RULE_POLICIES = {
    "easy":   ["impute_mean", "impute_mode", "drop_missing_rows"],
    "medium": ["remove_duplicates", "fix_type_errors", "drop_missing_rows"],
    "hard":   [
        "fill_quantity_mean",
        "drop_missing_rows",
        "remove_duplicates",
        "fix_type_errors",
        "remove_outliers",
        "normalize_text",
    ],
}

def run_rule_baseline(base_url: str) -> dict[str, float]:
    """Run deterministic rule-based baseline — no LLM required."""
    scores = {}
    for task in ["easy", "medium", "hard"]:
        # Reset
        resp = requests.post(f"{base_url}/reset", json={"task": task}, timeout=10)
        resp.raise_for_status()

        # Apply each operation in the policy
        for op in RULE_POLICIES[task]:
            resp = requests.post(
                f"{base_url}/step",
                json={"action": {"operation": op}},
                timeout=10,
            )
            resp.raise_for_status()
            data = resp.json()
            if data.get("done"):
                break

        # Grade
        resp = requests.post(f"{base_url}/grader", timeout=10)
        resp.raise_for_status()
        result = resp.json()
        scores[task] = result["score"]

    return scores

# ---------------------------------------------------------------------------
# LLM agent (uses OpenAI API)
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """You are a data cleaning agent. You will be shown a dirty dataset
as a text table and must choose ONE cleaning operation to apply per turn.

Available operations:
  impute_mean        – Fill numeric missing values with the column mean
  impute_mode        – Fill categorical missing values with the most common value
  drop_missing_rows  – Drop all rows that have any missing value
  remove_duplicates  – Remove exact duplicate rows
  fix_type_errors    – Coerce non-numeric values in numeric columns to float
  remove_outliers    – Drop rows where price <= 0 or price >= 500
  normalize_text     – Strip whitespace and title-case all string columns
  fill_quantity_mean – Fill missing quantity values with the column mean

Respond ONLY with a JSON object like:
  {"operation": "remove_duplicates"}
or with an optional column:
  {"operation": "impute_mean", "column": "age"}

No explanation. JSON only."""

def run_llm_baseline(base_url: str, api_key: str, max_steps: int = 10) -> dict[str, float]:
    """Run an LLM agent (GPT-4o-mini) against the environment."""
    client = OpenAI(api_key=api_key)
    # client = InferenceClient(api_key=api_key)
    scores = {}

    for task in ["easy", "medium", "hard"]:
        print(f"\n  [LLM] Task: {task}")
        resp = requests.post(f"{base_url}/reset", json={"task": task}, timeout=10)
        resp.raise_for_status()
        obs = resp.json()

        for step in range(max_steps):
            current_text = obs["observation"].get("current_text", "")
            metadata = obs["observation"].get("metadata", {})
            quality = metadata.get("quality_score", "?")
            valid_ops = metadata.get("valid_operations", [])

            user_msg = (
                f"Current dataset (quality score: {quality}):\n"
                f"{current_text}\n\n"
                f"Valid operations: {valid_ops}\n"
                f"Choose ONE operation to improve data quality."
            )

            completion = client.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": user_msg},
                ],
                temperature=0,
                max_tokens=64,
            )

            raw = completion.choices[0].message.content.strip()
            try:
                action = json.loads(raw)
            except json.JSONDecodeError:
                # Extract JSON from response if wrapped in markdown
                import re
                match = re.search(r"\{.*\}", raw, re.DOTALL)
                action = json.loads(match.group()) if match else {"operation": "drop_missing_rows"}

            print(f"    step {step+1}: {action}")

            resp = requests.post(
                f"{base_url}/step",
                json={"action": action},
                timeout=10,
            )
            resp.raise_for_status()
            obs = resp.json()

            if obs.get("done"):
                print(f"    Episode done at step {step+1}")
                break

        # Grade
        resp = requests.post(f"{base_url}/grader", timeout=10)
        resp.raise_for_status()
        result = resp.json()
        scores[task] = result["score"]
        print(f"  [LLM] {task} score: {scores[task]}")

    return scores

# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Data Cleaning Env baseline script")
    parser.add_argument(
        "--base-url",
        default="http://localhost:8000",
        help="Base URL of the running environment server",
    )
    parser.add_argument(
        "--mode",
        choices=["rule", "llm", "both"],
        default="rule",
        help="Baseline mode: 'rule' (no API key needed), 'llm' (needs OPENAI_API_KEY), 'both'",
    )
    args = parser.parse_args()

    base_url = args.base_url.rstrip("/")

    # Health check
    try:
        r = requests.get(f"{base_url}/health", timeout=5)
        r.raise_for_status()
        print(f"✓ Server healthy at {base_url}")
    except Exception as e:
        print(f"✗ Cannot reach server at {base_url}: {e}")
        sys.exit(1)

    # ── Rule-based baseline (always runs) ──────────────────────────────────
    if args.mode in ("rule", "both"):
        print("\n=== Rule-based Baseline ===")
        try:
            scores = run_rule_baseline(base_url)
            print("\nScores:")
            for task, score in scores.items():
                bar = "█" * int(score * 20)
                print(f"  {task:<8} {score:.4f}  {bar}")
            print(f"\n  Mean: {sum(scores.values()) / len(scores):.4f}")
        except Exception as e:
            print(f"Rule baseline failed: {e}")

    # ── LLM baseline ───────────────────────────────────────────────────────
    if args.mode in ("llm", "both"):
        api_key = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN")
        if not api_key:
            print("\nSkipping LLM baseline: OPENAI_API_KEY not set.")
        else:
            print("\n=== LLM Baseline (gpt-4o-mini) ===")
            try:
                scores = run_llm_baseline(base_url, api_key)
                print("\nScores:")
                for task, score in scores.items():
                    bar = "█" * int(score * 20)
                    print(f"  {task:<8} {score:.4f}  {bar}")
                print(f"\n  Mean: {sum(scores.values()) / len(scores):.4f}")
            except Exception as e:
                print(f"LLM baseline failed: {e}")


if __name__ == "__main__":
    main()