inference.py

"""
Inference script for the Code Review Environment.

Environment variables:
    API_BASE_URL  — LLM API endpoint (e.g. https://openrouter.ai/api/v1)
    MODEL_NAME    — Model identifier (e.g. openai/gpt-4o-mini)
    HF_TOKEN      — API key for the LLM provider
    ENV_URL       — Environment base URL (default: localhost:7860)

Usage:
    export API_BASE_URL=https://openrouter.ai/api/v1
    export MODEL_NAME=openai/gpt-4o-mini
    export HF_TOKEN=sk-...
    python inference.py
"""
from __future__ import annotations

import os
import sys
import json
import time

import httpx

API_BASE_URL: str = os.environ.get("API_BASE_URL", "").rstrip("/")
MODEL_NAME: str = os.environ.get("MODEL_NAME", "gpt-4o-mini")
HF_TOKEN: str = os.environ.get("HF_TOKEN", "")
ENV_URL: str = os.environ.get("ENV_URL", "http://localhost:7860").rstrip("/")

TASK_IDS = ["bug-detection", "security-audit", "comprehensive-review"]

SYSTEM_PROMPT = """\
You are an expert software engineer performing a thorough code review.

Your job is to identify bugs, security vulnerabilities, and performance issues in code.

For each issue you find, respond with a single JSON object:
  {"action_type": "flag_issue", "line_number": <int>, "filename": "<file>", "issue_type": "bug|security|performance|logic", "severity": "low|medium|high|critical", "description": "<explanation>", "fix_suggestion": "<fix>"}

When done, respond with:
  {"action_type": "submit_review"}

Rules:
- Respond with raw JSON only — no markdown fences, no extra text
- One action per response
- Be precise with line numbers (count from line 1)
- Only flag real issues, not style preferences
"""


def chat_completion(messages: list) -> str:
    try:
        from openai import OpenAI
    except ImportError:
        raise ImportError("pip install openai")

    kwargs = {"api_key": HF_TOKEN or "no-key"}
    if API_BASE_URL:
        kwargs["base_url"] = API_BASE_URL

    client = OpenAI(**kwargs)
    response = client.chat.completions.create(
        model=MODEL_NAME,
        messages=messages,
        temperature=0.0,
        max_tokens=400,
    )
    return response.choices[0].message.content.strip()


def parse_action(text: str) -> dict:
    text = text.strip()

    if "```" in text:
        parts = text.split("```")
        for part in parts:
            part = part.strip()
            if part.startswith("json"):
                part = part[4:].strip()
            if part.startswith("{") or part.startswith("["):
                text = part
                break

    decoder = json.JSONDecoder()
    for i, ch in enumerate(text):
        if ch in ("{", "["):
            try:
                obj, _ = decoder.raw_decode(text, i)
                if isinstance(obj, dict):
                    return obj
                if isinstance(obj, list):
                    for item in obj:
                        if isinstance(item, dict):
                            return item
            except json.JSONDecodeError:
                continue

    return {"action_type": "submit_review"}


def run_keyword_fallback(base_url: str, task_id: str) -> dict:
    """Fallback: use the built-in /baseline endpoint (no LLM needed)."""
    with httpx.Client(timeout=30) as client:
        resp = client.post(f"{base_url}/baseline")
        resp.raise_for_status()
        results = resp.json()
        score = results["baseline_scores"].get(task_id, {}).get("score", 0.0)
        return {"task_id": task_id, "score": score, "steps": 0, "method": "keyword_heuristic"}


def run_task(task_id: str, http_client: httpx.Client) -> dict:
    resp = http_client.post(f"{ENV_URL}/reset", json={"task_id": task_id}, timeout=30)
    resp.raise_for_status()
    obs = resp.json()

    code_display = "\n\n".join(
        f"=== {fname} ===\n{code}"
        for fname, code in obs.get("code_files", {}).items()
    )

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": (
                f"Task: {obs.get('task_description', '')}\n\n"
                f"{code_display}\n\n"
                f"Review this code carefully. Flag every issue you find. "
                f"You have {obs.get('max_steps', 20)} steps total."
            ),
        },
    ]

    done = False
    step_count = 0
    max_steps = obs.get("max_steps", 20)
    final_score = 0.0

    while not done and step_count < max_steps:
        action_text = chat_completion(messages)
        action = parse_action(action_text)

        try:
            step_resp = http_client.post(f"{ENV_URL}/step", json=action, timeout=30)
            step_resp.raise_for_status()
            obs = step_resp.json()
        except Exception as e:
            print(f"    Step error: {e}")
            break

        done = obs.get("done", False)
        step_count += 1
        final_score = obs.get("current_score", 0.0)
        reward = obs.get("reward")

        messages.append({"role": "assistant", "content": action_text})
        messages.append({
            "role": "user",
            "content": (
                f"Feedback: {obs.get('feedback', '')} "
                f"(step {step_count}/{max_steps}, score: {obs.get('current_score', 0.0):.3f})"
            ),
        })

        atype = action.get("action_type", "")
        print(f"    Step {step_count:2d}: {atype:20s} | reward={str(reward):8s} | score={obs.get('current_score', 0.0):.3f}")

        if atype == "submit_review":
            final_score = obs.get("reward", obs.get("current_score", 0.0)) or 0.0
            break

        time.sleep(0.3)

    return {
        "task_id": task_id,
        "score": float(final_score),
        "steps": step_count,
        "method": "llm",
    }


def main():
    use_llm = bool(HF_TOKEN and API_BASE_URL)

    print("Code Review Environment — Inference")
    print(f"  Model   : {MODEL_NAME}")
    print(f"  API URL : {API_BASE_URL or '(not set — using keyword heuristic)'}")
    print(f"  Env URL : {ENV_URL}")
    print(f"  Tasks   : {TASK_IDS}\n")

    try:
        with httpx.Client(timeout=10) as probe:
            health = probe.get(f"{ENV_URL}/health")
            health.raise_for_status()
            print(f"  Health: {health.json()}\n")
    except Exception as e:
        print(f"ERROR: Cannot reach environment at {ENV_URL}: {e}")
        sys.exit(1)

    results = {}

    if use_llm:
        with httpx.Client(timeout=60) as client:
            for task_id in TASK_IDS:
                print(f"Running task: {task_id}")
                result = run_task(task_id, client)
                results[task_id] = result
                print(f"  → score: {result['score']:.4f}  ({result['steps']} steps)\n")
    else:
        print("HF_TOKEN / API_BASE_URL not set — using built-in keyword heuristic baseline.\n")
        for task_id in TASK_IDS:
            print(f"Running task: {task_id}")
            result = run_keyword_fallback(ENV_URL, task_id)
            results[task_id] = result
            print(f"  → score: {result['score']:.4f}\n")

    print("=" * 50)
    print("INFERENCE RESULTS")
    print("=" * 50)
    for task_id, r in results.items():
        print(f"  {task_id:30s}  score={r['score']:.4f}")

    overall = sum(r["score"] for r in results.values()) / len(results)
    print(f"\n  Overall average: {overall:.4f}")
    print("=" * 50)

    return results


if __name__ == "__main__":
    main()