inference.py

"""Baseline inference script for PM-Ops environment."""
import json
import os

API_BASE_URL = os.getenv("API_BASE_URL", "https://adityaguntur-pm-ops.hf.space")
MODEL_NAME = os.getenv("MODEL_NAME", "claude-opus-4-5")
HF_TOKEN = os.getenv("HF_TOKEN")
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")

MAX_STEPS = 40
TASK_IDS = ["triage", "incident_routing", "release_notes", "dep_update"]


def baseline_agent(obs: dict, org_config: dict) -> tuple[str, dict]:
    """Heuristic baseline: read runbook -> create ticket -> assign -> notify -> finish."""
    step = obs.get("step", 0)
    task_brief = obs.get("task_brief", "")
    last = obs.get("last_action_result", {})
    data = last.get("data", {})

    if step == 0:
        return "meta.read_runbook", {}

    if step == 1 and isinstance(data, dict) and "org_config" in data:
        org = data["org_config"]
        org_config.update(org)
        label_taxonomy = org.get("label_taxonomy", {})
        priority_levels = org.get("priority_levels", ["P1"])
        label = list(label_taxonomy.values())[0] if label_taxonomy else "bug"
        priority = priority_levels[1] if len(priority_levels) > 1 else priority_levels[0]
        teams = list(org_config.get("team_map", {}).values())
        assignee = teams[0] if teams else "backend"
        return "ticketing.create_ticket", {
            "summary": f"Issue: {task_brief[:80]}",
            "description": task_brief,
            "label": label,
            "priority": priority,
            "assignee": assignee,
        }

    if step == 2 and isinstance(data, dict) and "id" in data:
        ticket_id = data["id"]
        teams = list(org_config.get("team_map", {}).values())
        team = teams[0] if teams else "backend"
        return "ticketing.assign_ticket", {"ticket_id": ticket_id, "team": team}

    if step == 3:
        return "chat.list_channels", {}

    if step == 4:
        oncall = org_config.get("oncall_channels", {})
        oncall_names = list(oncall.values())
        channels = data if isinstance(data, list) else []
        channel = oncall_names[0] if oncall_names else (channels[0] if channels else "general")
        return "chat.post_message", {
            "channel": channel,
            "text": f"Incident notification: {task_brief[:150]}. Please review and respond.",
        }

    return "meta.finish", {}


def run_episode(task_id: str, env) -> float:
    result = env.reset()
    obs = result.observation if hasattr(result, "observation") else result

    org_config: dict = {}
    rewards = []
    n = 0
    done = False
    score = 0.0

    print(f"[START] task={task_id} env=pm_ops model={MODEL_NAME}", flush=True)

    while not done and n < MAX_STEPS:
        action_type, args = baseline_agent(obs, org_config)

        try:
            result = env.step({"action_type": action_type, "args": args})
            obs = result.observation if hasattr(result, "observation") else result
            r = float(result.reward if hasattr(result, "reward") else obs.get("reward", 0.0))
            done = bool(result.done if hasattr(result, "done") else obs.get("done", False))
            err = obs.get("last_action_result", {}).get("error", "none") if isinstance(obs, dict) else "none"
        except Exception as e:
            r = 0.0
            done = True
            err = str(e)

        rewards.append(r)
        n += 1
        print(f"[STEP] step={n} action={action_type} reward={r:.2f} done={done} error={err}", flush=True)

        if done:
            score = r

    rewards_str = ",".join(f"{x:.2f}" for x in rewards)
    success = score > 0.5
    print(f"[END] success={success} steps={n} score={score:.2f} rewards={rewards_str}", flush=True)
    return score


def main():
    from openenv.core import GenericEnvClient

    async_client = GenericEnvClient(base_url=API_BASE_URL)
    env = async_client.sync()

    total_score = 0.0
    with env:
        for task_id in TASK_IDS:
            score = run_episode(task_id, env)
            total_score += score

    print(f"[SUMMARY] avg_score={total_score / len(TASK_IDS):.3f}", flush=True)


if __name__ == "__main__":
    main()