Spaces:
Sleeping
Sleeping
| """ | |
| environment.py | |
| ============== | |
| OpenEnv-compliant InformationIntegrityEnv. | |
| The environment is 100% offline and deterministic once a scenario has been | |
| collected by data_collector.py. It accepts any scenario dict produced by | |
| DataCollector.collect() — no hardcoded content. | |
| OpenEnv contract | |
| ---------------- | |
| env = InformationIntegrityEnv(scenario) | |
| obs = env.reset() | |
| obs, rew, done, info = env.step(action) | |
| state_dict = env.state() | |
| Action space | |
| ------------ | |
| verify_claim claim_id → VerificationResult | |
| check_source source_id → SourceEvidence | |
| decide decision (allow/flag/remove) → episode ends | |
| Observation | |
| ----------- | |
| content, claims (status masked), available_sources, | |
| interaction_history, status_flags, step_count | |
| Design invariants | |
| ----------------- | |
| - Ground truth NEVER exposed in Observation | |
| - All randomness eliminated (same action sequence → same result) | |
| - Reward computed by reward.py (configurable, not hardcoded here) | |
| - Scenario loaded once at __init__; reset() clears episode state only | |
| """ | |
| from __future__ import annotations | |
| import copy | |
| import os | |
| import time | |
| import uuid | |
| from typing import Any, Dict, List, Optional, Tuple | |
| import yaml | |
| from models import ( | |
| Action, | |
| ActionType, | |
| Claim, | |
| DecisionType, | |
| EpisodeResult, | |
| InteractionRecord, | |
| Observation, | |
| Reward, | |
| RewardBreakdown, | |
| Scenario, | |
| ScenarioClaim, | |
| ScenarioSource, | |
| SourceInfo, | |
| Stance, | |
| StatusFlags, | |
| VerificationStatus, | |
| ) | |
| from reward import EpisodeRewardState, RewardFunction | |
| try: | |
| from retrieval import RetrievalPipeline | |
| from reasoning import DeterministicReasoningEngine, reason_over_claims | |
| _RETRIEVAL_AVAILABLE = True | |
| except Exception: | |
| _RETRIEVAL_AVAILABLE = False | |
| RELIABILITY_CONFIG_PATH = os.getenv("RELIABILITY_CONFIG", "domain_reliability.yaml") | |
| # --------------------------------------------------------------------------- | |
| # Helpers | |
| # --------------------------------------------------------------------------- | |
| def _load_evidence_thresholds() -> dict: | |
| if os.path.exists(RELIABILITY_CONFIG_PATH): | |
| with open(RELIABILITY_CONFIG_PATH, "r", encoding="utf-8") as fh: | |
| cfg = yaml.safe_load(fh) or {} | |
| return cfg.get("evidence_thresholds", {}) | |
| return {} | |
| def _scenario_from_dict(d: dict) -> Scenario: | |
| """Deserialise a raw scenario dict (e.g. from DataCollector) into a Scenario model.""" | |
| claims = [ScenarioClaim(**c) for c in d["claims"]] | |
| sources = [ScenarioSource(**s) for s in d["sources"]] | |
| return Scenario( | |
| id=d.get("id", str(uuid.uuid4())), | |
| difficulty=d.get("difficulty", "medium"), | |
| content=d["content"], | |
| claims=claims, | |
| sources=sources, | |
| ground_truth=d["ground_truth"], | |
| topic=d.get("topic", ""), | |
| tags=d.get("tags", []), | |
| explanation_hint=d.get("explanation_hint", ""), | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # Main environment class | |
| # --------------------------------------------------------------------------- | |
| class InformationIntegrityEnv: | |
| """ | |
| OpenEnv environment for reasoning-based content moderation. | |
| Parameters | |
| ---------- | |
| scenario : dict | Scenario | |
| A collected scenario (from DataCollector) or a pre-built Scenario object. | |
| max_steps : int | |
| Maximum steps per episode (default 20). | |
| reward_fn : RewardFunction | None | |
| Pluggable reward function (default: RewardFunction()). | |
| """ | |
| # Public metadata (read by openenv validate) | |
| ENV_NAME = "InformationIntegrityEnv" | |
| ENV_VERSION = "1.0.0" | |
| ENV_AUTHOR = "openenv-content-moderation" | |
| def __init__( | |
| self, | |
| scenario: dict | Scenario, | |
| max_steps: int = 20, | |
| reward_fn: Optional[RewardFunction] = None, | |
| ) -> None: | |
| if isinstance(scenario, dict): | |
| scenario = _scenario_from_dict(scenario) | |
| self._scenario : Scenario = scenario | |
| self._max_steps : int = max_steps | |
| self._reward_fn : RewardFunction = reward_fn or RewardFunction() | |
| self._thresholds : dict = _load_evidence_thresholds() | |
| # Episode state (initialised by reset) | |
| self._episode_id : str = "" | |
| self._claims : List[Claim] = [] | |
| self._history : List[InteractionRecord] = [] | |
| self._flags : StatusFlags = StatusFlags() | |
| self._step_count : int = 0 | |
| self._done : bool = False | |
| self._final_decision: Optional[str] = None | |
| self._reward_state: EpisodeRewardState = EpisodeRewardState() | |
| # Index for fast lookup | |
| self._source_index : Dict[str, ScenarioSource] = { | |
| s.id: s for s in self._scenario.sources | |
| } | |
| self._claim_index : Dict[str, ScenarioClaim] = { | |
| c.id: c for c in self._scenario.claims | |
| } | |
| # Retrieval + reasoning engines (lazy init for pickle compat) | |
| self._retrieval_pipeline: Optional[Any] = None | |
| self._reasoning_engine: Optional[Any] = None | |
| self._last_reasoning: Optional[Any] = None # StructuredReasoning cached | |
| # ------------------------------------------------------------------ | |
| # OpenEnv API | |
| # ------------------------------------------------------------------ | |
| def reset( | |
| self, | |
| difficulty: Optional[str] = None, | |
| scenario_id: Optional[str | int] = None, | |
| ) -> Observation: | |
| """ | |
| Reset episode state and return initial observation. | |
| Programmatic difficulty control (no UI dependency): | |
| Parameters | |
| ---------- | |
| difficulty : str | None | |
| If provided, load a new scenario from the deterministic dataset | |
| for that difficulty tier ('easy', 'medium', 'hard'). | |
| Overrides the scenario supplied at __init__ time. | |
| scenario_id : str | int | None | |
| Index (int) or scenario id (str) within the difficulty dataset. | |
| If None and difficulty is provided, a random scenario is selected. | |
| Ignored when difficulty is None. | |
| """ | |
| if difficulty is not None: | |
| from tasks.registry import TASKS | |
| import random as _random | |
| task = TASKS[difficulty] | |
| if scenario_id is None: | |
| sc_dict = task.sample() | |
| else: | |
| sc_dict = task.get_scenario(scenario_id) | |
| self._scenario = _scenario_from_dict(sc_dict) | |
| self._source_index = {s.id: s for s in self._scenario.sources} | |
| self._claim_index = {c.id: c for c in self._scenario.claims} | |
| self._episode_id = str(uuid.uuid4()) | |
| self._step_count = 0 | |
| self._done = False | |
| self._final_decision = None | |
| self._history = [] | |
| self._flags = StatusFlags() | |
| self._reward_state = EpisodeRewardState() | |
| self._last_reasoning = None | |
| # Reinitialise mutable claim states (no ground-truth exposed) | |
| self._claims = [ | |
| Claim(id=c.id, text=c.text) | |
| for c in self._scenario.claims | |
| ] | |
| return self._build_observation() | |
| def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict[str, Any]]: | |
| """ | |
| Process one agent action. | |
| Returns | |
| ------- | |
| observation : Observation | |
| reward : Reward | |
| done : bool | |
| info : dict (extra diagnostics — not for agent use) | |
| """ | |
| if self._done: | |
| raise RuntimeError( | |
| "Episode has ended. Call reset() to start a new episode." | |
| ) | |
| self._step_count += 1 | |
| self._reward_state.step_count = self._step_count | |
| if action.type == ActionType.VERIFY_CLAIM: | |
| step_reward, result = self._handle_verify_claim(action) | |
| elif action.type == ActionType.CHECK_SOURCE: | |
| step_reward, result = self._handle_check_source(action) | |
| elif action.type == ActionType.DECIDE: | |
| step_reward, result = self._handle_decide(action) | |
| else: | |
| raise ValueError(f"Unknown action type: {action.type}") | |
| # Max-steps check | |
| if self._step_count >= self._max_steps and not self._done: | |
| self._done = True | |
| self._flags.max_steps_reached = True | |
| result["max_steps_reached"] = True | |
| # Update evidence sufficiency flag | |
| self._flags.sufficient_evidence = self._check_sufficient_evidence() | |
| # Log interaction | |
| record = InteractionRecord( | |
| step=self._step_count, | |
| action_type=action.type, | |
| action_payload=action.model_dump(exclude_none=True), | |
| result=result, | |
| reward_delta=step_reward, | |
| ) | |
| self._history.append(record) | |
| obs = self._build_observation() | |
| reward = self._build_reward(step_reward, result.get("message", "")) | |
| info = { | |
| "episode_id": self._episode_id, | |
| "step": self._step_count, | |
| "done": self._done, | |
| **result, | |
| } | |
| return obs, reward, self._done, info | |
| def state(self) -> Dict[str, Any]: | |
| """Return full internal state (for debugging / logging — not for agent).""" | |
| return { | |
| "episode_id": self._episode_id, | |
| "scenario_id": self._scenario.id, | |
| "difficulty": self._scenario.difficulty, | |
| "step_count": self._step_count, | |
| "done": self._done, | |
| "final_decision": self._final_decision, | |
| "ground_truth": self._scenario.ground_truth, # hidden from agent obs | |
| "claims": [ | |
| { | |
| "id": c.id, | |
| "text": c.text, | |
| "verification_status": c.verification_status, | |
| "confidence": c.confidence, | |
| } | |
| for c in self._claims | |
| ], | |
| "reward_breakdown": self._reward_fn.breakdown_dict(self._reward_state), | |
| "flags": self._flags.model_dump(), | |
| "history_len": len(self._history), | |
| "last_reasoning": ( | |
| {ct: r.to_dict() for ct, r in self._last_reasoning.items()} | |
| if self._last_reasoning else None | |
| ), | |
| } | |
| def episode_result(self, explanation: Optional[str] = None) -> EpisodeResult: | |
| """Build a final EpisodeResult (call after done=True).""" | |
| bd = self._reward_fn.breakdown_dict(self._reward_state) | |
| ground_truth = self._scenario.ground_truth | |
| correct = self._final_decision == ground_truth | |
| from tasks.graders import GraderFactory | |
| grader = GraderFactory.get(self._scenario.difficulty) | |
| grader_score = grader.score( | |
| decision=self._final_decision or "none", | |
| ground_truth=ground_truth, | |
| sources_checked=list(self._reward_state.checked_sources), | |
| claims_verified=list(self._reward_state.verified_claims), | |
| step_count=self._step_count, | |
| all_source_ids=[s.id for s in self._scenario.sources], | |
| all_claim_ids=[c.id for c in self._scenario.claims], | |
| ) | |
| return EpisodeResult( | |
| episode_id=self._episode_id, | |
| task_difficulty=self._scenario.difficulty, | |
| scenario_index=0, | |
| total_reward=self._reward_fn.final_score(self._reward_state), | |
| decision=self._final_decision, | |
| ground_truth=ground_truth, | |
| correct=correct, | |
| steps_taken=self._step_count, | |
| claims_verified=len(self._reward_state.verified_claims), | |
| sources_checked=len(self._reward_state.checked_sources), | |
| grader_score=grader_score, | |
| reward_breakdown=RewardBreakdown( | |
| verify_claim_reward=bd["verify_claim_reward"], | |
| check_source_reward=bd["check_source_reward"], | |
| multi_source_bonus=bd["multi_source_bonus"], | |
| decision_reward=bd["decision_reward"], | |
| premature_penalty=bd["premature_penalty"], | |
| redundant_penalty=bd["redundant_penalty"], | |
| inefficiency_penalty=bd["inefficiency_penalty"], | |
| aggression_penalty=bd["aggression_penalty"], | |
| ), | |
| explanation=explanation, | |
| ) | |
| # ------------------------------------------------------------------ | |
| # Action handlers | |
| # ------------------------------------------------------------------ | |
| def _handle_verify_claim(self, action: Action) -> Tuple[float, dict]: | |
| claim_id = action.claim_id | |
| if claim_id not in self._claim_index: | |
| return 0.0, {"error": f"Unknown claim_id: {claim_id}"} | |
| scenario_claim = self._claim_index[claim_id] | |
| agent_claim = self._get_agent_claim(claim_id) | |
| # Reveal verification status from scenario ground truth (masked from Observation) | |
| status = scenario_claim.ground_truth_status | |
| confidence = scenario_claim.ground_truth_confidence | |
| # Use RetrievalPipeline to find the most relevant sources for THIS claim | |
| if _RETRIEVAL_AVAILABLE: | |
| if self._retrieval_pipeline is None: | |
| self._retrieval_pipeline = RetrievalPipeline() | |
| raw_sources = [ | |
| { | |
| "id": s.id, | |
| "title": s.title, | |
| "text": s.text, | |
| "url": s.url, | |
| "stance": s.stance, | |
| "reliability": s.reliability, | |
| "source_type": s.source_type, | |
| } | |
| for s in self._scenario.sources | |
| ] | |
| ranked = self._retrieval_pipeline.retrieve(scenario_claim.text, raw_sources) | |
| relevant_sources = [s.id for s in ranked[:5]] | |
| else: | |
| relevant_sources = [ | |
| s.id for s in self._scenario.sources | |
| if s.stance in ("contradicting", "supporting") | |
| ] | |
| # Update mutable claim object | |
| agent_claim.verification_status = VerificationStatus(status) | |
| agent_claim.confidence = confidence | |
| agent_claim.verified_by_sources = relevant_sources[:3] | |
| # Reward | |
| r, msg = self._reward_fn.compute_verify_claim(claim_id, self._reward_state) | |
| result: dict = { | |
| "claim_id": claim_id, | |
| "status": status, | |
| "confidence": confidence, | |
| "message": msg, | |
| "relevant_sources": relevant_sources, | |
| "hint": f"Claim '{scenario_claim.text[:60]}…' assessed as {status}", | |
| } | |
| return r, result | |
| def _handle_check_source(self, action: Action) -> Tuple[float, dict]: | |
| source_id = action.source_id | |
| if source_id not in self._source_index: | |
| return 0.0, {"error": f"Unknown source_id: {source_id}"} | |
| src = self._source_index[source_id] | |
| # Reward | |
| r, msg = self._reward_fn.compute_check_source(source_id, self._reward_state) | |
| result: dict = { | |
| "source_id": source_id, | |
| "title": src.title, | |
| "text": src.text, | |
| "url": src.url, | |
| "stance": src.stance, | |
| "reliability": src.reliability, | |
| "source_type": src.source_type, | |
| "message": msg, | |
| } | |
| return r, result | |
| def _handle_decide(self, action: Action) -> Tuple[float, dict]: | |
| decision = action.decision | |
| ground_truth = self._scenario.ground_truth | |
| num_claims = len(self._scenario.claims) | |
| # Run deterministic reasoning engine if available | |
| reasoning_summary = {} | |
| if _RETRIEVAL_AVAILABLE: | |
| if self._retrieval_pipeline is None: | |
| self._retrieval_pipeline = RetrievalPipeline() | |
| if self._reasoning_engine is None: | |
| self._reasoning_engine = DeterministicReasoningEngine() | |
| raw_sources = [ | |
| { | |
| "id": s.id, | |
| "title": s.title, | |
| "text": s.text, | |
| "url": s.url, | |
| "stance": s.stance, | |
| "reliability": s.reliability, | |
| "source_type": s.source_type, | |
| } | |
| for s in self._scenario.sources | |
| ] | |
| claim_texts = [c.text for c in self._scenario.claims] | |
| claim_results = self._retrieval_pipeline.retrieve_for_all_claims( | |
| claim_texts, raw_sources | |
| ) | |
| agg_decision, per_claim_reasoning = reason_over_claims( | |
| claim_results, self._reasoning_engine | |
| ) | |
| # Cache reasoning for episode_result() | |
| self._last_reasoning = per_claim_reasoning | |
| reasoning_summary = { | |
| "engine_decision": agg_decision, | |
| "agent_decision": decision, | |
| "per_claim": { | |
| ct: r.to_dict() for ct, r in per_claim_reasoning.items() | |
| }, | |
| } | |
| r, msg = self._reward_fn.compute_decision( | |
| decision=decision, | |
| ground_truth=ground_truth, | |
| state=self._reward_state, | |
| num_claims=num_claims, | |
| ) | |
| self._final_decision = decision | |
| self._done = True | |
| self._flags.decision_made = True | |
| if "premature" in msg: | |
| self._flags.premature_decision = True | |
| result: dict = { | |
| "decision": decision, | |
| "message": msg, | |
| "episode_done": True, | |
| "reasoning": reasoning_summary, | |
| } | |
| return r, result | |
| # ------------------------------------------------------------------ | |
| # State construction helpers | |
| # ------------------------------------------------------------------ | |
| def _build_observation(self) -> Observation: | |
| """Build agent-facing Observation — ground truth is NEVER included.""" | |
| sources_view = [ | |
| SourceInfo( | |
| id=s.id, | |
| title=s.title, | |
| text=s.text, | |
| url=s.url, | |
| stance=Stance(s.stance), | |
| reliability=s.reliability, | |
| source_type=s.source_type, | |
| ) | |
| for s in self._scenario.sources | |
| ] | |
| return Observation( | |
| content=self._scenario.content, | |
| claims=copy.deepcopy(self._claims), | |
| available_sources=sources_view, | |
| interaction_history=list(self._history), | |
| status_flags=copy.deepcopy(self._flags), | |
| step_count=self._step_count, | |
| episode_id=self._episode_id, | |
| task_difficulty=self._scenario.difficulty, | |
| ) | |
| def _build_reward(self, score: float, message: str) -> Reward: | |
| bd = self._reward_fn.breakdown_dict(self._reward_state) | |
| return Reward( | |
| score=score, | |
| breakdown=RewardBreakdown( | |
| verify_claim_reward=bd["verify_claim_reward"], | |
| check_source_reward=bd["check_source_reward"], | |
| multi_source_bonus=bd["multi_source_bonus"], | |
| decision_reward=bd["decision_reward"], | |
| premature_penalty=bd["premature_penalty"], | |
| redundant_penalty=bd["redundant_penalty"], | |
| inefficiency_penalty=bd["inefficiency_penalty"], | |
| aggression_penalty=bd["aggression_penalty"], | |
| ), | |
| message=message, | |
| is_terminal=self._done, | |
| ) | |
| def _get_agent_claim(self, claim_id: str) -> Claim: | |
| for c in self._claims: | |
| if c.id == claim_id: | |
| return c | |
| raise KeyError(f"Claim {claim_id} not found in agent state") | |
| def _check_sufficient_evidence(self) -> bool: | |
| thresh = self._thresholds | |
| min_sources = int(thresh.get("min_sources_for_decision", 2)) | |
| min_verified = float(thresh.get("min_claims_verified_fraction", 0.5)) | |
| n_claims = len(self._scenario.claims) | |
| n_verified = len(self._reward_state.verified_claims) | |
| n_sources = len(self._reward_state.checked_sources) | |
| frac_ok = (n_verified / n_claims >= min_verified) if n_claims > 0 else True | |
| return n_sources >= min_sources and frac_ok | |
| # ------------------------------------------------------------------ | |
| # Convenience helpers | |
| # ------------------------------------------------------------------ | |
| def scenario(self) -> Scenario: | |
| return self._scenario | |
| def max_steps(self) -> int: | |
| return self._max_steps | |
| def done(self) -> bool: | |
| return self._done | |
| def from_dataset( | |
| cls, | |
| difficulty: str, | |
| scenario_id: Optional[str | int] = None, | |
| max_steps: int = 20, | |
| reward_fn: Optional[RewardFunction] = None, | |
| ) -> "InformationIntegrityEnv": | |
| """ | |
| Build an env directly from the deterministic dataset — no DataCollector needed. | |
| Parameters | |
| ---------- | |
| difficulty : 'easy' | 'medium' | 'hard' | |
| scenario_id : int index or str id; None = random selection | |
| max_steps : episode step limit | |
| reward_fn : optional custom RewardFunction | |
| Example | |
| ------- | |
| env = InformationIntegrityEnv.from_dataset("medium", scenario_id=2) | |
| obs = env.reset() | |
| """ | |
| from tasks.registry import TASKS | |
| task = TASKS[difficulty] | |
| if scenario_id is None: | |
| sc_dict = task.sample() | |
| else: | |
| sc_dict = task.get_scenario(scenario_id) | |
| return cls(scenario=sc_dict, max_steps=max_steps, reward_fn=reward_fn) | |
| def from_collector( | |
| cls, | |
| content: str, | |
| claims: List[str], | |
| difficulty: str = "medium", | |
| **collector_kwargs, | |
| ) -> "InformationIntegrityEnv": | |
| """ | |
| Convenience factory: collect a scenario on-the-fly and build an env. | |
| Example | |
| ------- | |
| env = InformationIntegrityEnv.from_collector( | |
| content="Post text here", | |
| claims=["Claim A", "Claim B"], | |
| difficulty="hard", | |
| ) | |
| obs = env.reset() | |
| """ | |
| from data_collector import CollectionRequest, DataCollector | |
| collector = DataCollector() | |
| request = CollectionRequest( | |
| content=content, | |
| claims=claims, | |
| difficulty=difficulty, | |
| **collector_kwargs, | |
| ) | |
| scenario_dict = collector.collect(request) | |
| return cls(scenario=scenario_dict) | |