server/citeGuardian_environment.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
CiteGuardian Environment Implementation.

Simulates a professional peer-review / journal-editing workflow.
The agent audits a research paper for:
  Task A (Easy)   – Structural omissions (missing mandatory section)
  Task B (Medium) – Citation orphans (cited but not in References, or vice-versa)
  Task C (Hard)   – Factual contradictions across sections (numeric mismatches)

Reward structure (cumulative, max 1.0):
  +0.02  Exploration  – first visit to each required section
  +0.30  Accuracy     – correct FLAG_ERROR on a seeded mistake
  -0.10  False Pos.   – FLAG_ERROR where no error exists
  -0.01  Efficiency   – every step taken
  +0.10–0.40 Completion – SUBMIT after finding all errors (scales with recall)
  Clamped to 1.0 on perfect run (100 % recall, 0 false positives).
"""

import random
import re
from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

try:
    from ..models import CiteguardianAction, CiteguardianObservation
except ImportError:
    from models import CiteguardianAction, CiteguardianObservation


# ---------------------------------------------------------------------------
# Paper templates
# ---------------------------------------------------------------------------

_TASK_A_PAPER = {
    "Abstract": (
        "We present a novel deep-learning approach for protein folding prediction. "
        "Our model achieves state-of-the-art accuracy on the CASP14 benchmark. [1][3]"
    ),
    "Introduction": (
        "Protein structure prediction has long been a grand challenge in biology. "
        "Recent advances in transformer architectures [1] have enabled breakthroughs. "
        "This paper extends the work of Jumper et al. (2021) [2] by incorporating "
        "multi-scale attention. We recruited 120 domain experts to evaluate outputs."
    ),
    "Methods": (
        "We trained a 48-layer transformer on the UniRef90 database. "
        "Training used 128 TPU-v4 chips for 14 days. "
        "Hyperparameters were tuned via Bayesian optimisation [3]. "
        "We recruited 120 subjects for the human-evaluation study."
    ),
    # "Results" section intentionally MISSING — this is the seeded error
    "Discussion": (
        "Our approach outperforms prior methods on all benchmarks. "
        "Limitations include high compute cost and limited generalisation to "
        "intrinsically disordered proteins."
    ),
    "References": (
        "[1] Vaswani et al. (2017) Attention is All You Need. NeurIPS.\n"
        "[2] Jumper et al. (2021) Highly accurate protein structure prediction. Nature.\n"
        "[3] Snoek et al. (2012) Practical Bayesian Optimization. NeurIPS."
    ),
}

_TASK_A_ERRORS = [
    {
        "id": "err_A1",
        "error_type": "STRUCTURAL_ERROR",
        "section": None,  # not section-specific
        "hint": "Results",  # the missing section name
        "description": "The mandatory 'Results' section is absent from the paper.",
    }
]

# ----

_TASK_B_PAPER = {
    "Abstract": (
        "We introduce CiteNet, a citation-graph neural network. "
        "Experiments on three benchmarks confirm superior performance [1][2]."
    ),
    "Introduction": (
        "Citation networks encode rich relational information [1]. "
        "Prior work by Kipf & Welling (2017) [2] used GCNs; we extend this "
        "with dynamic edge weighting. See also the survey by Hamilton et al. [4]."
        # [4] is cited here but NOT in References — orphan citation
    ),
    "Methods": (
        "Our model stacks three graph-attention layers [3]. "
        "We use the Adam optimiser with lr=0.001. "
        "Datasets: Cora, Citeseer, PubMed."
    ),
    "Results": (
        "CiteNet achieves 89.2 % accuracy on Cora, surpassing all baselines. "
        "Full results are in Table 1."
    ),
    "Discussion": (
        "The dynamic edge weighting is the key contributor to performance gains. "
        "Future work will explore temporal citation graphs."
    ),
    "References": (
        "[1] Perozzi et al. (2014) DeepWalk. KDD.\n"
        "[2] Kipf & Welling (2017) Semi-supervised classification with GCNs. ICLR.\n"
        "[3] Velickovic et al. (2018) Graph Attention Networks. ICLR.\n"
        # [4] Hamilton et al. intentionally MISSING from References
        # [5] listed here but never cited in text — extra orphan
        "[5] Xu et al. (2019) How Powerful are Graph Neural Networks? ICLR."
    ),
}

_TASK_B_ERRORS = [
    {
        "id": "err_B1",
        "error_type": "ORPHAN_CITATION",
        "section": "Introduction",
        "hint": "[4]",
        "description": "[4] is cited in Introduction but has no entry in References.",
    },
    {
        "id": "err_B2",
        "error_type": "ORPHAN_CITATION",
        "section": "References",
        "hint": "[5]",
        "description": "[5] appears in References but is never cited in the paper body.",
    },
]

# ----

_TASK_C_PAPER = {
    "Abstract": (
        "We conduct a randomised controlled trial on a novel cognitive training "
        "programme. Results show significant improvement in working memory. [1][2]"
    ),
    "Introduction": (
        "Cognitive decline affects millions worldwide [1]. "
        "Intervention studies [2] show promise but lack rigorous controls. "
        "We address this gap with a pre-registered RCT."
    ),
    "Methods": (
        "We recruited 100 subjects aged 60–75 from three urban clinics. "
        "Participants were randomised 1:1 to treatment and control arms (50 each). "
        "Primary outcome: digit-span score at 12 weeks."
    ),
    "Results": (
        # Contradiction: Methods says 100 subjects, Results says 85
        "Table 1 shows data for 85 subjects who completed the 12-week assessment. "
        "The treatment group (n=44) showed a mean improvement of 2.3 points (p<0.001). "
        "No adverse events were recorded."
        # No explanation for the 15-subject drop — this is the seeded error
    ),
    "Discussion": (
        "The significant improvement supports the efficacy of the programme. "
        "Limitations include the urban-only sample and short follow-up period."
    ),
    "References": (
        "[1] WHO (2023) Global status report on the public health response to dementia.\n"
        "[2] Smith et al. (2021) Cognitive interventions in older adults. Lancet."
    ),
}

_TASK_C_ERRORS = [
    {
        "id": "err_C1",
        "error_type": "LOGICAL_INCONSISTENCY",
        "section": "Results",
        "hint": "85",
        "description": (
            "Methods states 100 subjects were recruited, but Results reports data "
            "for only 85 subjects with no explanation for the 15-subject discrepancy."
        ),
    }
]

_TASKS = [
    ("A", _TASK_A_PAPER, _TASK_A_ERRORS),
    ("B", _TASK_B_PAPER, _TASK_B_ERRORS),
    ("C", _TASK_C_PAPER, _TASK_C_ERRORS),
]

MANDATORY_SECTIONS = ["Abstract", "Introduction", "Methods", "Results", "Discussion", "References"]

# Reward constants
_EXPLORATION_REWARD = 0.02
_ACCURACY_REWARD = 0.30
_FALSE_POSITIVE_PENALTY = -0.10
_STEP_PENALTY = -0.01
_MAX_COMPLETION_BONUS = 0.40
_MIN_COMPLETION_BONUS = 0.10
_VIEW_MAX_CHARS = 1000


def _extract_citations(text: str) -> list[str]:
    """Return all citation markers like [1], [2], [12] found in text."""
    return list(dict.fromkeys(re.findall(r"\[\d+\]", text)))


class CiteguardianEnvironment(Environment):
    """
    CiteGuardian: a peer-review RL environment.

    The agent navigates a research paper, uses audit tools, flags errors,
    and submits when done. Rewards are cumulative and capped at 1.0.
    """

    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        self._state = State(episode_id=str(uuid4()), step_count=0)
        # These are set properly in reset()
        self._task_level: str = "A"
        self._paper: dict[str, str] = {}
        self._hidden_errors: list[dict] = []
        self._current_section: str = ""
        self._visited_sections: set[str] = set()
        self._flagged_errors: list[dict] = []  # what the agent has flagged
        self._cumulative_reward: float = 0.0
        self._done: bool = False
        self._audit_log: list[dict] = []

    # ------------------------------------------------------------------
    # Public interface
    # ------------------------------------------------------------------

    def reset(self) -> CiteguardianObservation:
        level, paper, errors = random.choice(_TASKS)
        self._task_level = level
        self._paper = dict(paper)
        self._hidden_errors = list(errors)
        self._current_section = list(self._paper.keys())[0]
        self._visited_sections = set()
        self._flagged_errors = []
        self._cumulative_reward = 0.0
        self._done = False
        self._audit_log = []
        self._state = State(episode_id=str(uuid4()), step_count=0)

        return self._make_observation(
            message=f"[Task {level}] Paper loaded. Begin your audit. "
                    f"Available sections: {list(self._paper.keys())}",
            tool_result=None,
        )

    def step(self, action: CiteguardianAction) -> CiteguardianObservation:  # type: ignore[override]
        if self._done:
            return self._make_observation(
                message="Episode already finished. Call reset() to start a new one.",
                tool_result=None,
            )

        self._state.step_count += 1
        # Step penalty every action
        self._cumulative_reward += _STEP_PENALTY

        atype = action.action_type
        tool_result = None
        message = ""

        if atype == "GO_TO":
            message, tool_result = self._handle_go_to(action)

        elif atype == "SCAN_CITATIONS":
            message, tool_result = self._handle_scan_citations()

        elif atype == "COMPARE_VALUES":
            message, tool_result = self._handle_compare_values(action)

        elif atype == "FLAG_ERROR":
            message, tool_result = self._handle_flag_error(action)

        elif atype == "SUBMIT":
            message, tool_result = self._handle_submit()

        else:
            message = f"Unknown action_type '{atype}'."

        self._audit_log.append(
            {"step": self._state.step_count, "action": atype, "detail": message}
        )

        obs = self._make_observation(message=message, tool_result=tool_result)
        obs.reward = round(self._cumulative_reward, 4)
        return obs

    @property
    def state(self) -> State:
        return self._state

    # ------------------------------------------------------------------
    # Action handlers
    # ------------------------------------------------------------------

    def _handle_go_to(self, action: CiteguardianAction):
        section = action.section_name or ""
        if section not in self._paper:
            return (
                f"Section '{section}' not found. "
                f"Available: {list(self._paper.keys())}",
                None,
            )
        self._current_section = section
        # Exploration reward — first visit only
        if section not in self._visited_sections:
            self._visited_sections.add(section)
            self._cumulative_reward += _EXPLORATION_REWARD
            extra = f" [+{_EXPLORATION_REWARD} exploration reward]"
        else:
            extra = ""
        return f"Navigated to '{section}'.{extra}", None

    def _handle_scan_citations(self):
        text = self._paper.get(self._current_section, "")
        citations = _extract_citations(text)
        return (
            f"SCAN_CITATIONS in '{self._current_section}': found {len(citations)} marker(s).",
            citations,
        )

    def _handle_compare_values(self, action: CiteguardianAction):
        v1 = str(action.val1 or "").strip()
        v2 = str(action.val2 or "").strip()
        try:
            n1, n2 = float(v1), float(v2)
            conflict = abs(n1 - n2) > 1e-9
        except ValueError:
            conflict = v1.lower() != v2.lower()

        result = {
            "val1": v1,
            "val2": v2,
            "conflict_detected": conflict,
        }
        msg = (
            f"COMPARE_VALUES: '{v1}' vs '{v2}' → "
            f"{'CONFLICT DETECTED' if conflict else 'no conflict'}."
        )
        return msg, result

    def _handle_flag_error(self, action: CiteguardianAction):
        etype = action.error_type or ""
        snippet = action.text_snippet or ""

        # Check against hidden errors
        matched = self._match_hidden_error(etype, snippet)

        if matched:
            if matched["id"] in [f["matched_id"] for f in self._flagged_errors if "matched_id" in f]:
                # Already flagged this one — treat as redundant false positive
                self._cumulative_reward += _FALSE_POSITIVE_PENALTY
                return (
                    f"FLAG_ERROR: '{etype}' — already flagged. "
                    f"Duplicate flag penalised ({_FALSE_POSITIVE_PENALTY}).",
                    None,
                )
            self._cumulative_reward += _ACCURACY_REWARD
            self._flagged_errors.append(
                {"error_type": etype, "snippet": snippet, "matched_id": matched["id"]}
            )
            return (
                f"FLAG_ERROR: '{etype}' — CORRECT. "
                f"Matched seeded error '{matched['id']}'. "
                f"[+{_ACCURACY_REWARD} accuracy reward]",
                None,
            )
        else:
            self._cumulative_reward += _FALSE_POSITIVE_PENALTY
            self._flagged_errors.append(
                {"error_type": etype, "snippet": snippet, "matched_id": None}
            )
            return (
                f"FLAG_ERROR: '{etype}' — FALSE POSITIVE. "
                f"No matching seeded error found. [{_FALSE_POSITIVE_PENALTY} penalty]",
                None,
            )

    def _handle_submit(self):
        self._done = True
        total_errors = len(self._hidden_errors)
        correct_flags = len([f for f in self._flagged_errors if f.get("matched_id")])
        false_positives = len([f for f in self._flagged_errors if not f.get("matched_id")])

        recall = correct_flags / total_errors if total_errors > 0 else 0.0

        # Completion bonus scales linearly with recall
        bonus = _MIN_COMPLETION_BONUS + recall * (_MAX_COMPLETION_BONUS - _MIN_COMPLETION_BONUS)
        self._cumulative_reward += bonus

        # Perfect run clamp
        if recall == 1.0 and false_positives == 0:
            self._cumulative_reward = 1.0
            verdict = "PERFECT AUDIT"
        else:
            self._cumulative_reward = min(self._cumulative_reward, 1.0)
            verdict = "AUDIT COMPLETE"

        msg = (
            f"SUBMIT — {verdict}. "
            f"Errors found: {correct_flags}/{total_errors}. "
            f"False positives: {false_positives}. "
            f"Final reward: {round(self._cumulative_reward, 4)}."
        )
        return msg, {
            "recall": recall,
            "false_positives": false_positives,
            "completion_bonus": round(bonus, 4),
            "final_reward": round(self._cumulative_reward, 4),
        }

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _match_hidden_error(self, etype: str, snippet: str) -> dict | None:
        """
        Try to match a FLAG_ERROR call against the hidden error list.
        Matching rules:
          - error_type must match exactly.
          - snippet must contain the error's hint string (case-insensitive).
        """
        for err in self._hidden_errors:
            if err["error_type"] != etype:
                continue
            hint = err.get("hint", "")
            if hint and hint.lower() not in snippet.lower():
                continue
            return err
        return None

    def _make_observation(self, message: str, tool_result) -> CiteguardianObservation:
        section_text = self._paper.get(self._current_section, "")
        view = section_text[:_VIEW_MAX_CHARS]

        # Build full-paper citation index for metadata
        all_citations: list[str] = []
        for text in self._paper.values():
            all_citations.extend(_extract_citations(text))
        all_citations = list(dict.fromkeys(all_citations))

        metadata = {
            "current_section": self._current_section,
            "available_sections": list(self._paper.keys()),
            "word_count": len(section_text.split()),
            "citation_markers_in_view": _extract_citations(section_text),
            "all_paper_citations": all_citations,
            "visited_sections": list(self._visited_sections),
            "flags_raised": len(self._flagged_errors),
            "step": self._state.step_count,
        }

        return CiteguardianObservation(
            current_view=view,
            metadata=metadata,
            audit_log=list(self._audit_log),
            tool_result=tool_result,
            message=message,
            task_level=self._task_level,
            done=self._done,
            reward=round(self._cumulative_reward, 4),
        )