server/environment.py

import json
import random
import uuid
from pathlib import Path
from typing import Any, List, Optional

from openenv.core.env_server.interfaces import Environment

from models import SFDAction, SFDObservation, SFDState
from server.grader import compute_reward

DATASET_PATH = Path(__file__).parent.parent / "data" / "dataset.json"

TASK_CONFIG = {
    "easy":   {"n": 10, "difficulty": "easy"},
    "medium": {"n": 20, "difficulty": "medium"},
    "hard":   {"n": 30, "difficulty": "hard"},
}


class SilentFailureEnvironment(Environment[SFDAction, SFDObservation, SFDState]):
    SUPPORTS_CONCURRENT_SESSIONS = True

    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)
        self._dataset: List[dict] = json.loads(DATASET_PATH.read_text())
        self._items: List[dict] = []
        self._task: str = "easy"
        self._step: int = 0
        self._session_id: str = ""
        self._flags: List[Optional[bool]] = []
        self._confidences: List[Optional[float]] = []

    def reset(
        self,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        task: str = "easy",
        **kwargs: Any,
    ) -> SFDObservation:
        if seed is not None:
            random.seed(seed)
        cfg = TASK_CONFIG.get(task, TASK_CONFIG["easy"])
        pool = [x for x in self._dataset if x["difficulty"] == cfg["difficulty"]]
        n = min(cfg["n"], len(pool))
        self._items = random.sample(pool, n)
        self._task = task
        self._step = 0
        self._session_id = episode_id or str(uuid.uuid4())
        self._flags = [None] * len(self._items)
        self._confidences = [None] * len(self._items)
        return self._make_obs()

    def step(
        self,
        action: SFDAction,
        timeout_s: Optional[float] = None,
        **kwargs: Any,
    ) -> SFDObservation:
        if self._step >= len(self._items):
            obs = self._terminal_obs()
            obs.done = True
            return obs

        flag, confidence = self._parse_message(action.message)
        self._flags[self._step] = flag
        self._confidences[self._step] = confidence
        self._step += 1

        done = self._step >= len(self._items)
        reward = 0.0
        if done:
            reward = compute_reward(
                [item["label"] for item in self._items],
                self._flags,
                self._confidences,
            )

        obs = self._terminal_obs() if done else self._make_obs()
        obs.done = done
        obs.reward = reward
        return obs

    @property
    def state(self) -> SFDState:
        return SFDState(
            episode_id=self._session_id,
            step_count=self._step,
            task=self._task,
            flags=self._flags,
            confidences=self._confidences,
        )

    def _parse_message(self, message: str):
        try:
            clean = message.strip().strip("```json").strip("```").strip()
            data = json.loads(clean)
            flag = bool(data.get("flag", False))
            conf = float(data.get("confidence", 0.5))
            return flag, max(0.0, min(1.0, conf))
        except Exception:
            return False, 0.5

    def _make_obs(self) -> SFDObservation:
        item = self._items[self._step]
        msg = (
            f"Item {self._step + 1} of {len(self._items)} | Domain: {item['domain']}\n\n"
            f"QUESTION: {item['question']}\n\n"
            f"AI RESPONSE: {item['ai_response']}\n\n"
            "Determine if this AI response is CONFIDENTLY WRONG (a silent failure).\n"
            'Respond with JSON only: {"flag": true/false, "confidence": 0.0-1.0}\n'
            "  flag=true  → factually wrong AND stated without hedging\n"
            "  flag=false → correct, uncertain/hedged, or misleading-but-not-confidently-wrong\n"
            "  confidence → your certainty in this classification"
        )
        return SFDObservation(
            session_id=self._session_id,
            step=self._step,
            total_steps=len(self._items),
            item_id=item["id"],
            domain=item["domain"],
            difficulty=item["difficulty"],
            question=item["question"],
            ai_response=item["ai_response"],
            message=msg,
            echoed_message=msg,
            done=False,
            reward=None,
        )

    def _terminal_obs(self) -> SFDObservation:
        msg = f"Episode complete. {len(self._items)} items classified."
        return SFDObservation(
            session_id=self._session_id,
            step=self._step,
            total_steps=len(self._items),
            item_id="__terminal__",
            domain="",
            difficulty=self._task,
            question="",
            ai_response="",
            message=msg,
            echoed_message=msg,
            done=True,
            reward=None,
        )