Spaces:
Configuration error
Configuration error
| from __future__ import annotations | |
| import json | |
| import random | |
| import uuid | |
| from dataclasses import replace | |
| from pathlib import Path | |
| from .models import CommitGuardAction, CommitGuardObservation, CommitGuardState, ContextSnippet, DevignSample | |
| from .reward import compute_reward | |
| class CommitGuardEnvironment: | |
| def __init__(self, *, data_path: Path) -> None: | |
| self._data_path = data_path | |
| self._samples: list[DevignSample] = [] | |
| self._state: CommitGuardState | None = None | |
| self._rng = random.Random(0) | |
| self._cwe_keywords: dict[str, list[str]] = {} | |
| def load(self) -> None: | |
| if self._samples: | |
| return | |
| # Load CWE keywords from data directory (matching instructions) | |
| try: | |
| kw_path = self._data_path.parent / "cwe_keywords.json" | |
| if not kw_path.exists(): | |
| # Fallback to current directory or data subfolder if needed | |
| kw_path = self._data_path.parent / "data" / "cwe_keywords.json" | |
| self._cwe_keywords = json.loads(kw_path.read_text(encoding="utf-8")) | |
| except Exception: | |
| self._cwe_keywords = {} | |
| raw = self._data_path.read_text(encoding="utf-8").strip().splitlines() | |
| for line in raw: | |
| obj = json.loads(line) | |
| # Support both original and mvd schemas | |
| sample_id = str(obj.get("commit_id") or obj.get("sample_id", "unknown")) | |
| # Synthesize diff if missing (mvd branch data schema) | |
| diff = obj.get("diff") | |
| if not diff and "code_before" in obj and "code_after" in obj: | |
| diff = f"--- code_before\n+++ code_after\n{obj['code_before']}\n{obj['code_after']}" | |
| self._samples.append( | |
| DevignSample( | |
| sample_id=sample_id, | |
| diff=str(diff or ""), | |
| available_files=list(obj.get("available_files") or []), | |
| is_vulnerable=obj.get("is_vulnerable"), | |
| cwe=obj.get("cwe") or obj.get("cwe_type"), | |
| target_file=obj.get("target_file"), | |
| files=obj.get("files"), | |
| ) | |
| ) | |
| if not self._samples: | |
| raise RuntimeError("no_samples_loaded") | |
| def reset(self, sample_id: str | None = None) -> CommitGuardObservation: | |
| self.load() | |
| if sample_id: | |
| sample = next((s for s in self._samples if s.sample_id == sample_id), None) | |
| if not sample: | |
| raise ValueError(f"sample_id {sample_id} not found") | |
| else: | |
| sample = self._rng.choice(self._samples) | |
| episode_id = str(uuid.uuid4()) | |
| self._state = CommitGuardState( | |
| episode_id=episode_id, | |
| current_sample_id=sample.sample_id, | |
| step_count=0, | |
| context_requests=0, | |
| history=[], | |
| ) | |
| return CommitGuardObservation( | |
| episode_id=episode_id, | |
| diff=sample.diff, | |
| available_files=sample.available_files, | |
| step_idx=0, | |
| budget_remaining=5, | |
| ) | |
| def step(self, action: CommitGuardAction) -> tuple[CommitGuardObservation, float, bool]: | |
| if self._state is None: | |
| _ = self.reset() | |
| assert self._state is not None | |
| next_step = self._state.step_count + 1 | |
| sample = next(s for s in self._samples if s.sample_id == self._state.current_sample_id) | |
| context_snippets: list[ContextSnippet] = [] | |
| context_requests = self._state.context_requests | |
| if action.action_type == "request_context": | |
| context_requests += 1 | |
| if action.file_path and sample.files and action.file_path in sample.files: | |
| content = sample.files[action.file_path] | |
| lines = content.splitlines() | |
| start = 1 | |
| end = min(len(lines), 80) | |
| context_snippets = [ | |
| ContextSnippet( | |
| file_path=action.file_path, | |
| start_line=start, | |
| end_line=end, | |
| content="\n".join(lines[start - 1 : end]), | |
| ) | |
| ] | |
| reward = compute_reward( | |
| action=action, | |
| is_vulnerable=sample.is_vulnerable, | |
| cwe=sample.cwe, | |
| target_file=sample.target_file, | |
| cwe_keywords=self._cwe_keywords, | |
| context_requests=context_requests, | |
| ) | |
| done = bool(action.action_type == "verdict" or next_step >= 5) | |
| self._state = replace( | |
| self._state, | |
| step_count=next_step, | |
| context_requests=context_requests, | |
| history=[ | |
| *self._state.history, | |
| { | |
| "step": next_step, | |
| "action_type": action.action_type, | |
| "parse_error": action.parse_error, | |
| }, | |
| ], | |
| ) | |
| obs = CommitGuardObservation( | |
| episode_id=self._state.episode_id, | |
| diff=sample.diff, | |
| available_files=sample.available_files, | |
| context_snippets=context_snippets, | |
| step_idx=next_step, | |
| budget_remaining=max(0, 5 - next_step), | |
| error=action.parse_error or (None if context_snippets else ("context_unavailable" if action.action_type == "request_context" else None)), | |
| ) | |
| return obs, reward, done | |
| def state(self) -> CommitGuardState: | |
| if self._state is None: | |
| return CommitGuardState(episode_id="", current_sample_id="", step_count=0, context_requests=0, history=[]) | |
| return self._state | |