inference.py

#!/usr/bin/env python3
"""
Baseline inference script for OrgSim environment.

Runs all 3 tasks (solo_bug_fix, cross_team_launch, startup_crisis) against the
OrgSim environment using an LLM agent via the OpenAI client.

Required env vars:
    API_BASE_URL  - LLM API endpoint
    MODEL_NAME    - Model identifier
    HF_TOKEN      - HuggingFace / API key (used as OpenAI api_key)

Optional:
    ORGSIM_ENV_URL - Environment base URL (default: http://localhost:8000)
"""

import json
import os
import sys
import textwrap
from typing import Optional

try:
    from openai import OpenAI
except ImportError:
    print("ERROR: openai package not installed", file=sys.stderr)
    sys.exit(1)

from org_sim import OrgSimEnv, OrgAction

API_BASE_URL = os.getenv("API_BASE_URL")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o")
HF_TOKEN = os.getenv("HF_TOKEN")

if not API_BASE_URL or not HF_TOKEN:
    print("ERROR: API_BASE_URL and HF_TOKEN must be set", file=sys.stderr)
    sys.exit(1)

client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)

TASKS = ["solo_bug_fix", "cross_team_launch", "startup_crisis"]
ENV_NAME = "org_sim"


# --------------------------------------------------------------------------- #
# Exact log format — do not change field names, ordering, or format strings   #
# --------------------------------------------------------------------------- #

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,
    )


# --------------------------------------------------------------------------- #
# LLM decision logic                                                           #
# --------------------------------------------------------------------------- #

def get_model_action(step: int, obs, last_reward: float, history: list[str]) -> OrgAction:
    """Use LLM to decide next action."""
    history_block = "\n".join(history[-4:]) if history else "None"

    system_prompt = textwrap.dedent(f"""
        You are an agent in an organization simulation (OrgSim).
        Agent ID: {obs.my_agent_id}
        Team: {obs.my_team}
        Role: {obs.my_role}

        Available tasks: {obs.available_tasks}
        Active task: {obs.active_task}
        Inbox: {obs.inbox}
        Team status: {obs.team_status}
        Resources: {obs.resources}
        Metrics: {obs.metrics}

        Available actions:
        - REQUEST_TASK: Get next task from your team queue (no payload needed)
        - ACCEPT_TASK: payload={{"task_id": "<id>"}}
        - COMPLETE_TASK: payload={{"task_id": "<id>"}} — only when you have an active task
        - REQUEST_HELP: payload={{"task_id": "<id>"}} — advances progress on your task
        - PROVIDE_HELP: payload={{"task_id": "<id>"}}
        - ESCALATE: payload={{"task_id": "<id>"}} — for cross-team or stuck tasks
        - REQUEST_RESOURCE: payload={{"resource_id": "<id>"}} — lock senior_engineer before feature tasks
        - REPORT_STATUS: (no payload)

        Strategy hints:
        1. For startup_crisis: REQUEST_RESOURCE(senior_engineer) FIRST, then tackle the critical incident.
        2. For cross-team tasks you can't do yourself, ESCALATE them.
        3. Use REQUEST_HELP to build progress before attempting COMPLETE_TASK.

        Respond ONLY with valid JSON: {{"action_type": "...", "target_id": "...", "payload": {{}}}}
    """).strip()

    user_prompt = textwrap.dedent(f"""
        Step: {step}
        Last reward: {last_reward:.2f}
        Previous steps:
        {history_block}
        Send your next action.
    """).strip()

    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.3,
        )
        content = response.choices[0].message.content.strip()
        # Strip markdown code blocks if present
        if content.startswith("```"):
            content = content.split("```")[1]
            if content.startswith("json"):
                content = content[4:]
        action_data = json.loads(content)
        return OrgAction(
            action_type=action_data.get("action_type", "REQUEST_TASK"),
            target_id=action_data.get("target_id", ""),
            payload=action_data.get("payload", {}),
        )
    except Exception:
        return OrgAction(action_type="REQUEST_TASK", target_id="", payload={})


# --------------------------------------------------------------------------- #
# Main loop                                                                    #
# --------------------------------------------------------------------------- #

def run_task(env_url: str, task_id: str) -> tuple[bool, int, float, list[float]]:
    """Run one episode for a given task. Returns (success, steps, score, rewards)."""
    rewards: list[float] = []

    with OrgSimEnv(base_url=env_url).sync() as env:
        result = env.reset(task_id=task_id)

        step_count = 0
        history: list[str] = []
        last_reward = 0.0
        error_msg = None

        while not result.done:
            step_count += 1
            obs = result.observation

            try:
                action = get_model_action(step_count, obs, last_reward, history)
                error_msg = None
            except Exception as e:
                action = OrgAction(action_type="REQUEST_TASK", target_id="", payload={})
                error_msg = str(e)

            try:
                result = env.step(action)
                last_reward = result.reward
                rewards.append(result.reward)
                history.append(f"step={step_count} action={action.action_type} reward={result.reward:.2f}")
            except Exception as e:
                error_msg = str(e)
                last_reward = 0.0
                rewards.append(0.0)

            log_step(
                step=step_count,
                action=action.action_type,
                reward=last_reward,
                done=result.done,
                error=error_msg,
            )

        # Get final grade from /grade endpoint
        try:
            import httpx
            resp = httpx.get(f"{env_url}/grade", timeout=10.0)
            score = resp.json().get("score", 0.0)
        except Exception:
            # Fallback: compute from completion ratio
            metrics = result.observation.metrics if result else {}
            completed = metrics.get("tasks_completed", 0)
            total = completed + metrics.get("tasks_failed", 0) + metrics.get("tasks_escalated", 0)
            score = completed / max(1, total)

        success = score > 0.0
        return success, step_count, score, rewards


def main():
    env_url = os.getenv("ORGSIM_ENV_URL", "http://localhost:8000")

    for task_id in TASKS:
        log_start(task=task_id, env=ENV_NAME, model=MODEL_NAME)
        success, steps, score, rewards = run_task(env_url, task_id)
        log_end(success=success, steps=steps, score=score, rewards=rewards)


if __name__ == "__main__":
    main()