Jivan01/agentBox

1from fastapi import FastAPI2from typing import Any, Dict, Tuple3import os4 5from src.reward import compute_reward6from src.tasks import GRADERS, TASKS7 8app = FastAPI()9 10 11class CodeGuardEnv:12	def __init__(self) -> None:13		self.state: Dict[str, Any] = {}14		self.current_step: int = 015		self.max_steps: int = 5016		self.threshold: float = 0.9517		self.done: bool = False18		requested_task = os.getenv("TASK", "easy").strip().lower()19		self.task_key: str = requested_task if requested_task in GRADERS else "easy"20 21	def _get_task_score(self, action: str) -> float:22		grader = GRADERS[self.task_key]23		base_score = float(grader(action))24		return max(0.01, min(0.99, base_score))25 26	def _get_all_task_scores(self, action: str) -> Dict[str, float]:27		scores: Dict[str, float] = {}28		for key, grader in GRADERS.items():29			score = float(grader(action))30			scores[key] = max(0.01, min(0.99, score))31		return scores32 33	def reset(self) -> Dict[str, Any]:34		self.state = {35			"score": 0.01,36			"history": [],37			"task": TASKS[self.task_key],38			"tasks": list(TASKS.values()),39			"task_scores": {k: 0.01 for k in GRADERS.keys()},40		}41		self.current_step = 042		self.done = False43		return self.state44 45	def _is_valid_action(self, action: str) -> bool:46		if not isinstance(action, str):47			return False48		if len(action.strip()) == 0:49			return False50		if len(action) > 1000:51			return False52		return True53 54	def step(self, action: str) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]:55		if self.done:56			return self.state, 0.0, True, {"error": "episode_done"}57 58		self.current_step += 159		info: Dict[str, Any] = {"error": None}60 61		if not self._is_valid_action(action):62			self.done = True63			return self.state, -1.0, True, {"error": "invalid_action"}64 65		base_score: float = self._get_task_score(action)66		reward: float = compute_reward(self.state, action, base_score)67 68		self.state["score"] = max(0.01, min(0.99, base_score))69		self.state["task_scores"] = self._get_all_task_scores(action)70		self.state["history"].append(71			{72				"step": self.current_step,73				"action": action,74				"reward": reward,75			}76		)77 78		if reward <= -2.0:79			self.done = True80		elif reward >= self.threshold:81			self.done = True82		elif self.current_step >= self.max_steps:83			self.done = True84 85		return self.state, reward, self.done, info86 87 88env_instance = CodeGuardEnv()89 90 91@app.get("/")92def health_check() -> Dict[str, str]:93	return {"status": "ok"}94 95 96@app.post("/reset")97def reset() -> Dict[str, Any]:98	return env_instance.reset()99 100 101@app.get("/tasks")102def tasks() -> Dict[str, Any]:103	return {104		"count": len(TASKS),105		"tasks": list(TASKS.values()),106		"graders": sorted(GRADERS.keys()),107	}108 109 110@app.post("/grade")111def grade(task_id: str, candidate_code: str) -> Dict[str, Any]:112	key = task_id.strip().lower()113	if key not in GRADERS:114		return {"error": "unknown_task", "task_id": task_id}115 116	score = float(GRADERS[key](candidate_code))117	score = max(0.01, min(0.99, score))118	return {"task_id": key, "score": score}119 120 121@app.post("/step")122def step(action: str) -> Dict[str, Any]:123	state, reward, done, info = env_instance.step(action)124	return {125		"state": state,126		"reward": reward,127		"done": done,128		"info": info,129	}130 131 132@app.get("/state")133def state() -> Dict[str, Any]:134	return env_instance.state135