Spaces:
Sleeping
Sleeping
| """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, | |
| } | |