inference.py

"""
Quota-aware baseline inference script for the ML Training Optimizer environment.

This baseline uses the OpenAI Python client against Gemini's OpenAI-compatible
endpoint and keeps request pacing under a hard RPM ceiling.

Environment Variables:
    API_BASE_URL           - LLM API endpoint
    MODEL_NAME             - model identifier
    HF_TOKEN               - API key for auth
    ENV_URL                - environment server URL (default: http://localhost:8000)
    LLM_RPM_LIMIT          - max model requests per minute (default: 5)
    LLM_MAX_RETRIES        - max rate-limit retries per request (default: 3)
    LLM_REASONING_EFFORT   - reasoning effort sent to the model (default: minimal)
    LLM_MAX_STEPS_EASY     - max model decisions for easy task (default: 5)
    LLM_MAX_STEPS_MEDIUM   - max model decisions for medium task (default: 6)
    LLM_MAX_STEPS_HARD     - max model decisions for hard task (default: 7)
"""

from __future__ import annotations

import json
import os
import sys
import time
from collections import deque
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional

from dotenv import load_dotenv
from openai import OpenAI, RateLimitError
from openenv.core.env_server.mcp_types import CallToolAction

from client import MLTrainerEnv
from server.tasks import TASKS as SERVER_TASKS

load_dotenv()

API_BASE_URL = os.environ.get(
    "API_BASE_URL",
    "https://api.openai.com/v1",
)
MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini")

HF_TOKEN = (
    os.environ.get("HF_TOKEN")
    or os.environ.get("OPENAI_API_KEY")
    or os.environ.get("OPENROUTER_API_KEY")
    or os.environ.get("GEMINI_API_KEY")
)

ENV_URL = os.environ.get("ENV_URL", "http://localhost:8000")

LLM_RPM_LIMIT = max(1, int(os.environ.get("LLM_RPM_LIMIT", "5")))
LLM_MAX_RETRIES = max(0, int(os.environ.get("LLM_MAX_RETRIES", "3")))
LLM_REASONING_EFFORT = os.environ.get("LLM_REASONING_EFFORT", "minimal")
LLM_FORCE_TOOL_CHOICE = os.environ.get("LLM_FORCE_TOOL_CHOICE", "true").lower() == "true"
LLM_REASONING_MODELS = ["gemini", "o1", "deepseek-r1"]

LLM_MAX_STEPS = {
    "easy_mnist": max(1, int(os.environ.get("LLM_MAX_STEPS_EASY", "5"))),
    "medium_fashion": max(1, int(os.environ.get("LLM_MAX_STEPS_MEDIUM", "6"))),
    "hard_cifar": max(1, int(os.environ.get("LLM_MAX_STEPS_HARD", "7"))),
}
MIN_REQUEST_GAP_SECONDS = (60.0 / LLM_RPM_LIMIT) + 0.5

TASKS = ["easy_mnist", "medium_fashion", "hard_cifar"]
VALID_TOOLS = {
    "configure_training",
    "run_epochs",
    "adjust_learning_rate",
    "toggle_augmentation",
    "get_training_status",
    "submit_model",
}

RETRY_BACKOFF_SECONDS = [15, 30, 60]

LOCAL_TASK_METADATA = {
    task_id: {
        "task_id": task.task_id,
        "task_name": task.name,
        "task_description": task.description,
        "difficulty": task.difficulty,
        "model_type": task.model_type,
        "dataset": task.dataset_name,
        "max_epochs": task.max_epochs,
        "target_metric": task.target_metric,
        "target_value": task.target_value,
    }
    for task_id, task in SERVER_TASKS.items()
}

SYSTEM_PROMPT = """You are an ML training optimization agent.

You are controlling a real CPU PyTorch training environment through tool calls.
Your objective is to maximize validation quality with as few actions as possible.

Rules:
- Always respond by calling exactly one tool.
- First decision must configure training.
- After run_epochs, use the returned metrics instead of calling get_training_status immediately.
- Prefer larger epoch chunks early to conserve request budget.
- Only adjust learning rate or augmentation when metrics justify it.
- Submit once progress has stalled or the decision budget is nearly exhausted.

Task-specific epoch guidance:
- easy_mnist: prefer run_epochs(10-15) early
- medium_fashion: prefer run_epochs(8-12) early
- hard_cifar: prefer run_epochs(5-8) early
"""

TOOL_DEFINITIONS = [
    {
        "type": "function",
        "function": {
            "name": "configure_training",
            "description": "Configure training hyperparameters before training starts or when retuning.",
            "parameters": {
                "type": "object",
                "properties": {
                    "optimizer": {"type": "string", "enum": ["sgd", "adam", "adamw"]},
                    "learning_rate": {"type": "number"},
                    "batch_size": {"type": "integer", "enum": [32, 64, 128, 256]},
                    "weight_decay": {"type": "number"},
                    "dropout": {"type": "number"},
                    "lr_schedule": {
                        "type": "string",
                        "enum": ["constant", "step", "cosine", "warmup_cosine"],
                    },
                    "warmup_epochs": {"type": "integer"},
                    "augmentation": {"type": "boolean"},
                    "augmentation_strength": {"type": "number"},
                },
                "required": [
                    "optimizer",
                    "learning_rate",
                    "batch_size",
                    "weight_decay",
                    "dropout",
                    "lr_schedule",
                    "warmup_epochs",
                    "augmentation",
                    "augmentation_strength",
                ],
                "additionalProperties": False,
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "run_epochs",
            "description": "Run real training for a chunk of epochs.",
            "parameters": {
                "type": "object",
                "properties": {
                    "num_epochs": {"type": "integer", "minimum": 1, "maximum": 20},
                },
                "required": ["num_epochs"],
                "additionalProperties": False,
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "adjust_learning_rate",
            "description": "Update the live optimizer learning rate.",
            "parameters": {
                "type": "object",
                "properties": {
                    "new_lr": {"type": "number"},
                },
                "required": ["new_lr"],
                "additionalProperties": False,
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "toggle_augmentation",
            "description": "Enable or disable data augmentation and set strength.",
            "parameters": {
                "type": "object",
                "properties": {
                    "enabled": {"type": "boolean"},
                    "strength": {"type": "number"},
                },
                "required": ["enabled", "strength"],
                "additionalProperties": False,
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_training_status",
            "description": "Read the current metrics without advancing training.",
            "parameters": {
                "type": "object",
                "properties": {},
                "additionalProperties": False,
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "submit_model",
            "description": "Submit the current best model for grading.",
            "parameters": {
                "type": "object",
                "properties": {},
                "additionalProperties": False,
            },
        },
    },
]


class InferenceError(RuntimeError):
    """Raised when the baseline cannot continue safely."""


@dataclass
class TaskStats:
    """Tracks quota and runtime stats for a single task."""

    requests: int = 0
    retries: int = 0
    decisions: int = 0
    elapsed_seconds: float = 0.0


@dataclass
class RequestScheduler:
    """Ensures model requests stay under the configured RPM limit."""

    min_gap_seconds: float
    rpm_limit: int
    time_fn: Callable[[], float] = time.monotonic
    sleep_fn: Callable[[float], None] = time.sleep
    last_request_started_at: Optional[float] = None
    request_timestamps: deque[float] = None

    def __post_init__(self) -> None:
        if self.request_timestamps is None:
            self.request_timestamps = deque()

    def wait_for_turn(self) -> None:
        """Sleep until the next request is allowed."""
        now = self.time_fn()
        while self.request_timestamps and now - self.request_timestamps[0] >= 60.0:
            self.request_timestamps.popleft()

        waits = []
        if self.last_request_started_at is not None:
            waits.append(self.min_gap_seconds - (now - self.last_request_started_at))
        if len(self.request_timestamps) >= self.rpm_limit:
            waits.append(60.0 - (now - self.request_timestamps[0]) + 0.5)

        wait_seconds = max([w for w in waits if w > 0], default=0.0)
        if wait_seconds > 0:
            self.sleep_fn(wait_seconds)
            now = self.time_fn()
            while self.request_timestamps and now - self.request_timestamps[0] >= 60.0:
                self.request_timestamps.popleft()

        self.last_request_started_at = now
        self.request_timestamps.append(now)


def create_client() -> OpenAI:
    """Create an OpenAI client configured for the selected LLM provider."""
    # OpenRouter specific headers for better visibility and ranking
    extra_headers = {}
    if "openrouter.ai" in API_BASE_URL.lower():
        extra_headers = {
            "HTTP-Referer": "https://openenv.google.github.com",
            "X-Title": "OpenEnv ML Trainer",
        }

    return OpenAI(
        api_key=HF_TOKEN,
        base_url=API_BASE_URL,
        default_headers=extra_headers,
    )


def extract_reset_metadata(reset_result: Dict[str, Any]) -> Dict[str, Any]:
    """Normalize reset payloads across OpenEnv HTTP server variants."""
    if not isinstance(reset_result, dict):
        return {}

    observation = reset_result.get("observation")
    if isinstance(observation, dict) and isinstance(observation.get("metadata"), dict):
        metadata = observation["metadata"]
        if metadata:
            return metadata

    if isinstance(observation, dict):
        parsed = extract_result_dict(observation.get("result", {}))
        if parsed:
            return parsed

    if isinstance(reset_result.get("metadata"), dict):
        metadata = reset_result["metadata"]
        if metadata:
            return metadata

    result = reset_result.get("result", {})
    parsed = extract_result_dict(result)
    if parsed:
        return parsed

    return {}


def extract_result_data(payload: Dict[str, Any]) -> Dict[str, Any]:
    """Normalize tool results across OpenEnv HTTP payload variants."""
    if not isinstance(payload, dict):
        return {}

    observation = payload.get("observation", payload)
    if not isinstance(observation, dict):
        return {}

    tool_result = observation.get("result", {})
    parsed = extract_result_dict(tool_result)
    if parsed:
        return parsed

    metadata = observation.get("metadata", {})
    return metadata if isinstance(metadata, dict) else {}


def extract_result_dict(tool_result: Any) -> Dict[str, Any]:
    """Extract structured JSON-ish payloads from OpenEnv tool results."""
    if not isinstance(tool_result, dict):
        return {}

    for key in ("data", "structured_content"):
        value = tool_result.get(key)
        if isinstance(value, dict) and value:
            return value

    content = tool_result.get("content")
    if isinstance(content, list):
        for item in content:
            if not isinstance(item, dict):
                continue
            text = item.get("text")
            if not isinstance(text, str):
                continue
            try:
                parsed = json.loads(text)
            except json.JSONDecodeError:
                continue
            if isinstance(parsed, dict) and parsed:
                return parsed

    return {}


def extract_observation_metadata(observation: Any) -> Dict[str, Any]:
    """Read metadata from either object or dict observations."""
    if isinstance(observation, dict):
        metadata = observation.get("metadata", {})
        return metadata if isinstance(metadata, dict) else {}

    metadata = getattr(observation, "metadata", {})
    return metadata if isinstance(metadata, dict) else {}


def merge_task_metadata(task_id: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
    """Fill missing task fields from local task definitions."""
    merged = dict(LOCAL_TASK_METADATA.get(task_id, {}))
    merged.update({k: v for k, v in metadata.items() if v not in (None, "", [], {})})
    return merged


def extract_tool_result_from_observation(observation: Any) -> Dict[str, Any]:
    """Read tool result content from either object or dict observations."""
    if isinstance(observation, dict):
        return extract_result_data({"observation": observation})

    # Try the structured result first
    result = getattr(observation, "result", None)
    parsed = extract_result_dict(result)
    if parsed:
        return parsed

    # Fallback: if result itself is already a dict with metrics keys, use it directly
    if isinstance(result, dict) and result:
        return result

    # Fallback: check metadata for embedded tool results
    metadata = getattr(observation, "metadata", None)
    if isinstance(metadata, dict) and metadata:
        # Check if metadata itself contains training metrics
        if any(k in metadata for k in ("current_epoch", "train_loss", "val_accuracy", "best_val_accuracy")):
            return metadata

    return {}


def normalize_tool_result(result_data: Dict[str, Any]) -> Dict[str, Any]:
    """Flatten tool payloads so prompts receive a consistent metrics view."""
    if not isinstance(result_data, dict):
        return {}

    final_metrics = result_data.get("final_metrics")
    if isinstance(final_metrics, dict):
        merged = dict(final_metrics)
        if "status" in result_data:
            merged["status"] = result_data["status"]
        if "grade" in result_data:
            merged["grade"] = result_data["grade"]
        return merged

    metrics = result_data.get("metrics")
    if isinstance(metrics, dict):
        merged = dict(metrics)
        if "status" in result_data:
            merged["status"] = result_data["status"]
        if "grade" in result_data:
            merged["grade"] = result_data["grade"]
        if "final_metrics" in result_data:
            merged["final_metrics"] = result_data["final_metrics"]
        return merged

    return result_data


def apply_tool_context(
    tool_name: str,
    arguments: Dict[str, Any],
    latest_data: Dict[str, Any],
    normalized_data: Dict[str, Any],
) -> Dict[str, Any]:
    """Carry forward config state that the environment does not echo on every tool."""
    merged = dict(latest_data)
    merged.update(normalized_data)

    if tool_name == "configure_training":
        merged["current_config"] = dict(arguments)
    elif tool_name == "adjust_learning_rate":
        current_config = dict(merged.get("current_config") or {})
        current_config["learning_rate"] = arguments.get("new_lr")
        merged["current_config"] = current_config
    elif tool_name == "toggle_augmentation":
        current_config = dict(merged.get("current_config") or {})
        current_config["augmentation"] = arguments.get("enabled")
        current_config["augmentation_strength"] = arguments.get("strength")
        merged["current_config"] = current_config

    return merged


def compact_config_summary(data: Dict[str, Any]) -> str:
    """Render the active config into a short deterministic string."""
    config = data.get("current_config")
    if not isinstance(config, dict):
        return "none"

    ordered_keys = [
        "optimizer",
        "learning_rate",
        "batch_size",
        "weight_decay",
        "dropout",
        "lr_schedule",
        "augmentation",
        "augmentation_strength",
    ]
    parts = []
    for key in ordered_keys:
        if key in config:
            parts.append(f"{key}={config[key]}")
    return ", ".join(parts) if parts else "none"


def compact_state_summary(metadata: Dict[str, Any], latest_data: Dict[str, Any]) -> str:
    """Create a compact prompt-friendly environment summary."""
    if latest_data:
        fields = {
            "epoch": latest_data.get("current_epoch", 0),
            "remaining_budget": latest_data.get("remaining_budget", metadata.get("max_epochs", "?")),
            "train_loss": latest_data.get("train_loss", 0.0),
            "val_loss": latest_data.get("val_loss", 0.0),
            "train_accuracy": latest_data.get("train_accuracy", 0.0),
            "val_accuracy": latest_data.get("val_accuracy", 0.0),
            "best_val_accuracy": latest_data.get("best_val_accuracy", 0.0),
            "convergence_signal": latest_data.get("convergence_signal", "not_started"),
            "is_diverged": latest_data.get("is_diverged", False),
            "status": latest_data.get("status"),
            "error": latest_data.get("error"),
        }
        summary = (
            "State: "
            f"epoch={fields['epoch']}, remaining_budget={fields['remaining_budget']}, "
            f"train_loss={fields['train_loss']}, val_loss={fields['val_loss']}, "
            f"train_accuracy={fields['train_accuracy']}, val_accuracy={fields['val_accuracy']}, "
            f"best_val_accuracy={fields['best_val_accuracy']}, "
            f"convergence_signal={fields['convergence_signal']}, "
            f"is_diverged={fields['is_diverged']}, "
            f"current_config={compact_config_summary(latest_data)}"
        )
        if fields["status"]:
            summary += f", status={fields['status']}"
        if fields["error"]:
            summary += f", error={fields['error']}"
        return summary

    return (
        "State: not_started, "
        f"max_epochs={metadata.get('max_epochs', '?')}, "
        f"target={metadata.get('target_metric', 'val_accuracy')} >= {metadata.get('target_value', '?')}, "
        "current_config=none"
    )


def history_block(action_summaries: List[str]) -> str:
    """Render the last two action summaries for the prompt."""
    if not action_summaries:
        return "None"
    return "\n".join(action_summaries[-2:])


def task_brief(metadata: Dict[str, Any]) -> str:
    """Create the stable task brief sent on every request."""
    return (
        f"Task: {metadata.get('task_name', metadata.get('task_id', 'unknown'))}\n"
        f"Description: {metadata.get('task_description', '')}\n"
        f"Difficulty: {metadata.get('difficulty', '')}\n"
        f"Model: {metadata.get('model_type', '')}\n"
        f"Dataset: {metadata.get('dataset', '')}\n"
        f"Max epochs: {metadata.get('max_epochs', '')}\n"
        f"Target: {metadata.get('target_metric', '')} >= {metadata.get('target_value', '')}"
    )


def build_messages(
    metadata: Dict[str, Any],
    latest_data: Dict[str, Any],
    action_summaries: List[str],
    decision_index: int,
    max_decisions: int,
) -> List[Dict[str, str]]:
    """Build the compact chat payload for a single model decision."""
    remaining_decisions = max(0, max_decisions - decision_index)
    user_prompt = (
        f"{task_brief(metadata)}\n\n"
        f"{compact_state_summary(metadata, latest_data)}\n"
        f"Recent actions:\n{history_block(action_summaries)}\n\n"
        f"Decision budget remaining: {remaining_decisions}\n"
        "You are under a hard request limit. Minimize tool calls, prefer larger run_epochs chunks early, "
        "and reserve the final decision for submit_model if needed."
    )

    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_prompt},
    ]


def parse_retry_after_seconds(error: RateLimitError) -> Optional[float]:
    """Read Retry-After from the SDK error response when available."""
    response = getattr(error, "response", None)
    headers = getattr(response, "headers", None)
    if not headers:
        return None

    raw_value = headers.get("retry-after")
    if raw_value is None:
        return None

    try:
        return max(0.0, float(raw_value))
    except (TypeError, ValueError):
        return None


def parse_tool_call(message: Any) -> Dict[str, Any]:
    """Extract a single valid tool call from a chat completion message."""
    tool_calls = getattr(message, "tool_calls", None) or []
    if len(tool_calls) != 1:
        raise InferenceError("Expected exactly one tool call from the model.")

    tool_call = tool_calls[0]
    function = getattr(tool_call, "function", None)
    if function is None:
        raise InferenceError("Tool call did not include function details.")

    name = getattr(function, "name", "")
    if name not in VALID_TOOLS:
        raise InferenceError(f"Unsupported tool selected by model: {name}")

    raw_arguments = getattr(function, "arguments", "") or "{}"
    try:
        arguments = json.loads(raw_arguments)
    except json.JSONDecodeError as exc:
        raise InferenceError(f"Tool arguments were not valid JSON: {raw_arguments}") from exc

    if not isinstance(arguments, dict):
        raise InferenceError("Tool arguments must decode to a JSON object.")

    return {"tool_name": name, "arguments": arguments}


def action_summary(tool_name: str, arguments: Dict[str, Any], result_data: Dict[str, Any], reward: float) -> str:
    """Create a compact history line for prompt reuse."""
    best_val = result_data.get("best_val_accuracy", "-")
    val_acc = result_data.get("val_accuracy", "-")
    signal = result_data.get("convergence_signal", "-")
    remaining_budget = result_data.get("remaining_budget", "-")
    return (
        f"- {tool_name}({json.dumps(arguments, sort_keys=True)}) -> "
        f"best_val_accuracy={best_val}, val_accuracy={val_acc}, "
        f"signal={signal}, remaining_budget={remaining_budget}, reward={reward}"
    )


def model_tool_choice(decision_index: int, max_decisions: int) -> Any:
    """Choose the tool policy for the current decision."""
    if decision_index == 0:
        return {"type": "function", "function": {"name": "configure_training"}}
    if decision_index == max_decisions - 1:
        return {"type": "function", "function": {"name": "submit_model"}}
    return "required"


def request_action(
    client: OpenAI,
    scheduler: RequestScheduler,
    messages: List[Dict[str, str]],
    stats: TaskStats,
    decision_index: int,
    max_decisions: int,
) -> Dict[str, Any]:
    """Request a single model-directed tool call with strict rate-limit handling."""
    for attempt in range(LLM_MAX_RETRIES + 1):
        scheduler.wait_for_turn()
        stats.requests += 1

        # Only use reasoning_effort for known reasoning models to avoid 400s
        call_kwargs = {
            "model": MODEL_NAME,
            "messages": messages,
            "tools": TOOL_DEFINITIONS,
            "parallel_tool_calls": False,
            "temperature": 0,
            "max_completion_tokens": 200,
        }
        
        # Check if the model likely supports reasoning effort
        if any(rm in MODEL_NAME.lower() for rm in LLM_REASONING_MODELS):
            call_kwargs["reasoning_effort"] = LLM_REASONING_EFFORT

        # Determine tool_choice policy
        preferred_choice = model_tool_choice(decision_index, max_decisions)
        if not LLM_FORCE_TOOL_CHOICE and preferred_choice != "auto":
             preferred_choice = "auto"
        
        try:
            # Attempt with preferred choice
            response = client.chat.completions.create(
                tool_choice=preferred_choice,
                **call_kwargs
            )
            return parse_tool_call(response.choices[0].message)
        except Exception as exc:
            # Check if error is related to tool_choice compatibility
            err_msg = str(exc).lower()
            if ("tool_choice" in err_msg or "endpoints found" in err_msg) and preferred_choice != "auto":
                print(f"  [Model/Provider does not support forced tool choice, falling back to 'auto'...]")
                stats.retries += 1
                try:
                    response = client.chat.completions.create(
                        tool_choice="auto",
                        **call_kwargs
                    )
                    return parse_tool_call(response.choices[0].message)
                except Exception as nested_exc:
                    raise InferenceError(f"Model failed even with tool_choice='auto': {nested_exc}") from nested_exc
            
            if isinstance(exc, RateLimitError):
                if attempt >= LLM_MAX_RETRIES:
                    raise InferenceError(
                        f"Rate limit persisted after {LLM_MAX_RETRIES} retries."
                    ) from exc

                stats.retries += 1
                retry_after = parse_retry_after_seconds(exc)
                backoff = retry_after if retry_after is not None else RETRY_BACKOFF_SECONDS[min(attempt, len(RETRY_BACKOFF_SECONDS) - 1)]
                print(f"  [Rate limited, waiting {backoff:.1f} seconds before retry...]")
                time.sleep(backoff)
            else:
                raise exc


def run_task(client: OpenAI, scheduler: RequestScheduler, task_id: str) -> Dict[str, Any]:
    """Run one environment task under the quota-aware baseline."""
    max_decisions = LLM_MAX_STEPS[task_id]
    task_stats = TaskStats()
    task_started_at = time.time()

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

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

    try:
        env_instance = MLTrainerEnv(base_url=ENV_URL).sync()
        with env_instance as env:
            try:
                reset_result = env.reset(task_id=task_id)
            except Exception as e:
                # Need to swallow to trigger finally block properly
                print(f"[DEBUG] env.reset() failed: {e}", flush=True)
                raise e

            metadata = merge_task_metadata(task_id, extract_observation_metadata(reset_result.observation))
            latest_data: Dict[str, Any] = {}
            action_summaries: List[str] = []

            for decision_index in range(max_decisions):
                try:
                    messages = build_messages(
                        metadata=metadata,
                        latest_data=latest_data,
                        action_summaries=action_summaries,
                        decision_index=decision_index,
                        max_decisions=max_decisions,
                    )
                    tool_call = request_action(
                        client=client,
                        scheduler=scheduler,
                        messages=messages,
                        stats=task_stats,
                        decision_index=decision_index,
                        max_decisions=max_decisions,
                    )
                except Exception as e:
                    print(f"[DEBUG] model request failed: {e}", flush=True)
                    break
                    
                task_stats.decisions += 1
                tool_name = tool_call["tool_name"]
                arguments = tool_call["arguments"]
                
                args_str = json.dumps(arguments, separators=(',', ':'))
                action_str = f"{tool_name}({args_str})"
                
                error = None
                reward = 0.0
                done = False

                try:
                    step_result = env.step(CallToolAction(tool_name=tool_name, arguments=arguments))
                    observation = step_result.observation
                    reward = step_result.reward or 0.0
                    done = step_result.done
                    
                    result_data = normalize_tool_result(extract_tool_result_from_observation(observation))
                    if result_data:
                        latest_data = apply_tool_context(tool_name, arguments, latest_data, result_data)
                    else:
                        latest_data = apply_tool_context(
                            tool_name,
                            arguments,
                            latest_data,
                            extract_observation_metadata(observation),
                        )
                    
                    # Update score
                    metadata_grade = extract_observation_metadata(observation).get("grade", {})
                    result_grade = latest_data.get("grade", {}) if isinstance(latest_data, dict) else {}
                    if result_grade:
                        final_score = result_grade.get("score") or 0.0
                    elif metadata_grade:
                        final_score = metadata_grade.get("score") or 0.0
                        
                except Exception as exc:
                    error = str(exc)
                    done = True

                rewards.append(reward)
                total_reward += reward
                steps_taken = decision_index + 1
                
                error_val = error if error else "null"
                error_val = error_val.replace('\n', ' ')
                done_val = str(done).lower()
                
                print(f"[STEP] step={steps_taken} action={action_str} reward={reward:.2f} done={done_val} error={error_val}", flush=True)

                action_summaries.append(action_summary(tool_name, arguments, latest_data, reward))

                if done:
                    break
            else:
                # auto-submit
                try:
                    action_str = "submit_model({})"
                    step_result = env.step(CallToolAction(tool_name="submit_model", arguments={}))
                    observation = step_result.observation
                    reward = step_result.reward or 0.0
                    total_reward += reward
                    done = step_result.done
                    
                    result_data = normalize_tool_result(extract_tool_result_from_observation(observation))
                    if result_data:
                        latest_data = apply_tool_context("submit_model", {}, latest_data, result_data)

                    metadata_grade = extract_observation_metadata(observation).get("grade", {})
                    result_grade = latest_data.get("grade", {}) if isinstance(latest_data, dict) else {}
                    if result_grade:
                        final_score = result_grade.get("score") or 0.0
                    elif metadata_grade:
                        final_score = metadata_grade.get("score") or 0.0
                        
                    steps_taken += 1
                    rewards.append(reward)
                    print(f"[STEP] step={steps_taken} action={action_str} reward={reward:.2f} done={str(done).lower()} error=null", flush=True)

                except Exception as exc:
                    pass

    except Exception as e:
        print(f"[DEBUG] task container failed: {e}", flush=True)

    finally:
        task_stats.elapsed_seconds = round(time.time() - task_started_at, 1)
        if final_score is None:
            final_score = 0.001
        final_score = min(max(float(final_score), 0.001), 0.999)
        success = final_score >= 0.1
        
        rewards_str = ",".join(f"{float(r):.2f}" for r in rewards) if rewards else "0.00"
        print(f"[END] success={str(success).lower()} steps={steps_taken} score={final_score:.3f} rewards={rewards_str}", flush=True)

    return {
        "task_id": task_id,
        "final_score": final_score,
        "total_reward": round(total_reward, 4),
        "llm_decisions": task_stats.decisions,
        "requests": task_stats.requests,
        "retries": task_stats.retries,
        "elapsed_seconds": task_stats.elapsed_seconds,
    }


def main() -> None:
    """Run the baseline across all tasks."""
    if not HF_TOKEN:
        print("ERROR: Set HF_TOKEN, OPENAI_API_KEY, OPENROUTER_API_KEY or GEMINI_API_KEY environment variable")
        sys.exit(1)

    print(f"API Base: {API_BASE_URL}", file=sys.stderr)
    print(f"Model: {MODEL_NAME}", file=sys.stderr)
    print(f"Environment: {ENV_URL}", file=sys.stderr)
    print(f"Tasks: {TASKS}", file=sys.stderr)
    print(f"RPM limit: {LLM_RPM_LIMIT} (min gap {MIN_REQUEST_GAP_SECONDS:.1f}s)", file=sys.stderr)
    print(f"Retries: {LLM_MAX_RETRIES}", file=sys.stderr)

    client = create_client()
    scheduler = RequestScheduler(
        min_gap_seconds=MIN_REQUEST_GAP_SECONDS,
        rpm_limit=LLM_RPM_LIMIT,
    )
    started_at = time.time()
    results = []

    try:
        for task_id in TASKS:
            results.append(run_task(client, scheduler, task_id))
    except InferenceError as exc:
        print(f"ERROR: {exc}")
        sys.exit(1)

    elapsed = round(time.time() - started_at, 1)
    total_requests = sum(result["requests"] for result in results)
    total_retries = sum(result["retries"] for result in results)

    print(f"\n{'=' * 60}", file=sys.stderr)
    print("BASELINE RESULTS", file=sys.stderr)
    print(f"{'=' * 60}", file=sys.stderr)
    for result in results:
        print(
            "  "
            f"{result['task_id']:20s} score={result['final_score']} "
            f"decisions={result['llm_decisions']} requests={result['requests']} "
            f"retries={result['retries']} elapsed={result['elapsed_seconds']}s",
            file=sys.stderr
        )
    print(f"  Total requests: {total_requests}", file=sys.stderr)
    print(f"  Total retries: {total_retries}", file=sys.stderr)
    print(f"  Total elapsed: {elapsed}s ({elapsed / 60:.1f} min)", file=sys.stderr)
    print(f"{'=' * 60}", file=sys.stderr)

    output = {
        "results": results,
        "total_requests": total_requests,
        "total_retries": total_retries,
        "total_time_seconds": elapsed,
        "model": MODEL_NAME,
        "rpm_limit": LLM_RPM_LIMIT,
        "min_request_gap_seconds": round(MIN_REQUEST_GAP_SECONDS, 1),
    }
    print(f"\n{json.dumps(output, indent=2, default=str)}", file=sys.stderr)


if __name__ == "__main__":
    main()