env/models.py

from pydantic import BaseModel
from typing import Optional

class Action(BaseModel):
    answer: str
    suspicion_flags: list[dict]
    raw_text: str

class Observation(BaseModel):
    query: str
    context: list[str]
    step: int
    difficulty: str
    n_passages: int

class State(BaseModel):
    task_id: str
    injected_indices: list[int]
    injection_types: list[str]
    ground_truth: dict
    difficulty: str
    episode_id: str