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,
            }