inference.py

#!/usr/bin/env python3
"""
DataQA Inference Script — Two-Phase Agent
------------------------------------------
LLM agent that plays the DataQA environment in two phases:
  Phase 1: Identify all data quality issues
  Phase 2: Propose fixes for identified issues

Uses the OpenAI client to interact with any OpenAI-compatible LLM API.

Required environment variables:
    API_BASE_URL  - LLM API endpoint (e.g., https://router.huggingface.co/v1)
    MODEL_NAME    - Model identifier (e.g., Qwen/Qwen2.5-72B-Instruct)
    HF_TOKEN      - HuggingFace token / API key

STDOUT FORMAT (mandatory for evaluation):
    [START] task=<task_name> env=<benchmark> model=<model_name>
    [STEP]  step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
    [END]   success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>
"""

from __future__ import annotations

import os
import re
import sys
import time
from typing import List, Optional

import requests
from openai import OpenAI

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
ENV_URL = os.getenv("ENV_URL", "http://localhost:8000")

BENCHMARK = "dataqa_env"
TASKS = ["easy", "medium", "hard", "alignment", "coding", "toolcalling", "moderation"]
MAX_STEPS_PER_TASK = 3


# ---------------------------------------------------------------------------
# Logging helpers (structured stdout — exact format required by evaluation)
# ---------------------------------------------------------------------------

def log_start(task: str, env: str, model: str) -> None:
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
    error_val = error if error else "null"
    done_val = str(done).lower()
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
        flush=True,
    )


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
        flush=True,
    )


# ---------------------------------------------------------------------------
# Environment HTTP client
# ---------------------------------------------------------------------------

class EnvHTTPClient:
    """Minimal HTTP client for the DataQA environment."""

    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip("/")
        self.session = requests.Session()

    def health(self) -> bool:
        try:
            r = self.session.get(f"{self.base_url}/health", timeout=10)
            return r.status_code == 200
        except Exception:
            return False

    def reset(self, task_id: str = "easy") -> dict:
        r = self.session.post(
            f"{self.base_url}/reset",
            json={"task_id": task_id},
            timeout=30,
        )
        r.raise_for_status()
        return r.json()

    def step(self, issues: list[str], fixes: list[str], task_id: str = "easy") -> dict:
        r = self.session.post(
            f"{self.base_url}/step",
            json={"action": {"issues": issues, "fixes": fixes, "task_id": task_id}},
            timeout=30,
        )
        r.raise_for_status()
        return r.json()


# ---------------------------------------------------------------------------
# LLM Prompts
# ---------------------------------------------------------------------------

IDENTIFY_SYSTEM_PROMPT = """You are a data quality analyst. Your job is to inspect datasets and identify data quality issues.

You will be given:
1. A dataset in CSV format
2. A schema describing expected column types and constraints
3. Validation rules that the data should satisfy

You must identify ALL data quality issues and report each one in EXACTLY this format:
row:<row_number>,col:<column_name>,issue:<issue_type>

Supported issue types:
- missing_value (null, empty, or whitespace-only)
- wrong_type (value doesn't match expected type)
- duplicate_row (exact duplicate or duplicate key)
- out_of_range (value outside valid range)
- format_violation (wrong format, invalid enum value)
- inconsistent_value (computed field doesn't match, logical inconsistency)
- statistical_outlier (value is unreasonable given context)
- referential_integrity (foreign key violation)

CRITICAL INSTRUCTIONS FOR ROW NUMBERING:
- Row numbers refer to the ROW POSITION in the CSV data, NOT the value of any ID column
- Row 1 = the FIRST data row after the header
- Row 2 = the SECOND data row after the header
- DO NOT use the employee_id, order_id, or experiment_id as the row number
- Column names must match exactly (use the CSV header names, lowercase)
- Check EVERY row and EVERY column systematically
- Consider cross-column consistency (e.g., total = quantity * price)
- Look for subtle issues like whitespace-only values, near-duplicates
- Report ALL issues you find, even if uncertain

Respond with ONLY the list of issues, one per line. No other text.
Example: row:3,col:salary,issue:missing_value"""


FIX_SYSTEM_PROMPT = """You are a data repair specialist. You have already identified data quality issues in a dataset. Now you must propose the correct values to fix each issue.

For each issue you identified, propose a fix in EXACTLY this format:
row:<row_number>,col:<column_name>,fix:<corrected_value>

Guidelines for proposing fixes:
- For missing_value: infer the correct value from context, schema, and other rows
- For wrong_type: convert to the correct type (e.g., "seventy-five thousand" → "75000")
- For out_of_range: propose a value within the valid range that makes sense in context
- For format_violation: correct the format (e.g., "26/01/2024" → "2024-01-26")
- For inconsistent_value: compute the correct value from related fields
- For duplicate_row: propose a corrected unique key or indicate removal
- For statistical_outlier: propose a reasonable value given the model/context

Use the schema, validation rules, and surrounding data to determine the correct fix.
Respond with ONLY the list of fixes, one per line. No other text.
Example: row:3,col:salary,fix:75000"""


def build_user_prompt(observation: dict, include_fixes: bool = False) -> str:
    obs = observation if isinstance(observation, dict) else observation
    parts = []

    if obs.get("task_description"):
        parts.append(f"TASK: {obs['task_description']}")

    parts.append(f"SCHEMA:\n{obs.get('schema_description', '')}")
    parts.append(f"VALIDATION RULES:\n{obs.get('validation_rules', '')}")
    parts.append(f"DATASET:\n{obs.get('dataset_csv', '')}")

    hint = obs.get("num_issues_hint", 0)
    if hint:
        parts.append(f"HINT: There are exactly {hint} issues to find.")

    feedback = obs.get("feedback", "")
    if feedback and "reset" not in feedback.lower():
        parts.append(f"FEEDBACK FROM PREVIOUS ATTEMPT:\n{feedback}")

    if include_fixes:
        parts.append(
            "Now propose fixes for ALL issues. "
            "Use format: row:<N>,col:<name>,fix:<corrected_value>"
        )

    return "\n\n".join(parts)


def parse_llm_response(response: str) -> list[str]:
    """Extract issue lines from LLM response."""
    issues = []
    for line in response.strip().split("\n"):
        line = line.strip()
        if not line:
            continue
        line = re.sub(r"^\s*[\d]+[.\)]\s*", "", line)
        line = re.sub(r"^\s*[-*]\s*", "", line)
        line = line.strip()
        if "row" in line.lower() and "col" in line.lower():
            match = re.search(
                r"row\s*[:=]\s*(\d+)\s*[,;\s]+col(?:umn)?\s*[:=]\s*([\w_]+)\s*[,;\s]+issue\s*[:=]\s*([\w_]+)",
                line,
                re.IGNORECASE,
            )
            if match:
                normalized = f"row:{match.group(1)},col:{match.group(2).lower()},issue:{match.group(3).lower()}"
                issues.append(normalized)
    return issues


def parse_fix_response(response: str) -> list[str]:
    """Extract fix lines from LLM response."""
    fixes = []
    for line in response.strip().split("\n"):
        line = line.strip()
        if not line:
            continue
        line = re.sub(r"^\s*[\d]+[.\)]\s*", "", line)
        line = re.sub(r"^\s*[-*]\s*", "", line)
        line = line.strip()
        if "row" in line.lower() and "fix" in line.lower():
            match = re.search(
                r"row\s*[:=]\s*(\d+)\s*[,;\s]+col(?:umn)?\s*[:=]\s*([\w_]+)\s*[,;\s]+fix\s*[:=]\s*(.+?)$",
                line,
                re.IGNORECASE,
            )
            if match:
                normalized = f"row:{match.group(1)},col:{match.group(2).lower()},fix:{match.group(3).strip()}"
                fixes.append(normalized)
    return fixes


def call_llm(client: OpenAI, system_prompt: str, user_prompt: str) -> str:
    """Call the LLM with retry on rate limit."""
    for attempt in range(3):
        try:
            response = client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                temperature=0.1,
                max_tokens=2048,
            )
            return response.choices[0].message.content or ""
        except Exception as e:
            if "rate_limit" in str(e).lower() or "429" in str(e):
                wait = 10 * (attempt + 1)
                print(f"[DEBUG] Rate limited, waiting {wait}s...", file=sys.stderr, flush=True)
                time.sleep(wait)
            else:
                print(f"[DEBUG] LLM call failed: {e}", file=sys.stderr, flush=True)
                return ""
    return ""


def run_task(client: OpenAI, env: EnvHTTPClient, task_id: str) -> float:
    """
    Run a single task with two-phase strategy:
      Step 1: Identify issues only
      Step 2: Identify + Fix (using feedback from step 1)
      Step 3: Refined identify + fix (if needed)
    """
    log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)

    rewards: List[float] = []
    steps_taken = 0
    best_score = 0.0
    success = False

    try:
        reset_response = env.reset(task_id=task_id)
        observation = reset_response.get("observation", reset_response)

        last_issues: list[str] = []
        last_llm_output = ""

        for step_num in range(1, MAX_STEPS_PER_TASK + 1):
            error_msg = None

            # ── Phase 1: Identify issues ──
            user_prompt = build_user_prompt(observation)
            identify_output = call_llm(client, IDENTIFY_SYSTEM_PROMPT, user_prompt)
            issues = parse_llm_response(identify_output)

            if not issues and not error_msg:
                error_msg = "no issues parsed from LLM response"

            # ── Phase 2: Propose fixes (from step 2 onward, or always if we have issues) ──
            fixes: list[str] = []
            if issues and step_num >= 2:
                # Build a fix prompt that includes the identified issues
                fix_prompt = build_user_prompt(observation, include_fixes=True)
                fix_prompt += f"\n\nISSUES FOUND:\n" + "\n".join(issues)
                fix_output = call_llm(client, FIX_SYSTEM_PROMPT, fix_prompt)
                fixes = parse_fix_response(fix_output)

            # ── Submit to environment ──
            action_str = ";".join(issues[:5]) if issues else "none"
            if fixes:
                action_str += "|fixes:" + ";".join(fixes[:3])

            step_response = env.step(issues, fixes, task_id=task_id)
            observation = step_response.get("observation", step_response)

            reward = float(step_response.get("reward", 0.0) or 0.0)
            done = bool(step_response.get("done", False))
            best_score = max(best_score, reward)
            rewards.append(reward)
            steps_taken = step_num

            log_step(
                step=step_num,
                action=action_str,
                reward=reward,
                done=done,
                error=error_msg,
            )

            if done:
                break

            last_issues = issues
            last_llm_output = identify_output

        success = best_score >= 0.5

    finally:
        log_end(success=success, steps=steps_taken, score=best_score, rewards=rewards)

    return best_score


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

def main():
    print(f"[DEBUG] DataQA Inference starting", file=sys.stderr, flush=True)
    print(f"[DEBUG] ENV_URL={ENV_URL}", file=sys.stderr, flush=True)
    print(f"[DEBUG] API_BASE_URL={API_BASE_URL}", file=sys.stderr, flush=True)
    print(f"[DEBUG] MODEL_NAME={MODEL_NAME}", file=sys.stderr, flush=True)

    env = EnvHTTPClient(ENV_URL)
    llm_client = OpenAI(
        base_url=API_BASE_URL,
        api_key=API_KEY or "no-key",
    )

    if not env.health():
        print("[DEBUG] Environment is not healthy. Exiting.", file=sys.stderr, flush=True)
        sys.exit(1)

    print(f"[DEBUG] Environment is healthy", file=sys.stderr, flush=True)

    scores = {}
    for task_id in TASKS:
        try:
            score = run_task(llm_client, env, task_id)
            scores[task_id] = score
        except Exception as e:
            print(f"[DEBUG] Task {task_id} failed: {e}", file=sys.stderr, flush=True)
            scores[task_id] = 0.0

    avg_score = sum(scores.values()) / len(scores) if scores else 0.0
    print(f"\n[DEBUG] FINAL RESULTS: {scores} avg={avg_score:.3f}", file=sys.stderr, flush=True)


if __name__ == "__main__":
    main()