inference.py

"""
Inference Script — CSV Cleaner Environment
=============================================
Baseline agent using OpenAI client to clean CSV datasets across 3 tasks.

MANDATORY ENV VARS:
  API_BASE_URL   The API endpoint for the LLM.
  MODEL_NAME     The model identifier to use for inference.
  HF_TOKEN       Your Hugging Face / API key.
  IMAGE_NAME     Docker image name (if using from_docker_image)

STDOUT FORMAT:
  [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>
"""

import asyncio
import json
import os
import textwrap
from typing import Any, Dict, List, Optional

from openai import OpenAI
from openenv.core.env_server.mcp_types import CallToolAction

from csv_cleaner_env import CsvCleanerEnv

IMAGE_NAME   = os.getenv("IMAGE_NAME")
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")  # default allowed
MODEL_NAME   = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
API_KEY      = os.environ.get("API_KEY", os.environ.get("HF_TOKEN", "dummy_key"))

BENCHMARK   = os.getenv("CSV_CLEANER_BENCHMARK", "csv_cleaner_env")
TEMPERATURE = 0.3
MAX_TOKENS  = 300

# Debug print to confirm env vars are loaded
print(f"[CONFIG] API_BASE_URL={API_BASE_URL} MODEL={MODEL_NAME} API_KEY={'SET' if API_KEY != 'dummy_key' else 'NOT SET'}", flush=True)

# Task configurations
TASKS = [
    {"name": "fix_column_types",          "max_steps": 10},
    {"name": "clean_missing_duplicates",  "max_steps": 15},
    {"name": "full_pipeline",             "max_steps": 20},
]

SYSTEM_PROMPT = textwrap.dedent("""
    You are a data cleaning agent. You interact with a CSV dataset through structured tool calls.

    Available tools:
    - get_dataset_info(): See current columns, types, null counts, samples
    - rename_column(old_name, new_name): Rename a column
    - cast_column(column, dtype): Cast column to int/float/str/datetime
    - fill_missing(column, strategy, value): Fill nulls. strategy: mean/median/mode/constant/zero
    - drop_missing(column): Drop rows with nulls (empty string for all columns)
    - drop_duplicates(columns): Remove duplicates (empty string for all columns)
    - filter_rows(column, operator, value): Filter rows. operator: ==/!=/>/</contains
    - strip_whitespace(column): Strip whitespace from string column
    - replace_values(column, old_value, new_value): Replace values in column

    You must respond with EXACTLY ONE tool call per turn as a JSON object:
    {"tool": "<tool_name>", "args": {"param1": "value1", ...}}

    Read the task description carefully and execute the cleaning steps one at a time.
    Start by calling get_dataset_info to understand the current state, then fix issues.
""").strip()


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:.2f} rewards={rewards_str}",
        flush=True,
    )


def parse_tool_call(text: str) -> Optional[Dict[str, Any]]:
    """Extract JSON tool call from model response."""
    text = text.strip()
    for start_char in ["{"]:
        start = text.find(start_char)
        if start == -1:
            continue
        depth = 0
        for i in range(start, len(text)):
            if text[i] == "{":
                depth += 1
            elif text[i] == "}":
                depth -= 1
            if depth == 0:
                try:
                    return json.loads(text[start : i + 1])
                except json.JSONDecodeError:
                    continue
    return None


def normalize_tool_call(tool_call: Dict[str, Any]) -> tuple[str, Dict[str, Any]]:
    """Normalize model output into a safe tool name + args payload."""
    tool_name = tool_call.get("tool", "get_dataset_info")
    tool_args = tool_call.get("args", {})

    if not isinstance(tool_args, dict):
        tool_args = {}

    # Model outputs sometimes include nulls for string fields; FastMCP rejects None for str args.
    normalized_args: Dict[str, Any] = {}
    for key, value in tool_args.items():
        normalized_args[key] = "" if value is None else value

    return tool_name, normalized_args


def parse_dataset_snapshot(result_str: str) -> Optional[Dict[str, Any]]:
    """Parse dataset info payload when a tool returns JSON snapshot text."""
    try:
        payload = json.loads(result_str)
    except Exception:
        return None

    if isinstance(payload, dict) and "columns" in payload:
        return payload
    return None


def get_model_response(
    client: OpenAI,
    task_desc: str,
    dataset_info: str,
    last_result: str,
    step: int,
    history: List[str],
) -> Optional[Dict[str, Any]]:
    """Get next tool call from the model."""
    history_block = "\n".join(history[-6:]) if history else "None"

    user_prompt = textwrap.dedent(f"""
        Task: {task_desc}

        Current Step: {step}
        Last Action Result: {last_result}

        Current Dataset State:
        {dataset_info}

        Previous Actions:
        {history_block}

        Respond with your next tool call as JSON: {{"tool": "tool_name", "args": {{...}}}}
    """).strip()

    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user",   "content": user_prompt},
            ],
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            stream=False,
        )
        text = (completion.choices[0].message.content or "").strip()
        return parse_tool_call(text)
    except Exception as exc:
        error_text = str(exc)
        print(f"[DEBUG] Model request failed: {error_text}", flush=True)

        # Stop the task early when provider quota is exhausted instead of
        # repeatedly falling back to no-op tool calls.
        if "Error code: 402" in error_text or "depleted your monthly included credits" in error_text:
            return {"tool": "__quota_exhausted__", "args": {}}

        return None


async def run_task(client: OpenAI, env: CsvCleanerEnv, task_config: Dict) -> None:
    """Run a single task and log stdout."""
    task_name = task_config["name"]
    max_steps = task_config["max_steps"]

    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)

    rewards:     List[float] = []
    steps_taken: int         = 0
    score:       float       = 0.0
    success:     bool        = False
    done:        bool        = False

    try:
        result     = await env.reset(task=task_name)
        obs_metadata = getattr(result.observation, "metadata", getattr(result, "metadata", {})) or {}
        task_desc  = obs_metadata.get("task_description", task_name)
        dataset_info = json.dumps(obs_metadata.get("columns", []), indent=2)
        last_result  = obs_metadata.get("last_action_result", "Ready")
        history: List[str] = []

        for step in range(1, max_steps + 1):
            if result.done:
                break

            # First step: always get dataset info
            if step == 1:
                tool_call = {"tool": "get_dataset_info", "args": {}}
            else:
                tool_call = get_model_response(
                    client, task_desc, dataset_info, last_result, step, history
                )
                if tool_call is None:
                    tool_call = {"tool": "get_dataset_info", "args": {}}

            tool_name, tool_args = normalize_tool_call(tool_call)

            if tool_name == "__quota_exhausted__":
                print("[DEBUG] Stopping task early due to provider quota exhaustion", flush=True)
                break

            try:
                action = CallToolAction(tool_name=tool_name, arguments=tool_args)
                result = await env.step(action)

                obs = result.observation
                obs_error = getattr(obs, "error", None)
                if obs_error is not None:
                    result_str = f"Error: {getattr(obs_error, 'message', str(obs_error))}"
                else:
                    obs_result = getattr(obs, "result", None)
                    if hasattr(obs_result, "data"):
                        obs_result = obs_result.data
                    elif isinstance(obs_result, dict) and "data" in obs_result:
                        obs_result = obs_result["data"]
                    result_str = str(obs_result) if obs_result is not None else ""
            except Exception as e:
                result_str = f"Error: {e}"

            reward = result.reward if hasattr(result, "reward") and result.reward else 0.0
            done = result.done if hasattr(result, "done") else False

            obs_metadata = getattr(result.observation, "metadata", getattr(result, "metadata", {})) or {}
            snapshot = parse_dataset_snapshot(result_str)

            state_payload: Dict[str, Any] = {}
            if isinstance(obs_metadata, dict):
                state_payload.update(obs_metadata)
            if isinstance(snapshot, dict):
                state_payload.update(snapshot)

            if state_payload:
                progress = state_payload.get("progress")
                if isinstance(progress, (int, float)):
                    score = float(progress)

                columns = state_payload.get("columns")
                if isinstance(columns, list):
                    dataset_info = json.dumps(columns, indent=2)

                task_desc = state_payload.get("task_description", task_desc)
                last_result = state_payload.get("last_action_result", result_str)
            else:
                last_result = result_str

            rewards.append(reward)
            steps_taken = step
            action_str  = f"{tool_name}({json.dumps(tool_args)})"
            log_step(step=step, action=action_str, reward=reward, done=done, error=None)
            history.append(f"Step {step}: {action_str} -> {last_result[:100]}")

            if done:
                break

        if score == 0.0 and rewards:
            # Fallback when progress metadata is unavailable from the client payload.
            score = min(1.0, max(0.0, sum(rewards)))

        score = min(max(score, 0.0), 1.0)

        # Environment marks done=True on success or step-limit. If it ended before max steps,
        # that's a reliable success signal even when explicit progress metadata is absent.
        success = (done and steps_taken < max_steps) or score >= 0.95

    except Exception as e:
        print(f"[DEBUG] Task error: {e}", flush=True)
    finally:
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)


async def main() -> None:
    # Use API_KEY as the API key — injected by the hackathon validator
    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    if IMAGE_NAME:
        env = await CsvCleanerEnv.from_docker_image(IMAGE_NAME)
    else:
        # Since uvicorn is already running on port 8000 inside the HF Space container, connect locally
        env = CsvCleanerEnv(base_url="http://localhost:8000")
        await env.connect()

    try:
        for task_config in TASKS:
            await run_task(client, env, task_config)
    finally:
        try:
            await env.close()
        except Exception as e:
            print(f"[DEBUG] env.close() error: {e}", flush=True)


if __name__ == "__main__":
    asyncio.run(main())