File size: 3,246 Bytes
e298389 d5e0274 e298389 | 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 | from fastapi import FastAPI
from typing import Any, Dict, Tuple
import os
from src.reward import compute_reward
from src.tasks import GRADERS, TASKS
app = FastAPI()
class CodeGuardEnv:
def __init__(self) -> None:
self.state: Dict[str, Any] = {}
self.current_step: int = 0
self.max_steps: int = 50
self.threshold: float = 0.95
self.done: bool = False
requested_task = os.getenv("TASK", "easy").strip().lower()
self.task_key: str = requested_task if requested_task in GRADERS else "easy"
def _get_task_score(self, action: str) -> float:
grader = GRADERS[self.task_key]
base_score = float(grader(action))
return max(0.01, min(0.99, base_score))
def _get_all_task_scores(self, action: str) -> Dict[str, float]:
scores: Dict[str, float] = {}
for key, grader in GRADERS.items():
score = float(grader(action))
scores[key] = max(0.01, min(0.99, score))
return scores
def reset(self) -> Dict[str, Any]:
self.state = {
"score": 0.01,
"history": [],
"task": TASKS[self.task_key],
"tasks": list(TASKS.values()),
"task_scores": {k: 0.01 for k in GRADERS.keys()},
}
self.current_step = 0
self.done = False
return self.state
def _is_valid_action(self, action: str) -> bool:
if not isinstance(action, str):
return False
if len(action.strip()) == 0:
return False
if len(action) > 1000:
return False
return True
def step(self, action: str) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:
if self.done:
return self.state, 0.0, True, {"error": "episode_done"}
self.current_step += 1
info: Dict[str, Any] = {"error": None}
if not self._is_valid_action(action):
self.done = True
return self.state, -1.0, True, {"error": "invalid_action"}
base_score: float = self._get_task_score(action)
reward: float = compute_reward(self.state, action, base_score)
self.state["score"] = max(0.01, min(0.99, base_score))
self.state["task_scores"] = self._get_all_task_scores(action)
self.state["history"].append(
{
"step": self.current_step,
"action": action,
"reward": reward,
}
)
if reward <= -2.0:
self.done = True
elif reward >= self.threshold:
self.done = True
elif self.current_step >= self.max_steps:
self.done = True
return self.state, reward, self.done, info
env_instance = CodeGuardEnv()
@app.get("/")
def health_check() -> Dict[str, str]:
return {"status": "ok"}
@app.post("/reset")
def reset() -> Dict[str, Any]:
return env_instance.reset()
@app.get("/tasks")
def tasks() -> Dict[str, Any]:
return {
"count": len(TASKS),
"tasks": list(TASKS.values()),
"graders": sorted(GRADERS.keys()),
}
@app.post("/grade")
def grade(task_id: str, candidate_code: str) -> Dict[str, Any]:
key = task_id.strip().lower()
if key not in GRADERS:
return {"error": "unknown_task", "task_id": task_id}
score = float(GRADERS[key](candidate_code))
score = max(0.01, min(0.99, score))
return {"task_id": key, "score": score}
@app.post("/step")
def step(action: str) -> Dict[str, Any]:
state, reward, done, info = env_instance.step(action)
return {
"state": state,
"reward": reward,
"done": done,
"info": info,
}
@app.get("/state")
def state() -> Dict[str, Any]:
return env_instance.state
|