hanabhi/gridworld-env

1#!/usr/bin/env python32"""Solve a coding task with a hosted LLM via Hugging Face Inference.3 4This script mirrors ``textarena_wordle_inference.py`` but targets the Coding5environment. It launches the CodingEnv Docker image locally and asks an6OpenAI-compatible model served through Hugging Face's router to iteratively7produce Python code until the task is solved.8 9Prerequisites10-------------111. Build the Coding environment Docker image::12 13       docker build \14           -f envs/coding_env/server/Dockerfile \15           -t coding-env:latest .16 172. Set your Hugging Face token, or any other API key that is compatible with the OpenAI API:18 19       export HF_TOKEN=your_token_here20       export API_KEY=your_api_key_here21 223. Run the script::23 24       python examples/coding_env_inference.py25 26The script keeps sending execution feedback to the model until it prints27``Result: 338350`` or reaches the configured step limit.28"""29 30from __future__ import annotations31 32import os33import re34from typing import List, Tuple35 36from openai import OpenAI37 38from coding_env import CodeAction, CodingEnv39 40 41# ---------------------------------------------------------------------------42# Configuration43# ---------------------------------------------------------------------------44 45API_BASE_URL = "https://router.huggingface.co/v1"46API_KEY = os.getenv("API_KEY") or os.getenv("HF_TOKEN")47 48MODEL = "openai/gpt-oss-120b:novita"49MAX_STEPS = 550VERBOSE = True51 52CODING_TASK = (53    "Write Python code that prints the sum of squares of the integers from 1 "54    "to 100 inclusive. The final line must be exactly `Result: <value>` with "55    "the correct number substituted."56)57EXPECTED_SUBSTRING = "Result: 338350"58 59SYSTEM_PROMPT = (60    "You are an expert Python programmer. Respond with valid Python code that "61    "solves the user's task. Always wrap your final answer in a fenced code "62    "block starting with ```python. Provide a complete script that can be "63    "executed as-is, with no commentary outside the code block."64)65 66 67# ---------------------------------------------------------------------------68# Helpers69# ---------------------------------------------------------------------------70 71 72def extract_python_code(text: str) -> str:73    """Extract the first Python code block from the model output."""74 75    code_blocks = re.findall(76        r"```(?:python)?\s*(.*?)```",77        text,78        re.IGNORECASE | re.DOTALL,79    )80    if code_blocks:81        return code_blocks[0].strip()82    return text.strip()83 84 85def format_feedback(86    step: int,87    stdout: str,88    stderr: str,89    exit_code: int,90) -> str:91    """Generate feedback text describing the previous execution."""92 93    stdout_display = stdout if stdout.strip() else "<empty>"94    stderr_display = stderr if stderr.strip() else "<empty>"95    return (96        f"Execution feedback for step {step}:\n"97        f"exit_code={exit_code}\n"98        f"stdout:\n{stdout_display}\n"99        f"stderr:\n{stderr_display}\n"100        "If the task is not solved, return an improved Python script."101    )102 103 104def build_initial_prompt(task: str) -> str:105    """Construct the first user prompt for the coding task."""106 107    return (108        "You must write Python code to satisfy the following task. "109        "When executed, your script should behave exactly as described.\n\n"110        f"Task:\n{task}\n\n"111        "Reply with the full script in a single ```python code block."112    )113 114 115# ---------------------------------------------------------------------------116# Gameplay117# ---------------------------------------------------------------------------118 119 120def solve_coding_task(121    env: CodingEnv,122    client: OpenAI,123) -> Tuple[bool, List[str]]:124    """Iteratively ask the model for code until the task is solved."""125 126    history = [127        {"role": "system", "content": SYSTEM_PROMPT},128        {"role": "user", "content": build_initial_prompt(CODING_TASK)},129    ]130 131    obs = env.reset().observation132 133    transcripts: List[str] = []134 135    for step in range(1, MAX_STEPS + 1):136        response = client.chat.completions.create(137            model=MODEL,138            messages=history,139            max_tokens=2048,140            temperature=0.2,141        )142 143        assistant_message = response.choices[0].message.content.strip()144        history.append({"role": "assistant", "content": assistant_message})145 146        code = extract_python_code(assistant_message)147 148        if VERBOSE:149            print(f"\n🛠️  Step {step}: executing model-produced code")150            print(code)151 152        result = env.step(CodeAction(code=code))153        obs = result.observation154 155        transcripts.append(156            (157                f"Step {step} | exit_code={obs.exit_code}\nstdout:\n{obs.stdout}\nstderr:\n{obs.stderr}\n"158            )159        )160 161        if VERBOSE:162            print("   ▶ exit_code:", obs.exit_code)163            if obs.stdout:164                print("   ▶ stdout:\n" + obs.stdout)165            if obs.stderr:166                print("   ▶ stderr:\n" + obs.stderr)167 168        solved = obs.exit_code == 0 and EXPECTED_SUBSTRING in obs.stdout169        if solved:170            return True, transcripts171 172        history.append(173            {174                "role": "user",175                "content": format_feedback(176                    step,177                    obs.stdout,178                    obs.stderr,179                    obs.exit_code,180                ),181            }182        )183 184        # Keep conversation history compact to avoid exceeding context limits185        if len(history) > 20:186            history = [history[0]] + history[-19:]187 188    return False, transcripts189 190 191# ---------------------------------------------------------------------------192# Entrypoint193# ---------------------------------------------------------------------------194 195 196def main() -> None:197    if not API_KEY:198        raise SystemExit(199            "HF_TOKEN (or API_KEY) must be set to query the model."200        )201 202    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)203 204    env = CodingEnv.from_docker_image(205        "coding-env:latest",206        ports={8000: 8000},207    )208 209    try:210        success, transcripts = solve_coding_task(env, client)211    finally:212        env.close()213 214    print(215        "\n✅ Session complete"216        if success217        else "\n⚠️ Session finished without solving the task"218    )219    print("--- Execution transcripts ---")220    for entry in transcripts:221        print(entry)222 223 224if __name__ == "__main__":225    main()226