inference.py

#!/usr/bin/env python3
"""Baseline inference script for autotest-env following OpenEnv spec.

This script runs an LLM agent through all 3 tasks in sequence,
using the OpenAI-compatible API to generate pytest test code.

Usage:
    python inference.py [--seed SEED] [--task TASK]

Arguments:
    --seed SEED    Random seed for reproducible task generation (default: 42)
    --task TASK    Run specific task only: unit_test_writer, coverage_audit, regression_audit

Environment variables (required by OpenEnv spec):
    HF_TOKEN: HuggingFace token (primary API key)
    API_BASE_URL: Base URL for OpenAI-compatible API
    MODEL_NAME: Model to use for inference

Optional environment variables:
    OPENAI_API_KEY: OpenAI API key (fallback if HF_TOKEN not set)
    ENV_URL: URL of the autotest-env server (default: http://localhost:7860)
"""

import argparse
import json
import os
import re
import sys
import time
import traceback
import urllib.request
import urllib.error
from typing import Optional

from openai import OpenAI

# =============================================================================
# Environment Variables (required by OpenEnv spec)
# =============================================================================
# These MUST be top-level module variables for the validator to detect them

HF_TOKEN = os.getenv("HF_TOKEN")  # Required - no default
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")

# Optional environment variables
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")  # Fallback if HF_TOKEN not set
ENV_URL = os.getenv("ENV_URL", "http://localhost:7860")
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")  # Optional for local testing

# Tasks to run in sequence
TASKS = ["unit_test_writer", "coverage_audit", "regression_audit"]

# Environment name for logging
ENV_NAME = "rl-testing-env"

# Maximum total reward for score normalization (max_steps * max_reward_per_step)
# With 10 max steps and 1.0 max reward per step
MAX_TOTAL_REWARD = 10.0


def get_api_key() -> str:
    """Get API key, preferring HF_TOKEN over OPENAI_API_KEY (per OpenEnv spec)."""
    if HF_TOKEN:
        return HF_TOKEN
    if OPENAI_API_KEY:
        return OPENAI_API_KEY
    raise ValueError(
        "No API key found. Set HF_TOKEN or OPENAI_API_KEY environment variable."
    )


def create_client() -> OpenAI:
    """Create OpenAI client with configured base URL and token."""
    return OpenAI(
        base_url=API_BASE_URL,
        api_key=get_api_key(),
        timeout=60.0,  # 60 second timeout for API calls
    )


def _post_json(url: str, data: dict, timeout: int = 30) -> dict:
    """Make a POST request with JSON body using urllib."""
    json_data = json.dumps(data).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=json_data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as response:
            response_data = response.read().decode("utf-8")
            try:
                return json.loads(response_data)
            except json.JSONDecodeError as e:
                raise RuntimeError(f"JSON parse error: {e}")
    except urllib.error.HTTPError as e:
        raise RuntimeError(f"HTTP {e.code}: {e.reason}")
    except urllib.error.URLError as e:
        raise RuntimeError(f"URL Error: {e.reason}")
    except TimeoutError:
        raise RuntimeError(f"Request timeout after {timeout}s")
    except Exception as e:
        raise RuntimeError(f"Network error: {str(e)}")


def check_env_health(max_retries: int = 3, retry_delay: float = 2.0) -> bool:
    """Check if the environment server is reachable and healthy."""
    for attempt in range(max_retries):
        try:
            req = urllib.request.Request(f"{ENV_URL}/health", method="GET")
            with urllib.request.urlopen(req, timeout=10) as response:
                if response.status == 200:
                    return True
        except Exception as e:
            if attempt < max_retries - 1:
                print(f"# Health check attempt {attempt + 1} failed: {e}, retrying...", file=sys.stderr)
                time.sleep(retry_delay)
            else:
                print(f"# Health check failed after {max_retries} attempts: {e}", file=sys.stderr)
    return False


def reset_env(task_id: str, seed: int = 42) -> dict:
    """Call /reset endpoint to start a new episode."""
    return _post_json(
        f"{ENV_URL}/reset",
        {"task_id": task_id, "seed": seed},
        timeout=30,
    )


def step_env(action_type: str, test_code: str, notes: str = "") -> dict:
    """Call /step endpoint to execute an action."""
    return _post_json(
        f"{ENV_URL}/step",
        {
            "action_type": action_type,
            "test_code": test_code,
            "notes": notes,
        },
        timeout=60,
    )


def build_prompt(observation: dict, task_id: str) -> str:
    """Build the LLM prompt from the current observation."""
    code_under_test = observation.get("code_under_test", "")
    task_description = observation.get("task_description", "")
    previous_test_results = observation.get("previous_test_results", "")
    bugs_found_so_far = observation.get("bugs_found_so_far", 0)
    coverage_pct = observation.get("coverage_pct", 0.0)
    step_number = observation.get("step_number", 0)
    hint = observation.get("hint", "")

    prompt = f"""You are an expert software tester. Your task is to write pytest tests for the following code.

## Task
{task_description}

## Code Under Test
```python
{code_under_test}
```

## Current Status
- Step: {step_number}
- Bugs found so far: {bugs_found_so_far}
- Coverage: {coverage_pct:.1f}%
"""

    if previous_test_results:
        prompt += f"""
## Previous Test Results
```
{previous_test_results}
```
"""

    if hint:
        prompt += f"""
## Hint
{hint}
"""

    prompt += """
## Instructions
Write a complete pytest test file that thoroughly tests the code above.
- Include all necessary imports
- Write multiple test functions covering different cases
- Test edge cases and boundary conditions
- Use descriptive test names

Respond with ONLY a Python code block containing your pytest tests.
```python
# Your test code here
```
"""

    return prompt


def extract_code_block(response_text: str) -> str:
    """Extract Python code block from LLM response."""
    # Try to find ```python ... ``` block
    pattern = r"```python\s*(.*?)\s*```"
    matches = re.findall(pattern, response_text, re.DOTALL)
    if matches:
        return matches[-1].strip()

    # Try to find ``` ... ``` block
    pattern = r"```\s*(.*?)\s*```"
    matches = re.findall(pattern, response_text, re.DOTALL)
    if matches:
        return matches[-1].strip()

    # If no code block, try to extract anything that looks like Python
    lines = response_text.strip().split("\n")
    code_lines = []
    in_code = False
    for line in lines:
        if line.strip().startswith(("import ", "from ", "def test_", "class Test")):
            in_code = True
        if in_code:
            code_lines.append(line)

    if code_lines:
        return "\n".join(code_lines)

    # Last resort: return the whole response
    return response_text.strip()


def truncate(s: str, max_len: int = 80) -> str:
    """Truncate string to max length, replacing newlines with spaces."""
    s = s.replace("\n", " ").replace("\r", " ")
    if len(s) > max_len:
        return s[:max_len]
    return s


def format_rewards(rewards: list[float]) -> str:
    """Format rewards list as comma-separated string."""
    return ",".join(f"{r:.2f}" for r in rewards)


def run_task(
    client: OpenAI,
    task_id: str,
    seed: int = 42,
    task_timeout: float = 300.0,
    max_steps: int = 10,
) -> tuple[bool, int, float, list[float]]:
    """
    Run a single task to completion.

    Args:
        client: OpenAI client instance
        task_id: Task identifier
        seed: Random seed for reproducible task generation
        task_timeout: Maximum seconds allowed for this task (default: 5 minutes)
        max_steps: Maximum steps per task (default: 10)

    Returns:
        Tuple of (success, steps, final_score, rewards_list)
    """
    rewards = []
    step_count = 0
    total_reward = 0.0
    error_message: Optional[str] = None
    task_start = time.time()

    try:
        # Reset environment for this task with specified seed
        result = reset_env(task_id, seed=seed)
        observation = result.get("observation", {})
        done = result.get("done", False)

        while not done:
            # Check per-task timeout
            if time.time() - task_start > task_timeout:
                error_message = "task_timeout"
                print(
                    f"[STEP] step={step_count} action=null "
                    f"reward=0.00 done=true error=\"{error_message}\""
                )
                break

            step_count += 1
            error_message = None

            try:
                # Build prompt for LLM
                prompt = build_prompt(observation, task_id)

                # Call LLM with timeout
                completion = client.chat.completions.create(
                    model=MODEL_NAME,
                    messages=[
                        {
                            "role": "system",
                            "content": "You are an expert Python test engineer. Write pytest tests to find bugs in code.",
                        },
                        {"role": "user", "content": prompt},
                    ],
                    temperature=0.7,
                    max_tokens=2048,
                    timeout=60.0,  # Per-request timeout
                )

                response_text = completion.choices[0].message.content or ""

                # Extract code from response
                test_code = extract_code_block(response_text)

                # Submit tests to environment
                result = step_env("submit_tests", test_code)

                observation = result.get("observation", {})
                reward = result.get("reward", 0.0)
                done = result.get("done", False)

                rewards.append(reward)
                total_reward += reward

            except Exception as e:
                error_message = str(e)[:100]
                test_code = ""
                reward = 0.0
                # Try to continue despite error
                try:
                    result = step_env("submit_tests", "# Error occurred")
                    observation = result.get("observation", {})
                    done = result.get("done", False)
                except Exception:
                    done = True

            # Print step output
            action_str = truncate(test_code) if test_code else "null"
            error_str = f'"{error_message}"' if error_message else "null"
            print(
                f"[STEP] step={step_count} action={action_str} "
                f"reward={reward:.2f} done={str(done).lower()} error={error_str}"
            )

            # Safety limit to prevent infinite loops (reduced from 100)
            if step_count >= max_steps:
                done = True

    except Exception as e:
        error_message = str(e)[:100]
        print(
            f"[STEP] step={step_count} action=null "
            f"reward=0.00 done=true error=\"{error_message}\""
        )

    # Calculate final score as normalized cumulative reward (per OpenEnv spec)
    # score = sum(rewards) / MAX_TOTAL_REWARD, clamped to (0, 1) exclusive
    final_score = sum(rewards) / MAX_TOTAL_REWARD if MAX_TOTAL_REWARD > 0 else 0.001
    # Clamp to strictly between 0 and 1 (not 0.0 or 1.0)
    final_score = min(max(final_score, 0.001), 0.999)

    # Success if score >= 0.5
    success = final_score >= 0.5

    return success, step_count, final_score, rewards


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Run baseline inference agent against autotest-env",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    python inference.py                           # Run all tasks with seed=42
    python inference.py --seed 123                # Run all tasks with seed=123
    python inference.py --task unit_test_writer   # Run only easy task
    python inference.py --seed 42 --task coverage_audit  # Specific task and seed

Environment variables (per OpenEnv spec):
    HF_TOKEN          HuggingFace token (primary, required)
    API_BASE_URL      API endpoint URL (required)
    MODEL_NAME        Model to use (required)
    OPENAI_API_KEY    OpenAI API key (fallback if HF_TOKEN not set)
    ENV_URL           AutoTest-Env server URL (default: http://localhost:7860)
        """
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for reproducible task generation (default: 42)"
    )
    parser.add_argument(
        "--task",
        type=str,
        choices=TASKS,
        default=None,
        help="Run specific task only (default: run all tasks)"
    )
    # Suppress argparse help output from going to stdout
    return parser.parse_args()


def main():
    """Main entry point - run all tasks in sequence."""
    try:
        args = parse_args()
    except Exception as e:
        print(f"[START] task=init env={ENV_NAME} model={MODEL_NAME}")
        print(f"[STEP] step=0 action=null reward=0.00 done=true error=\"arg_parse_error: {str(e)[:50]}\"")
        print(f"[END] success=false steps=0 score=0.001 rewards=")
        return []

    start_time = time.time()
    # Total budget: 25 minutes to stay well under 30 minute limit
    total_timeout_seconds = 25 * 60
    # Per-task timeout: ~7 minutes each (allows some buffer for 3 tasks)
    task_timeout_seconds = 7 * 60
    # Max steps per task to limit LLM calls
    max_steps_per_task = 10

    # Determine which tasks to run
    tasks_to_run = [args.task] if args.task else TASKS
    all_results = []

    # Debug output goes to stderr (per OpenEnv spec - only [START]/[STEP]/[END] on stdout)
    print(f"# Baseline Inference: seed={args.seed}, model={MODEL_NAME}", file=sys.stderr)
    print(f"# Environment: {ENV_URL}", file=sys.stderr)

    # Check environment health before starting
    print("# Checking environment health...", file=sys.stderr)
    if not check_env_health():
        print("# ERROR: Environment not reachable", file=sys.stderr)
        for task_id in tasks_to_run:
            print(f"[START] task={task_id} env={ENV_NAME} model={MODEL_NAME}")
            print(f"[STEP] step=0 action=null reward=0.00 done=true error=\"env_not_reachable\"")
            print(f"[END] success=false steps=0 score=0.001 rewards=")
            all_results.append((task_id, False, 0.0))
        return all_results
    print("# Environment health check passed", file=sys.stderr)

    # Create OpenAI client
    try:
        client = create_client()
    except Exception as e:
        print(f"# ERROR: Failed to create client: {e}", file=sys.stderr)
        for task_id in tasks_to_run:
            print(f"[START] task={task_id} env={ENV_NAME} model={MODEL_NAME}")
            print(f"[STEP] step=0 action=null reward=0.00 done=true error=\"client_init_error: {str(e)[:50]}\"")
            print(f"[END] success=false steps=0 score=0.001 rewards=")
            all_results.append((task_id, False, 0.0))
        return all_results

    for task_id in tasks_to_run:
        # Check global timeout
        elapsed = time.time() - start_time
        remaining = total_timeout_seconds - elapsed
        if remaining <= 60:  # Less than 1 minute left
            print(f"[START] task={task_id} env={ENV_NAME} model={MODEL_NAME}")
            print("[STEP] step=0 action=null reward=0.00 done=true error=\"global_timeout\"")
            print(f"[END] success=false steps=0 score=0.001 rewards=")
            all_results.append((task_id, False, 0.0))
            continue

        # Compute effective task timeout (min of task limit and remaining time)
        effective_timeout = min(task_timeout_seconds, remaining - 30)

        # Print start marker
        print(f"[START] task={task_id} env={ENV_NAME} model={MODEL_NAME}")

        try:
            success, steps, score, rewards = run_task(
                client,
                task_id,
                seed=args.seed,
                task_timeout=effective_timeout,
                max_steps=max_steps_per_task,
            )

            # Print end marker
            rewards_str = format_rewards(rewards)
            print(
                f"[END] success={str(success).lower()} steps={steps} "
                f"score={score:.3f} rewards={rewards_str}"
            )

            all_results.append((task_id, success, score))

        except Exception as e:
            # Always print [END] even on crash
            error_msg = str(e)[:100]
            print(f"[END] success=false steps=0 score=0.001 rewards=")
            all_results.append((task_id, False, 0.0))
            # Print error to stderr for debugging
            print(f"Error in task {task_id}: {error_msg}", file=sys.stderr)
            traceback.print_exc(file=sys.stderr)

    # Print summary to stderr (per OpenEnv spec - only [START]/[STEP]/[END] on stdout)
    print("", file=sys.stderr)
    print("# Summary", file=sys.stderr)
    for task_id, success, score in all_results:
        status = "PASS" if success else "FAIL"
        print(f"#   {task_id}: {status} (score={score:.3f})", file=sys.stderr)

    return all_results


if __name__ == "__main__":
    main()