inference.py

"""Baseline inference script for CI/CD Debug Environment.

Uses OpenAI-compatible client to call Llama 3.1 70B via HuggingFace router.
Required by OpenEnv specification.

Usage:
    export API_BASE_URL=https://router.huggingface.co/v1
    export MODEL_NAME=meta-llama/Llama-3.1-70B-Instruct
    export HF_TOKEN=your_token_here
    python inference.py
"""


import json
import os
import re
import sys
import time
from typing import Any, Dict, List, Optional

import requests
from openai import OpenAI


API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Llama-3.1-70B-Instruct")
HF_TOKEN = os.getenv("HF_TOKEN")
ENV_URL = os.getenv("ENV_URL", "http://localhost:8000")
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
MAX_STEPS = 8  # leave 2 steps buffer before env hard-limit of 10

SYSTEM_PROMPT = """You are an expert DevOps engineer debugging CI/CD pipelines.
You will receive broken Dockerfile and/or GitHub Actions workflow files along with error messages.

Your job is to:
1. Analyze the error message carefully
2. Identify the root cause in the configuration files
3. Provide a precise fix

When you identify a fix, respond with a JSON object in this exact format:
{
  "reasoning": "Brief explanation of the bug and fix",
  "edits": [
    {
      "file_path": "path/to/file",
      "old_content": "exact broken line or block",
      "new_content": "corrected line or block"
    }
  ]
}

If you believe all issues are fixed and want to submit, respond with:
{"action": "submit"}

If you need a hint, respond with:
{"action": "hint"}

Rules:
- Match old_content EXACTLY as it appears in the file (whitespace matters)
- Fix one issue at a time for precision
- Focus on the error message — it tells you exactly what's wrong
- Common issues: typos, wrong syntax, missing fields, wrong secret references
- For GitHub Actions: check secret syntax (${{ }} not ${ }), env blocks, permissions
- For Dockerfiles: check instruction syntax, file paths, base image tags
- Always respond with valid JSON only, no markdown fences"""


def create_client() -> OpenAI:
    """Create OpenAI-compatible client for HuggingFace router."""
    return OpenAI(
        base_url=API_BASE_URL,
        api_key=HF_TOKEN,
    )


def env_request(method: str, endpoint: str, json_data: Optional[Dict] = None) -> Dict[str, Any]:
    """Make a request to the environment server."""
    url = f"{ENV_URL}{endpoint}"
    if method == "GET":
        resp = requests.get(url, timeout=30)
    else:
        resp = requests.post(url, json=json_data or {}, timeout=30)
    resp.raise_for_status()
    return resp.json()


def format_observation(obs: Dict[str, Any]) -> str:
    """Format observation into a prompt for the LLM."""
    parts = []
    parts.append(f"Task: {obs.get('task_description', 'Unknown')}")
    parts.append(f"Difficulty: {obs.get('difficulty', 'unknown')}")
    parts.append(f"Step: {obs.get('step_number', 0)}/{obs.get('max_steps', 10)}")
    parts.append(f"Issues fixed: {obs.get('issues_fixed', 0)}/{obs.get('total_issues', '?')}")

    error = obs.get("error", {})
    parts.append(f"\n--- ERROR ---")
    parts.append(f"Phase: {error.get('phase', 'unknown')}")
    parts.append(f"Message: {error.get('error_message', 'No error')}")
    if error.get("failed_step"):
        parts.append(f"Failed step: {error['failed_step']}")
    if error.get("line_hint"):
        parts.append(f"Line hint: {error['line_hint']}")

    parts.append(f"\n--- FILES ---")
    for f in obs.get("files", []):
        parts.append(f"\n=== {f['path']} ({f.get('file_type', 'unknown')}) ===")
        content = f.get("content", "")
        lines = content.split("\n")
        for i, line in enumerate(lines, 1):
            parts.append(f"{i:3d} | {line}")

    if obs.get("available_secrets"):
        parts.append(f"\n--- AVAILABLE SECRETS ---")
        parts.append(", ".join(obs["available_secrets"]))

    if obs.get("last_action_feedback"):
        parts.append(f"\n--- LAST ACTION FEEDBACK ---")
        parts.append(obs["last_action_feedback"])

    return "\n".join(parts)


def parse_llm_response(text: str) -> Dict[str, Any]:
    """Parse LLM response into an action dict."""
    text = text.strip()

    # Strip markdown code fences if present
    if text.startswith("```"):
        lines = text.split("\n")
        lines = [l for l in lines if not l.strip().startswith("```")]
        text = "\n".join(lines).strip()

    # Try to find JSON in the response
    json_match = re.search(r'\{[\s\S]*\}', text)
    if json_match:
        try:
            return json.loads(json_match.group())
        except json.JSONDecodeError:
            pass

    # Fallback: treat as submit
    return {"action": "submit"}


def build_action(parsed: Dict[str, Any]) -> Dict[str, Any]:
    """Convert parsed LLM response to environment action format."""
    if parsed.get("action") == "submit":
        return {"action_type": "submit"}
    if parsed.get("action") == "hint":
        return {"action_type": "request_hint"}

    edits = parsed.get("edits", [])
    if not edits:
        return {"action_type": "submit"}

    return {
        "action_type": "edit_file",
        "edits": [
            {
                "file_path": e.get("file_path", ""),
                "old_content": e.get("old_content", ""),
                "new_content": e.get("new_content", ""),
            }
            for e in edits
        ],
    }


def run_episode(client: OpenAI, task_id: Optional[str] = None, scenario_id: Optional[str] = None) -> Dict[str, Any]:
    """Run a single episode: reset, loop (observe -> LLM -> act), grade."""
    reset_payload: Dict[str, Any] = {}
    if task_id:
        reset_payload["task_id"] = task_id
    if scenario_id:
        reset_payload["scenario_id"] = scenario_id

    reset_resp = env_request("POST", "/reset", reset_payload)
    obs = reset_resp["observation"]
    info = reset_resp.get("info", {})

    actual_task_id = info.get("task_id", task_id or "unknown")
    actual_scenario_id = info.get("scenario_id", scenario_id or "unknown")

    print(f"[START] task_id={actual_task_id} scenario_id={actual_scenario_id}")

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    trajectory = []
    total_steps = 0

    for step_num in range(MAX_STEPS):
        user_msg = format_observation(obs)
        messages.append({"role": "user", "content": user_msg})

        try:
            completion = client.chat.completions.create(
                model=MODEL_NAME,
                messages=messages,
                temperature=0.1,
                max_tokens=1024,
            )
            llm_text = completion.choices[0].message.content or '{"action": "submit"}'
        except Exception as e:
            print(f"[STEP] step={step_num + 1} action=error reward=0.00 done=false issues_fixed=0 issues_total=0 error={e}")
            llm_text = '{"action": "submit"}'

        messages.append({"role": "assistant", "content": llm_text})

        parsed = parse_llm_response(llm_text)
        action = build_action(parsed)

        step_resp = env_request("POST", "/step", {"action": action})
        obs = step_resp["observation"]
        reward = step_resp.get("reward", 0.0)
        done = step_resp.get("done", False)
        step_info = step_resp.get("info", {})
        total_steps = step_num + 1

        issues_fixed = step_info.get("issues_fixed", 0)
        issues_total = step_info.get("issues_total", 0)

        print(f"[STEP] step={total_steps} action={action['action_type']} reward={reward:.2f} done={str(done).lower()} issues_fixed={issues_fixed} issues_total={issues_total}")

        trajectory.append({
            "step": total_steps,
            "action": action,
            "reward": reward,
            "done": done,
            "info": step_info,
        })

        if done:
            break

    # Grade the trajectory
    grade_resp = env_request("POST", "/grader", {
        "task_id": actual_task_id,
        "trajectory": trajectory,
    })
    result = grade_resp.get("result", {})
    score = result.get("score", 0.0)

    print(f"[END] task_id={actual_task_id} scenario_id={actual_scenario_id} score={score:.3f} steps={total_steps}")
    return result


def run_all_tasks(client: OpenAI) -> Dict[str, float]:
    """Run baseline on all tasks and report scores."""
    tasks_resp = env_request("GET", "/tasks")
    tasks = tasks_resp.get("tasks", [])

    scores: Dict[str, List[float]] = {}

    for task in tasks:
        task_id = task["id"]
        print(f"\n{'='*60}")
        print(f"Task: {task['name']} ({task['difficulty']})")
        print(f"{'='*60}")

        task_scores = []
        # Run one episode per task for baseline
        result = run_episode(client, task_id=task_id)
        task_scores.append(result.get("score", 0.0))
        scores[task_id] = task_scores

    # Summary
    print(f"\n{'='*60}")
    print("BASELINE RESULTS SUMMARY")
    print(f"{'='*60}")
    avg_scores = {}
    for task_id, task_scores in scores.items():
        avg = sum(task_scores) / len(task_scores) if task_scores else 0.0
        avg_scores[task_id] = avg
        print(f"  {task_id:40s} {avg:.3f}")

    overall = sum(avg_scores.values()) / len(avg_scores) if avg_scores else 0.0
    print(f"  {'OVERALL':40s} {overall:.3f}")

    return avg_scores


def main():
    """Entry point for baseline inference."""
    print("CI/CD Debug Environment - Baseline Inference")
    print(f"API: {API_BASE_URL}")
    print(f"Model: {MODEL_NAME}")
    print(f"Environment: {ENV_URL}")

    if not HF_TOKEN:
        print("\nWARNING: HF_TOKEN not set. Set it via: export HF_TOKEN=your_token_here")
        print("Continuing anyway (will fail if auth is required)...\n")

    # Verify environment is running
    try:
        health = env_request("GET", "/")
        print(f"Environment status: {health.get('status', 'unknown')}\n")
    except Exception as e:
        print(f"\nERROR: Cannot connect to environment at {ENV_URL}")
        print(f"  {e}")
        print("\nStart the server first:")
        print("  python -m uvicorn server.main:app --host 0.0.0.0 --port 8000")
        sys.exit(1)

    client = create_client()

    # If a specific task is requested via CLI arg
    if len(sys.argv) > 1:
        task_id = sys.argv[1]
        scenario_id = sys.argv[2] if len(sys.argv) > 2 else None
        run_episode(client, task_id=task_id, scenario_id=scenario_id)
    else:
        run_all_tasks(client)


if __name__ == "__main__":
    main()