Debdeep123/capability_forge_codeops

1import os2import asyncio3import json4from typing import List, Dict, Any5from openai import AsyncOpenAI6from client import CapabilityForgeCodeopsEnv7from models import Action8 9# Environment variables as required by the hackathon10API_BASE_URL = os.getenv("API_BASE_URL", "https://api-inference.huggingface.co/v1/")11API_KEY = os.getenv("API_KEY", os.getenv("HF_TOKEN", ""))12MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Meta-Llama-3-8B-Instruct")13IMAGE_NAME = "capability_forge_codeops-env:latest"14TASK_NAME = "Code Bug Fixing"15BENCHMARK = "CapabilityForge CodeOps"16MAX_STEPS = 517SUCCESS_SCORE_THRESHOLD = 0.918 19def log_start(task: str, env: str, model: str):20    print(f"[START] Task: {task} | Env: {env} | Model: {model}", flush=True)21 22def log_step(step: int, action: str, reward: float, done: bool, error: Any = None):23    err_str = f" | Error: {error}" if error else ""24    # Truncate action for logging if it's too long25    action_log = action[:100] + "..." if len(action) > 100 else action26    print(f"[STEP] {step} | Action: {action_log!r} | Reward: {reward:+.2f} | Done: {done}{err_str}", flush=True)27 28def log_end(success: bool, steps: int, score: float, rewards: List[float]):29    print(f"[END] Success: {success} | Steps: {steps} | Score: {score:.2f} | Rewards: {rewards}", flush=True)30 31async def get_model_message(client: AsyncOpenAI, obs, history: List[str]) -> Action:32    system_prompt = (33        "You are an expert Python developer fixing a bug. "34        "You must output valid JSON exactly matching this schema: "35        '{"reasoning": "your step-by-step reasoning", "corrected_code": "the corrected python code without markdown ticks"}'36    )37    38    prompt = f"{obs.task_prompt}\n\nBuggy Code:\n{obs.buggy_code}\n"39    if obs.execution_feedback:40        prompt += f"\nPrevious Error:\n{obs.execution_feedback}\n"41    if obs.hint:42        prompt += f"\nHint:\n{obs.hint}\n"43        44    messages = [{"role": "system", "content": system_prompt}]45    for h in history:46        messages.append({"role": "user", "content": "Previous interaction omitted for brevity."})47    messages.append({"role": "user", "content": prompt})48    49    try:50        response = await client.chat.completions.create(51            model=MODEL_NAME,52            messages=messages,53            response_format={"type": "json_object"},54            temperature=0.255        )56        content = response.choices[0].message.content57        data = json.loads(content)58        return Action(59            reasoning=data.get("reasoning", ""),60            corrected_code=data.get("corrected_code", "")61        )62    except Exception as exc:63        print(f"[DEBUG] Model request failed or failed to parse JSON: {exc}", flush=True)64        return Action(reasoning="Failed to parse model output.", corrected_code=obs.buggy_code)65 66async def main() -> None:67    client = AsyncOpenAI(base_url=API_BASE_URL, api_key=API_KEY)68    69    # Run against a local/remote API instance. 70    # For local testing, we assume the server is running on localhost:8000.71    # Note: the hackathon template used `from_docker_image`, but since openenv isn't fully set up with docker right now,72    # we connect via HTTP. If from_docker_image is required:73    # env = CapabilityForgeCodeopsEnv.from_docker_image(IMAGE_NAME)74    # Using HTTP client:75    env = CapabilityForgeCodeopsEnv(base_url="http://localhost:8000")76    77    history: List[str] = []78    rewards: List[float] = []79    steps_taken = 080    score = 0.081    success = False82 83    log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)84 85    try:86        result = env.reset() # env.reset is sync in client if using standard httpx, or async if openenv client is async. Wait, EnvClient uses requests natively.87        obs = result.observation88        last_reward = 0.089 90        for step in range(1, MAX_STEPS + 1):91            if result.done:92                break93 94            action = await get_model_message(client, obs, history)95            96            try:97                result = env.step(action)98                obs = result.observation99                reward = result.reward or 0.0100                done = result.done101                error = None102            except Exception as e:103                reward = -0.1104                done = True105                error = str(e)106                obs = obs # keep previous107 108            rewards.append(reward)109            steps_taken = step110            last_reward = reward111 112            log_step(step=step, action=action.model_dump_json(), reward=reward, done=done, error=error)113            history.append(f"Step {step}: ... -> reward {reward:+.2f}")114 115            if done:116                break117 118        # Max total reward for our env is ~1.3 (1.0 exec + 0.2 first attempt + 0.1 format)119        MAX_TOTAL_REWARD = 1.3120        score = sum(rewards) / MAX_TOTAL_REWARD if MAX_TOTAL_REWARD > 0 else 0.0121        score = min(max(score, 0.0), 1.0)  # clamp to [0, 1]122        success = score >= SUCCESS_SCORE_THRESHOLD123 124    finally:125        try:126            env.close()127        except Exception as e:128            print(f"[DEBUG] env.close() error: {e}", flush=True)129        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)130 131 132if __name__ == "__main__":133    asyncio.run(main())134