vasiuuu/DGX_AI
0
1from __future__ import annotations2 3import logging4import re5import uuid6from pathlib import Path7from typing import Any8 9from openenv.core.env_server.interfaces import Environment10 11from codeforge.audit.ledger import AuditLedger12from codeforge.grader import compute_reward13from codeforge.grounder import ground14from codeforge.interrogator.interrogator import Interrogator15from codeforge.kb.cluster import build_clusters16from codeforge.kb.indexer import SkillsIndex17from codeforge.models import AuditEntry, CodeForgeAction, CodeForgeActionType, CodeForgeObservation18from codeforge.observation import build_observation19from codeforge.ralph.loop import run_loop20from codeforge.ralph.models import LoopConfig21from codeforge.ralph.synthesizer import StubSynthesizer, Synthesizer22from codeforge.sandbox.sandbox import run_sandbox23from codeforge.shaping import citation_shaping_bonus24from codeforge.tasks import Task, get_task25 26_log = logging.getLogger(__name__)27_DEFAULT_CORPUS = Path(__file__).resolve().parent / "kb" / "skills_corpus.jsonl"28 29# ---------------------------------------------------------------------------30# Filename validation (SYSTEM_DESIGN §14.2, §14.3)31# ---------------------------------------------------------------------------32_FILENAME_RE = re.compile(r"^[a-z][a-z0-9_]*\.py$")33_FORBIDDEN_FILENAMES = frozenset({34 "conftest.py", "pytest.ini", "setup.cfg", "pyproject.toml", "tox.ini",35})36_MAX_FILES = 1037_MAX_FILE_SIZE = 50 * 1024 # 50 KB38_MAX_TOTAL_SIZE = 200 * 1024 # 200 KB39 40 41def _validate_files(files: dict[str, str]) -> str | None:42 """Return an error message if *files* violates submission constraints, else None."""43 if not files:44 return "files dict is empty"45 if len(files) > _MAX_FILES:46 return f"too many files ({len(files)} > {_MAX_FILES})"47 total_size = 048 for name, content in files.items():49 if name in _FORBIDDEN_FILENAMES:50 return f"filename '{name}' is not allowed"51 if not _FILENAME_RE.match(name):52 return f"filename '{name}' must match [a-z][a-z0-9_]*.py"53 size = len(content.encode("utf-8"))54 if size > _MAX_FILE_SIZE:55 return f"file '{name}' exceeds {_MAX_FILE_SIZE} bytes"56 total_size += size57 if total_size > _MAX_TOTAL_SIZE:58 return f"total size ({total_size}) exceeds {_MAX_TOTAL_SIZE} bytes"59 return None60 61 62# ---------------------------------------------------------------------------63# Valid action types (for fast membership check)64# ---------------------------------------------------------------------------65_VALID_ACTION_TYPES = frozenset(member.value for member in CodeForgeActionType)66 67 68# ---------------------------------------------------------------------------69# Environment70# ---------------------------------------------------------------------------71 72 73class CodeForgeEnvironment(Environment): # type: ignore[type-arg]74 """OpenEnv-compliant RL environment with all 6 CodeForge actions.75 76 Implements SYSTEM_DESIGN §4.9, §5.2, §17.77 """78 79 SUPPORTS_CONCURRENT_SESSIONS = True80 81 def __init__(82 self,83 *,84 corpus_path: Path | None = None,85 synthesizer: Synthesizer | None = None,86 ) -> None:87 super().__init__()88 self._corpus_path = corpus_path or _DEFAULT_CORPUS89 self._synthesizer = synthesizer90 self._index: SkillsIndex | None = None91 self._task: Task | None = None92 self._episode_id: str = ""93 self._budget_remaining: int = 094 self._current_files: dict[str, str] = {}95 self._previous_score: float = 0.096 self._is_done: bool = False97 98 # Per-step state99 self._last_citations: tuple[dict[str, object], ...] = ()100 self._last_grounding: dict[str, object] | None = None101 self._last_reward: float = 0.0102 self._last_cluster_hits: tuple[str, ...] = ()103 self._last_interrogation_questions: tuple[str, ...] = ()104 self._last_ralph_run_id: str | None = None105 self._last_ralph_iterations: tuple[dict[str, object], ...] = ()106 107 # Brier/quality tracking for audit entries108 self._last_brier_penalty: float | None = None109 self._last_quality: float = 0.0110 111 # Episode-level accumulators112 self._all_episode_citations: list[dict[str, object]] = []113 self._all_episode_cluster_hits: list[str] = []114 self._ledger: AuditLedger | None = None115 self._step_index: int = 0116 117 # ------------------------------------------------------------------118 # Index management119 # ------------------------------------------------------------------120 121 def _ensure_index(self) -> SkillsIndex:122 if self._index is None:123 if not self._corpus_path.is_file():124 msg = (125 f"corpus not found: {self._corpus_path}. "126 f"Run the skills scraper first."127 )128 raise FileNotFoundError(msg)129 idx = SkillsIndex(corpus_path=self._corpus_path)130 idx.build()131 # Build and attach clusters132 import json133 nodes: list[dict[str, Any]] = []134 with self._corpus_path.open(encoding="utf-8") as f:135 for line in f:136 line = line.strip()137 if line:138 nodes.append(json.loads(line))139 manifest = build_clusters(nodes)140 idx.attach_cluster_manifest(manifest)141 self._index = idx142 return self._index143 144 # ------------------------------------------------------------------145 # OpenEnv interface146 # ------------------------------------------------------------------147 148 def reset(149 self,150 seed: int | None = None,151 episode_id: str | None = None,152 **kwargs: Any,153 ) -> CodeForgeObservation:154 task_level: str = kwargs.get("task_level", "easy")155 task = get_task(task_level)156 self._task = task157 self._episode_id = episode_id or uuid.uuid4().hex[:12]158 self._budget_remaining = task.max_budget159 self._current_files = dict(task.initial_files)160 self._previous_score = 0.0161 self._is_done = False162 163 # Reset per-step164 self._last_citations = ()165 self._last_grounding = None166 self._last_reward = 0.0167 self._last_cluster_hits = ()168 self._last_interrogation_questions = ()169 self._last_ralph_run_id = None170 self._last_ralph_iterations = ()171 172 # Reset episode accumulators173 self._all_episode_citations = []174 self._all_episode_cluster_hits = []175 self._ledger = AuditLedger()176 self._step_index = 0177 178 _log.info(179 "reset id=%s task=%s budget=%s",180 self._episode_id, task.task_id, task.max_budget,181 )182 return self._build_obs()183 184 def step(185 self,186 action: CodeForgeAction,187 timeout_s: float | None = None,188 **kwargs: Any,189 ) -> CodeForgeObservation:190 # --- Pre-check: no active episode --------------------------------191 if self._task is None:192 return self._error_obs("No active episode — call reset() first")193 194 # --- Pre-check: episode already done -----------------------------195 if self._is_done:196 return self._build_obs()197 198 # --- Pre-check: valid action_type --------------------------------199 action_type_str = str(action.action_type)200 if action_type_str not in _VALID_ACTION_TYPES:201 return self._error_obs(f"Unknown action_type: {action_type_str!r}")202 203 # --- Budget check (variable cost) --------------------------------204 cost = self._action_cost(action)205 if cost > self._budget_remaining:206 return self._error_obs(207 f"Insufficient budget: need {cost}, have {self._budget_remaining}"208 )209 self._budget_remaining -= cost210 211 # --- Clear per-step state ----------------------------------------212 self._last_reward = 0.0213 self._last_citations = ()214 self._last_grounding = None215 self._last_cluster_hits = ()216 self._last_interrogation_questions = ()217 self._last_ralph_run_id = None218 self._last_ralph_iterations = ()219 error: str | None = None220 221 # --- Route to handler --------------------------------------------222 try:223 if action_type_str == CodeForgeActionType.QUERY_KB:224 error = self._handle_query_kb(action)225 elif action_type_str == CodeForgeActionType.QUERY_CLUSTER:226 error = self._handle_query_cluster(action)227 elif action_type_str == CodeForgeActionType.INTERROGATE:228 error = self._handle_interrogate(action)229 elif action_type_str == CodeForgeActionType.SUBMIT:230 error = self._handle_submit(action)231 elif action_type_str == CodeForgeActionType.RUN_RALPH:232 error = self._handle_run_ralph(action)233 elif action_type_str == CodeForgeActionType.GET_AUDIT:234 error = self._handle_get_audit(action)235 except Exception as exc:236 _log.exception("handler error: %s", exc)237 error = f"Internal error: {exc}"238 239 # --- Append audit entry ------------------------------------------240 assert self._ledger is not None241 _cited: list[str] = []242 _cite: dict[str, object]243 for _cite in self._last_citations:244 _cited.append(str(_cite.get("node_id", "")))245 cited_ids: tuple[str, ...] = tuple(_cited)246 self._ledger.append(247 AuditEntry(248 step_index=self._step_index,249 action_type=action_type_str,250 cited_skill_ids=cited_ids,251 cited_clusters=self._last_cluster_hits,252 grounding_report=(253 self._last_grounding if self._last_grounding else None254 ),255 reward=self._last_reward,256 brier_penalty=(257 self._last_brier_penalty258 if action_type_str == CodeForgeActionType.SUBMIT259 else None260 ),261 confidence_declared=(262 action.confidence263 if action_type_str == CodeForgeActionType.SUBMIT264 else None265 ),266 quality=(267 self._last_quality268 if action_type_str == CodeForgeActionType.SUBMIT269 else self._previous_score270 ),271 ),272 )273 self._step_index += 1274 275 # --- Check budget exhaustion -------------------------------------276 if self._budget_remaining <= 0:277 self._is_done = True278 279 return self._build_obs(error=error)280 281 @property282 def state(self) -> CodeForgeObservation:283 if self._task is None:284 return self._error_obs("No active episode — call reset() first")285 return self._build_obs()286 287 # ------------------------------------------------------------------288 # Cost computation289 # ------------------------------------------------------------------290 291 @staticmethod292 def _action_cost(action: CodeForgeAction) -> int:293 """Variable-cost budget accounting (SYSTEM_DESIGN §17.2)."""294 if str(action.action_type) == CodeForgeActionType.GET_AUDIT:295 return 0296 if str(action.action_type) == CodeForgeActionType.RUN_RALPH:297 return action.max_iters298 return 1299 300 # ------------------------------------------------------------------301 # Action handlers (each returns an error string or None)302 # ------------------------------------------------------------------303 304 def _handle_query_kb(self, action: CodeForgeAction) -> str | None:305 try:306 idx = self._ensure_index()307 except FileNotFoundError as e:308 _log.warning("query_kb: no corpus: %s", e)309 self._last_citations = ()310 return None311 tags = set(action.required_tags) if action.required_tags else None312 results = idx.search(313 action.claim or "", top_k=action.top_k, required_tags=tags,314 )315 self._last_citations = tuple(316 {317 "node_id": r.node_id,318 "skill_name": r.skill_name,319 "section_path": list(r.section_path),320 "section_body": r.section_body,321 "score": r.score,322 "rank": r.rank,323 }324 for r in results325 )326 self._all_episode_citations.extend(self._last_citations)327 return None328 329 def _handle_query_cluster(self, action: CodeForgeAction) -> str | None:330 try:331 idx = self._ensure_index()332 except FileNotFoundError as e:333 _log.warning("query_cluster: no corpus: %s", e)334 self._last_cluster_hits = ()335 return None336 label = action.cluster_label or ""337 results = idx.nodes_in_cluster(label)338 if not results:339 self._last_cluster_hits = ()340 return None341 self._last_cluster_hits = tuple(r.node_id for r in results)342 self._all_episode_cluster_hits.extend(self._last_cluster_hits)343 return None344 345 def _handle_interrogate(self, action: CodeForgeAction) -> str | None:346 idx: SkillsIndex | None347 try:348 idx = self._ensure_index()349 except FileNotFoundError:350 idx = None351 interrogator = Interrogator(idx)352 assert self._task is not None353 result = interrogator.generate(self._task.brief)354 self._last_interrogation_questions = result.questions355 return None356 357 def _handle_submit(self, action: CodeForgeAction) -> str | None:358 if action.files is None:359 return "files required for submit"360 file_err = _validate_files(action.files)361 if file_err is not None:362 return file_err363 364 self._current_files = dict(action.files)365 assert self._task is not None366 367 # Merge hidden correctness tests into sandbox files (agent cannot see these)368 sandbox_files = dict(action.files)369 if self._task.hidden_tests:370 sandbox_files.update(self._task.hidden_tests)371 372 # Run sandbox373 try:374 sandbox_result = run_sandbox(375 files=sandbox_files,376 tools=self._task.tools,377 timeout_per_tool=30.0,378 )379 sandbox_score = sandbox_result.composite_score380 except Exception as e:381 _log.exception("sandbox error: %s", e)382 sandbox_score = 0.0383 384 # Run grounder (pass local module names so they're not penalized)385 local_modules = frozenset(386 f.removesuffix(".py") for f in action.files if f.endswith(".py")387 )388 concatenated = "\n".join(action.files.values())389 grounding_report = ground(concatenated, local_modules=local_modules)390 self._last_grounding = grounding_report.model_dump()391 392 # Compute reward with Brier calibration393 quality = 0.6 * sandbox_score + 0.4 * grounding_report.groundedness394 effective_conf = action.confidence if action.confidence is not None else 0.5395 brier_penalty: float | None = min((effective_conf - quality) ** 2, 0.5)396 self._last_brier_penalty = brier_penalty397 self._last_quality = quality398 399 reward = compute_reward(400 sandbox_score=sandbox_score,401 groundedness=grounding_report.groundedness,402 confidence=action.confidence,403 )404 405 # Apply citation shaping bonus only on successful submits (§4.8.4)406 if reward > 0:407 shaping = citation_shaping_bonus(408 submit_files=action.files,409 prior_citations=self._all_episode_citations,410 prior_cluster_hits=self._all_episode_cluster_hits,411 )412 reward = round(min(1.0, reward + shaping), 3)413 414 self._last_reward = reward415 self._previous_score = reward416 417 # Check target score418 if reward >= self._task.target_score:419 self._is_done = True420 421 return None422 423 def _handle_run_ralph(self, action: CodeForgeAction) -> str | None:424 assert self._task is not None425 try:426 idx = self._ensure_index()427 except FileNotFoundError as e:428 return f"corpus not available: {e}"429 430 config = LoopConfig(431 max_iters=action.max_iters,432 target_score=self._task.target_score,433 tools=self._task.tools,434 )435 synthesizer = self._synthesizer or StubSynthesizer()436 result = run_loop(437 spec=self._task.brief,438 initial_files=self._current_files,439 index=idx,440 synthesizer=synthesizer,441 config=config,442 )443 444 self._last_ralph_run_id = result.run_id445 self._last_ralph_iterations = tuple(446 it.model_dump() for it in result.iterations447 )448 self._current_files = dict(result.final_files)449 450 # Compute ralph reward (SYSTEM_DESIGN §4.8.5)451 concatenated = "\n".join(result.final_files.values())452 grounding_report = ground(concatenated)453 self._last_grounding = grounding_report.model_dump()454 455 wasted = sum(456 1 for it in result.iterations if it.reason in ("score_regressed", "score_plateau")457 )458 base = compute_reward(459 sandbox_score=result.final_score,460 groundedness=grounding_report.groundedness,461 confidence=0.75,462 )463 waste_penalty = wasted * 0.05464 ralph_reward = round(max(0.0, min(1.0, base - waste_penalty)), 3)465 466 self._last_reward = ralph_reward467 self._previous_score = ralph_reward468 return None469 470 def _handle_get_audit(self, action: CodeForgeAction) -> str | None:471 # Audit data is populated in _build_obs via cumulative_audit_summary472 return None473 474 # ------------------------------------------------------------------475 # Observation helpers476 # ------------------------------------------------------------------477 478 def _build_obs(self, *, error: str | None = None) -> CodeForgeObservation:479 assert self._task is not None480 audit_summary: dict[str, object] | None = None481 if self._ledger is not None:482 audit_summary = self._ledger.serialize()483 return build_observation(484 episode_id=self._episode_id,485 task=self._task,486 current_files=self._current_files,487 budget_remaining=self._budget_remaining,488 previous_score=self._previous_score,489 last_citations=self._last_citations,490 last_grounding=self._last_grounding,491 is_done=self._is_done,492 last_reward=self._last_reward,493 last_cluster_hits=self._last_cluster_hits,494 last_interrogation_questions=self._last_interrogation_questions,495 last_ralph_run_id=self._last_ralph_run_id,496 last_ralph_iterations=self._last_ralph_iterations,497 cumulative_audit_summary=audit_summary,498 error=error,499 )500 501 def _error_obs(self, msg: str) -> CodeForgeObservation:502 """Return an error observation without modifying episode state."""503 if self._task is None:504 # No task set — use a dummy task for the observation structure505 dummy = get_task("easy")506 return build_observation(507 episode_id=self._episode_id or "none",508 task=dummy,509 current_files=self._current_files,510 budget_remaining=self._budget_remaining,511 previous_score=self._previous_score,512 is_done=self._is_done,513 error=msg,514 )515 return self._build_obs(error=msg)516