FlakeForge / client.py
random70249's picture
Upload folder using huggingface_hub
ee933ab verified
Raw
History Blame Contribute Delete
5.08 kB
"""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,
}