feat: implement actions, reward, env, server

.gitignore

@@ -0,0 +1,11 @@
+venv/
+__pycache__/
+*.pyc
+*.pyo
+.env
+*.egg-info/
+dist/
+build/
+.DS_Store
+data/facts.json
+assets/*.png

environment/actions.py

@@ -0,0 +1,42 @@
+from enum import Enum
+from typing import Optional
+from pydantic import BaseModel, field_validator
+
+
+class ActionType(str, Enum):
+    read_doc = "read_doc"
+    flag_suspicious = "flag_suspicious"
+    unflag_doc = "unflag_doc"
+    submit_answer = "submit_answer"
+
+
+class ContextCorruptionAction(BaseModel):
+    action_type: ActionType
+    doc_id: Optional[int] = None
+    answer: Optional[str] = None
+    confidence: Optional[float] = None
+
+    @field_validator("confidence")
+    @classmethod
+    def confidence_range(cls, v):
+        if v is not None and not (0.0 <= v <= 1.0):
+            raise ValueError("confidence must be between 0.0 and 1.0")
+        return v
+
+
+class Document(BaseModel):
+    id: int
+    title: str
+    content: str
+    is_flagged: bool = False
+
+
+class EpisodeObservation(BaseModel):
+    question: str
+    documents: list[Document]
+    flagged_ids: list[int]
+    budget_remaining: int
+    turn: int
+    episode_done: bool = False
+    reward: Optional[float] = None
+    message: Optional[str] = None

environment/env.py

@@ -0,0 +1,140 @@
+import json
+import random
+from pathlib import Path
+
+from environment.actions import ActionType, ContextCorruptionAction, Document, EpisodeObservation
+from environment.reward import compute_reward
+
+_FALLBACK_FACTS = [
+    {"question": "What is the capital of France?", "answer": "Paris"}
+]
+
+
+class ContextCorruptionEnv:
+    MAX_BUDGET = 12
+    NUM_DOCS = 8
+    DIFFICULTY_LEVELS = [1, 2, 3, 4]
+
+    def __init__(self, difficulty=None):
+        self.difficulty = difficulty
+        facts_path = Path(__file__).parent.parent / "data" / "facts.json"
+        if facts_path.exists():
+            with open(facts_path) as f:
+                self._facts = json.load(f)
+        else:
+            self._facts = _FALLBACK_FACTS
+        self._reset_state()
+
+    def _reset_state(self):
+        self._question = ""
+        self._ground_truth = ""
+        self._documents: list[dict] = []
+        self._corrupt_ids: list[int] = []
+        self._flagged_ids: list[int] = []
+        self._budget_used = 0
+        self._turn = 0
+        self._done = False
+        self._reward = None
+        self._breakdown = None
+
+    def reset(self) -> EpisodeObservation:
+        self._reset_state()
+        fact = random.choice(self._facts)
+        n_corrupt = self.difficulty if self.difficulty is not None else random.choice(self.DIFFICULTY_LEVELS)
+        self._corrupt_ids = random.sample(range(self.NUM_DOCS), n_corrupt)
+        self._question = fact["question"]
+        self._ground_truth = fact["answer"]
+
+        try:
+            from data.generator import generate_documents
+            raw_docs = generate_documents(fact, num_docs=self.NUM_DOCS, corrupt_positions=self._corrupt_ids)
+        except Exception:
+            raw_docs = [
+                {"id": i, "title": f"Document {i}", "content": fact["answer"], "is_corrupt": i in self._corrupt_ids}
+                for i in range(self.NUM_DOCS)
+            ]
+
+        self._documents = raw_docs
+        return self._build_observation()
+
+    def step(self, action: ContextCorruptionAction) -> EpisodeObservation:
+        if self._done:
+            return self._build_observation(message="Episode already done.")
+
+        self._turn += 1
+        self._budget_used += 1
+        reward = None
+
+        if action.action_type == ActionType.read_doc:
+            pass  # budget cost is the point; content already in observation
+
+        elif action.action_type == ActionType.flag_suspicious:
+            if action.doc_id is not None and action.doc_id not in self._flagged_ids:
+                self._flagged_ids.append(action.doc_id)
+
+        elif action.action_type == ActionType.unflag_doc:
+            if action.doc_id in self._flagged_ids:
+                self._flagged_ids.remove(action.doc_id)
+
+        elif action.action_type == ActionType.submit_answer:
+            reward, self._breakdown = compute_reward(
+                submitted_answer=action.answer or "",
+                ground_truth_answer=self._ground_truth,
+                flagged_ids=self._flagged_ids,
+                corrupt_ids=self._corrupt_ids,
+                confidence=action.confidence or 0.0,
+                budget_used=self._budget_used,
+                max_budget=self.MAX_BUDGET,
+            )
+            self._reward = reward
+            self._done = True
+
+        # Force-submit on budget exhaustion
+        if self._budget_used >= self.MAX_BUDGET and not self._done:
+            reward, self._breakdown = compute_reward(
+                submitted_answer="",
+                ground_truth_answer=self._ground_truth,
+                flagged_ids=self._flagged_ids,
+                corrupt_ids=self._corrupt_ids,
+                confidence=0.0,
+                budget_used=self._budget_used,
+                max_budget=self.MAX_BUDGET,
+            )
+            self._reward = reward
+            self._done = True
+
+        return self._build_observation(reward=reward)
+
+    def state(self) -> dict:
+        return {
+            "question": self._question,
+            "ground_truth": self._ground_truth,
+            "corrupt_ids": self._corrupt_ids,
+            "flagged_ids": self._flagged_ids,
+            "budget_used": self._budget_used,
+            "turn": self._turn,
+            "done": self._done,
+            "reward": self._reward,
+            "breakdown": self._breakdown,
+        }
+
+    def _build_observation(self, reward=None, message=None) -> EpisodeObservation:
+        docs = [
+            Document(
+                id=d["id"],
+                title=d["title"],
+                content=d["content"],
+                is_flagged=d["id"] in self._flagged_ids,
+            )
+            for d in self._documents
+        ]
+        return EpisodeObservation(
+            question=self._question,
+            documents=docs,
+            flagged_ids=list(self._flagged_ids),
+            budget_remaining=self.MAX_BUDGET - self._budget_used,
+            turn=self._turn,
+            episode_done=self._done,
+            reward=reward if reward is not None else self._reward,
+            message=message,
+        )

environment/reward.py

@@ -0,0 +1,50 @@
+import re
+
+
+def _normalize(text: str) -> str:
+    text = text.lower()
+    text = re.sub(r"[^\w\s]", "", text)
+    text = re.sub(r"\s+", " ", text).strip()
+    return text
+
+
+def compute_reward(
+    submitted_answer: str,
+    ground_truth_answer: str,
+    flagged_ids: list[int],
+    corrupt_ids: list[int],
+    confidence: float,
+    budget_used: int,
+    max_budget: int,
+) -> tuple[float, dict]:
+    # Answer correctness
+    correct = _normalize(submitted_answer) == _normalize(ground_truth_answer)
+    answer_score = 0.4 if correct else 0.0
+
+    # Flag recall
+    true_positives = [i for i in flagged_ids if i in corrupt_ids]
+    recall = len(true_positives) / len(corrupt_ids) if corrupt_ids else 0.0
+    recall_score = 0.3 * recall
+
+    # Precision (false positive penalty)
+    false_positives = [i for i in flagged_ids if i not in corrupt_ids]
+    precision_score = max(0.0, 0.2 - 0.1 * len(false_positives))
+
+    # Confidence calibration
+    confidence = confidence or 0.0
+    calibration_score = (0.1 * confidence) if correct else (-0.2 * confidence)
+
+    # Efficiency bonus
+    efficiency_score = 0.05 * (1 - budget_used / max_budget)
+
+    total = answer_score + recall_score + precision_score + calibration_score + efficiency_score
+
+    breakdown = {
+        "answer_correctness": round(answer_score, 4),
+        "flag_recall": round(recall_score, 4),
+        "false_positive_penalty": round(precision_score, 4),
+        "confidence_calibration": round(calibration_score, 4),
+        "efficiency": round(efficiency_score, 4),
+        "total": round(total, 4),
+    }
+    return round(total, 4), breakdown

environment/server.py

@@ -0,0 +1,56 @@
+import uuid
+from fastapi import FastAPI, HTTPException
+import uvicorn
+
+from environment.actions import ContextCorruptionAction, EpisodeObservation
+from environment.env import ContextCorruptionEnv
+
+app = FastAPI(title="ContextCorruption-Env")
+
+# session_id -> env instance
+_sessions: dict[str, ContextCorruptionEnv] = {}
+_MAX_SESSIONS = 64
+
+
+@app.post("/reset", response_model=EpisodeObservation)
+def reset(session_id: str | None = None):
+    if session_id is None:
+        if len(_sessions) >= _MAX_SESSIONS:
+            raise HTTPException(status_code=503, detail="Max concurrent sessions reached")
+        session_id = str(uuid.uuid4())
+    if session_id not in _sessions:
+        _sessions[session_id] = ContextCorruptionEnv()
+    obs = _sessions[session_id].reset()
+    return obs
+
+
+@app.post("/step/{session_id}", response_model=EpisodeObservation)
+def step(session_id: str, action: ContextCorruptionAction):
+    if session_id not in _sessions:
+        raise HTTPException(status_code=404, detail="Session not found")
+    obs = _sessions[session_id].step(action)
+    if obs.episode_done:
+        del _sessions[session_id]
+    return obs
+
+
+@app.get("/state/{session_id}")
+def state(session_id: str):
+    if session_id not in _sessions:
+        raise HTTPException(status_code=404, detail="Session not found")
+    return _sessions[session_id].state()
+
+
+@app.delete("/session/{session_id}")
+def close_session(session_id: str):
+    _sessions.pop(session_id, None)
+    return {"status": "closed"}
+
+
+@app.get("/health")
+def health():
+    return {"status": "ok", "active_sessions": len(_sessions)}
+
+
+if __name__ == "__main__":
+    uvicorn.run("environment.server:app", host="0.0.0.0", port=8000, reload=False)