inference.py

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional

from openai import OpenAI

_REPO_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(_REPO_ROOT))

from env.aether_env import AetherTaskFlowEnvironment


API_KEY = os.getenv("API_KEY", os.getenv("OPENAI_API_KEY", os.getenv("HF_TOKEN", "")))
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
TEMPERATURE = float(os.getenv("TEMPERATURE", "0.1"))
MAX_LLM_TOKENS = int(os.getenv("MAX_LLM_TOKENS", "128"))
USE_LLM = bool(API_KEY)

SYSTEM_PROMPT = """You are an expert workflow orchestration agent inside the AETHER-TaskFlow RL environment.

Each step you receive an observation and must output a single JSON action.

RULES:
- Output ONLY valid JSON with keys: action_type, task_id, reasoning
- action_type must be one of: execute, defer, delegate, optimize
- task_id must be an integer matching a pending task id
- reasoning should be brief and may be empty
- No explanation, no markdown, no extra text - raw JSON only

STRATEGY:
- execute: high-priority, low-uncertainty tasks with sufficient resources
- optimize: before executing high-uncertainty tasks (reduces failure risk)
- delegate: when resources are too low (free action, 35% reward)
- defer: tasks you can handle later when resources recover

Example: {"action_type": "execute", "task_id": 2, "reasoning": "highest value low risk"}"""


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,
) -> None:
    error_val = error if error else "null"
    done_val = str(done).lower()
    action_safe = action.replace("\n", " ").replace("\r", "")[:120]
    print(
        f"[STEP] step={step} action={action_safe} reward={reward:.2f} "
        f"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} "
        f"score={score:.2f} rewards={rewards_str}",
        flush=True,
    )


def _obs_to_prompt(obs_dict: Dict[str, Any]) -> str:
    tasks = obs_dict.get("tasks", [])
    task_lines = []
    for t in tasks:
        task_lines.append(
            f"  id={t['task_id']} name='{t['name']}' priority={t['priority']:.2f} "
            f"deadline={t['deadline']} uncertainty={t['uncertainty']:.2f} "
            f"value={t['value']:.1f} energy_cost={t['required_energy']:.1f} "
            f"budget_cost={t['required_budget']:.1f} status={t.get('status', 'pending')}"
        )
    return (
        f"OBSERVATION:\n"
        f"  time_remaining={obs_dict.get('time_remaining')} "
        f"energy={obs_dict.get('energy_remaining', 0):.1f} "
        f"budget={obs_dict.get('budget_remaining', 0):.1f} "
        f"system_health={obs_dict.get('system_health', 1):.2f}\n"
        f"PENDING TASKS:\n"
        + "\n".join(task_lines)
        + "\n\nOutput your action JSON:"
    )


def _call_llm(client: OpenAI, obs_text: str, history: List[dict]) -> Optional[Dict[str, Any]]:
    """
    Call the LLM using proper OpenAI client (required by judges).
    Falls back to heuristic if it fails.
    """
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages.extend(history[-4:])
    messages.append({"role": "user", "content": obs_text})

    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
            temperature=TEMPERATURE,
            max_tokens=MAX_LLM_TOKENS,
        )
        raw = (completion.choices[0].message.content or "").strip()
        if raw.startswith("```"):
            raw = raw.split("```")[1]
            if raw.startswith("json"):
                raw = raw[4:].strip()
        parsed = json.loads(raw)
        return {
            "task_id": int(parsed["task_id"]),
            "action_type": str(parsed["action_type"]),
            "reasoning": str(parsed.get("reasoning", "")),
        }
    except Exception as exc:
        print(f"[DEBUG] LLM call/parse failed: {exc}", flush=True)
        return None


def get_llm_action(
    client: OpenAI,
    obs_dict: Dict[str, Any],
    history: List[dict],
) -> Optional[Dict[str, Any]]:
    """Call the configured OpenAI-compatible endpoint for an LLM action."""
    obs_text = _obs_to_prompt(obs_dict)
    result = _call_llm(client, obs_text, history)
    if result is not None:
        history.append({"role": "user", "content": obs_text})
        history.append({"role": "assistant", "content": json.dumps(result)})
    return result


def get_heuristic_action(env: AetherTaskFlowEnvironment) -> Dict[str, Any]:
    """Built-in AETHER + RAPTOR heuristic - no API required."""
    action = env.message_to_action("")
    if hasattr(action, "model_dump"):
        return action.model_dump(exclude={"reasoning"})
    return {
        "action_type": getattr(action, "action_type", "execute").value
        if hasattr(action, "action_type")
        else "execute",
        "task_id": getattr(action, "task_id", 0),
    }


def get_action(
    env: AetherTaskFlowEnvironment,
    client: OpenAI,
    obs_dict: Dict[str, Any],
    history: List[dict],
) -> Dict[str, Any]:
    """Return LLM action if an API key is set, otherwise heuristic."""
    if USE_LLM:
        result = get_llm_action(client, obs_dict, history)
        if result is not None:
            return result

    heuristic = get_heuristic_action(env)
    obs_text = _obs_to_prompt(obs_dict)
    history.append({"role": "user", "content": obs_text})
    history.append({"role": "assistant", "content": json.dumps(heuristic)})
    return heuristic


def run_episode(difficulty: str, client: OpenAI) -> None:
    os.environ["AETHER_DIFFICULTY"] = difficulty
    env = AetherTaskFlowEnvironment(difficulty=difficulty)

    model_label = MODEL_NAME if USE_LLM else "HEURISTIC-AETHER-RAPTOR"
    log_start(task=difficulty, env="aether_taskflow", model=model_label)

    obs = env.reset()
    obs_dict = obs.model_dump() if hasattr(obs, "model_dump") else obs

    rewards: List[float] = []
    history: List[dict] = []
    step = 0

    while True:
        step += 1
        action_dict = get_action(env, client, obs_dict, history)
        next_obs = env.step(action_dict)

        next_obs_dict = next_obs.model_dump() if hasattr(next_obs, "model_dump") else next_obs
        reward = next_obs_dict.get("reward", 0.0)
        done = next_obs_dict.get("done", False)
        rewards.append(reward)

        action_str = (
            f"{action_dict.get('action_type', 'execute')}"
            f"(task_id={action_dict.get('task_id', 0)})"
        )
        log_step(step, action_str, reward, done)

        obs_dict = next_obs_dict
        if done:
            break

    score = env.compute_final_score()
    log_end(success=True, steps=step, score=score, rewards=rewards)


def main() -> None:
    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)

    parser = argparse.ArgumentParser(description="AETHER-TaskFlow Inference")
    parser.add_argument(
        "--single",
        choices=["easy", "medium", "hard"],
        default=None,
        help="Run a single difficulty (default: all three)",
    )
    args = parser.parse_args()

    difficulties = [args.single] if args.single else ["easy", "medium", "hard"]
    for diff in difficulties:
        run_episode(diff, client)


if __name__ == "__main__":
    main()