Spaces:
Sleeping
Sleeping
File size: 5,082 Bytes
ee933ab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | """V3 FlakeForge Client — OpenEnv client for the unified agent architecture.
Changes from V2:
- No judge calls
- No hypothesis gating
- Unified observe → generate → step loop
"""
from __future__ import annotations
import asyncio
import json
import os
from typing import Any, Dict, Optional
try:
from models import FlakeForgeAction, FlakeForgeObservation
from agent.unified_agent import (
UnifiedFlakeForgeAgent,
extract_think,
extract_patch,
extract_category_from_think,
extract_confidence_from_think,
)
except ImportError:
from .models import FlakeForgeAction, FlakeForgeObservation
from .agent.unified_agent import (
UnifiedFlakeForgeAgent,
extract_think,
extract_patch,
extract_category_from_think,
extract_confidence_from_think,
)
try:
from utils.logger import get_logger
except ImportError:
try:
from .utils.logger import get_logger
except ImportError:
import logging
get_logger = lambda n, **kw: logging.getLogger(n)
logger = get_logger(__name__)
class FlakeForgeClient:
"""V3 client for communicating with the FlakeForge environment.
Wraps the unified agent and environment interaction into a
simple interface for both inference and training.
"""
def __init__(
self,
agent: Optional[UnifiedFlakeForgeAgent] = None,
env_url: Optional[str] = None,
) -> None:
self.agent = agent
self.env_url = env_url or os.environ.get("ENV_BASE_URL", "http://localhost:8080")
def generate_action(self, observation: FlakeForgeObservation) -> FlakeForgeAction:
"""Generate a unified action from an observation."""
if self.agent is None:
raise RuntimeError("No agent configured. Pass agent to constructor.")
return self.agent.generate(observation)
def parse_raw_response(self, raw_response: str) -> FlakeForgeAction:
"""Parse a raw model response into a FlakeForgeAction."""
think = extract_think(raw_response)
patch = extract_patch(raw_response)
return FlakeForgeAction(
raw_response=raw_response,
think_text=think,
patch_text=patch,
predicted_category=extract_category_from_think(think),
predicted_confidence=extract_confidence_from_think(think),
)
async def run_episode_remote(
self,
test_identifier: str,
repo_path: str,
max_steps: int = 8,
) -> Dict[str, Any]:
"""Run an episode against a remote FlakeForge environment server."""
try:
import httpx
except ImportError:
raise ImportError("Remote client requires: pip install httpx")
async with httpx.AsyncClient(
base_url=self.env_url, timeout=120.0
) as client:
# Reset
reset_response = await client.post(
"/reset",
json={"test_identifier": test_identifier, "repo_path": repo_path},
)
reset_data = reset_response.json()
observation = FlakeForgeObservation(**reset_data["observation"])
trajectory = []
total_reward = 0.0
for step in range(max_steps):
if reset_data.get("done", False) and step > 0:
break
# Generate action
action = self.generate_action(observation)
# Send to environment
step_response = await client.post(
"/step",
json={
"raw_response": action.raw_response,
"think_text": action.think_text,
"patch_text": action.patch_text,
"predicted_category": action.predicted_category,
"predicted_confidence": action.predicted_confidence,
},
)
step_data = step_response.json()
observation = FlakeForgeObservation(**step_data["observation"])
reward = step_data.get("reward", 0.0)
total_reward += reward
trajectory.append({
"step": step + 1,
"category": action.predicted_category,
"confidence": action.predicted_confidence,
"reward": reward,
"pass_rate": step_data.get("state", {}).get("current_pass_rate", 0.0),
"done": step_data.get("done", False),
})
if step_data.get("done", False):
break
return {
"trajectory": trajectory,
"total_reward": total_reward,
"steps": len(trajectory),
"final_pass_rate": observation.current_pass_rate,
}
|