File size: 6,187 Bytes
4423d13
06b0d69
 
 
 
 
 
 
 
 
 
 
 
 
4423d13
 
06b0d69
 
 
 
 
3424779
06b0d69
 
3424779
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
06b0d69
 
 
 
5834385
06b0d69
3424779
 
 
06b0d69
 
 
 
 
 
 
 
 
3424779
06b0d69
3424779
06b0d69
3424779
06b0d69
 
 
 
3424779
06b0d69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4423d13
06b0d69
 
 
 
 
3424779
 
 
 
 
 
 
06b0d69
 
 
3424779
 
06b0d69
 
 
 
3424779
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
06b0d69
 
 
 
 
 
 
 
 
 
 
 
 
3424779
06b0d69
 
 
 
5834385
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
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
import os
import random
import httpx

try:
    from openenv.core.env_server import Environment
except ImportError:
    from openenv.core.env_server.interfaces import Environment

from models import APIAction, APIObservation, APIState
from tasks.registry import TASK_REGISTRY
from graders.grader import APIGrader

GRADER = APIGrader()
MOCK_BASE = os.getenv("MOCK_BASE_URL", "http://localhost:7860")
MAX_RESPONSE_BODY_LENGTH = 2000


class APIDebugEnvironment(Environment):

    def __init__(self):
        self._state = APIState(episode_id="init", step_count=0, task_id="easy", max_steps=10, solved=False, current_step_index=0, curriculum_level=0)
        self._current_task = None

    def reset(self, task_id: str = "auto", **kwargs) -> APIObservation:
        import copy
        import uuid

        old_level = getattr(self._state, "curriculum_level", 0)
        if getattr(self._state, "solved", False):
            old_level += 1

        effective_task_id = task_id
        if task_id == "auto":
            levels = ["easy", "medium", "hard", "expert"]
            effective_task_id = levels[min(old_level, len(levels) - 1)]

        if effective_task_id not in TASK_REGISTRY:
            effective_task_id = "easy"
            
        task = copy.deepcopy(random.choice(TASK_REGISTRY[effective_task_id]))

        dyn_config = {
            "demo_token": f"tok_{uuid.uuid4().hex[:8]}",
            "client_id": f"cli_{uuid.uuid4().hex[:8]}",
            "client_secret": f"sec_{uuid.uuid4().hex[:8]}"
        }

        try:
            with httpx.Client(base_url=MOCK_BASE, timeout=2.0) as http:
                http.post("/mock_api/_admin/reset", json=dyn_config)
        except Exception:
            pass

        try:
            task["description"] = task["description"].format(**dyn_config)
            req = task["broken_request"]
            if "headers" in req:
                req["headers"] = {k: (v.format(**dyn_config) if isinstance(v, str) else v) for k, v in req["headers"].items()}
        except Exception:
            pass

        self._current_task = task
        max_steps = self._current_task.get("max_steps", 10)
        
        self._state = APIState(
            episode_id=f"ep_{random.randint(10000, 99999)}",
            step_count=0,
            task_id=task_id,
            max_steps=max_steps,
            solved=False,
            current_step_index=0,
            visited_endpoints=set(),
            curriculum_level=old_level
        )
        return APIObservation(
            task_id=task_id,
            task_description=self._current_task["description"],
            broken_request=self._current_task["broken_request"],
            last_status_code=0,
            last_response_headers={},
            last_response_body="",
            step_feedback="Episode started. Read the task description and fix the broken request.",
            current_score=0.001,
            attempt=0,
            reward=0.001,
            done=False,
            current_step_index=0
        )

    def step(self, action: APIAction, **kwargs) -> APIObservation:
        if self._current_task is None:
            self.reset(task_id="auto")

        self._state.step_count += 1
        task = self._current_task

        try:
            with httpx.Client(base_url=MOCK_BASE, timeout=5.0) as http:
                resp = http.request(
                    method=action.method.upper(),
                    url=action.url,
                    headers=action.headers,
                    json=action.body if action.body else None,
                    params=action.query_params,
                )
            status = resp.status_code
            resp_headers = dict(resp.headers)
            resp_body = resp.text[:MAX_RESPONSE_BODY_LENGTH]
        except Exception as e:
            status = 0
            resp_headers = {}
            resp_body = f"Request error: {str(e)}"

        expected_schema = task["expected_schema"]
        expected_status = task["expected_status"]
        is_chain = isinstance(expected_schema, list)

        target_schema = expected_schema[self._state.current_step_index] if is_chain else expected_schema
        target_status = expected_status[self._state.current_step_index] if isinstance(expected_status, list) else expected_status

        reward = GRADER.grade(
            response_status=status,
            response_body=resp_body,
            expected_status=target_status,
            expected_schema=target_schema,
            attempt=self._state.step_count,
            max_steps=self._state.max_steps,
        )

        feedback = GRADER.get_feedback(status, target_status)
        step_solved = status == target_status and reward >= 0.8
        
        if is_chain and step_solved:
            self._state.current_step_index += 1
            if self._state.current_step_index >= len(expected_schema):
                self._state.solved = True
                feedback += " Chain completed successfully!"
            else:
                feedback += f" Step {self._state.current_step_index} complete. Moving to next step."
                reward = 0.5  # Partial reward for finishing a step
        else:
            self._state.solved = step_solved if not is_chain else False

        if action.url not in self._state.visited_endpoints and status > 0 and status < 500:
            self._state.visited_endpoints.add(action.url)
            reward += 0.05
        reward = round(min(0.999, max(0.001, reward)), 4)

        done = self._state.solved or self._state.step_count >= self._state.max_steps

        return APIObservation(
            task_id=self._state.task_id,
            task_description=task["description"],
            broken_request=task["broken_request"],
            last_status_code=status,
            last_response_headers=resp_headers,
            last_response_body=resp_body,
            step_feedback=feedback,
            current_score=reward,
            attempt=self._state.step_count,
            reward=reward,
            done=done,
            current_step_index=self._state.current_step_index
        )

    @property
    def state(self) -> APIState:
        return self._state