VivienB/data-quality-auditing-env

1"""2models.py — Typed data models for the Data Quality Auditing OpenEnv.3Uses Pydantic v2 for strict validation.4"""5 6from __future__ import annotations7 8from enum import Enum9from typing import Any, Dict, List, Optional10 11from pydantic import BaseModel, Field12 13 14# ---------------------------------------------------------------------------15# Enumerations16# ---------------------------------------------------------------------------17 18class IssueSeverity(str, Enum):19    LOW = "low"20    MEDIUM = "medium"21    HIGH = "high"22    CRITICAL = "critical"23 24 25class IssueType(str, Enum):26    MISSING_VALUES = "missing_values"27    DUPLICATES = "duplicates"28    TYPE_MISMATCH = "type_mismatch"29    OUTLIER = "outlier"30    INCONSISTENCY = "inconsistency"31    BIAS = "bias"32 33 34class ActionType(str, Enum):35    VIEW_SCHEMA = "view_schema"36    SAMPLE_ROWS = "sample_rows"37    QUERY_COLUMN = "query_column"38    REPORT_ISSUE = "report_issue"39    SUBMIT_AUDIT = "submit_audit"40 41 42class TaskDifficulty(str, Enum):43    EASY = "easy"44    MEDIUM = "medium"45    HARD = "hard"46 47 48# ---------------------------------------------------------------------------49# Core domain models50# ---------------------------------------------------------------------------51 52class ColumnSchema(BaseModel):53    """Schema definition for a single dataset column."""54    name: str55    expected_type: str  # e.g. "int", "float", "str", "date"56    nullable: bool = True57    unique: bool = False58    min_value: Optional[float] = None59    max_value: Optional[float] = None60    allowed_values: Optional[List[Any]] = None61    description: Optional[str] = None62 63 64class DatasetSchema(BaseModel):65    """Full schema for the dataset being audited."""66    name: str67    description: str68    columns: List[ColumnSchema]69    primary_key: Optional[str] = None70 71 72class DataIssue(BaseModel):73    """A single data quality issue identified by the agent."""74    issue_type: IssueType75    column: Optional[str] = None76    description: str77    severity: IssueSeverity78    affected_rows: Optional[List[int]] = Field(default=None, description="Row indices affected")79    confidence: float = Field(default=1.0, ge=0.0, le=1.0)80 81 82class AuditReport(BaseModel):83    """The agent's final audit report."""84    task_id: str85    issues: List[DataIssue] = Field(default_factory=list)86    summary: Optional[str] = None87    total_issues: int = 088 89    def add_issue(self, issue: DataIssue) -> None:90        self.issues.append(issue)91        self.total_issues = len(self.issues)92 93 94# ---------------------------------------------------------------------------95# Action models96# ---------------------------------------------------------------------------97 98class Action(BaseModel):99    """An action taken by the agent."""100    action_type: ActionType101    params: Dict[str, Any] = Field(default_factory=dict)102 103 104# ---------------------------------------------------------------------------105# Observation / State models106# ---------------------------------------------------------------------------107 108class Observation(BaseModel):109    """The observation returned to the agent after each step."""110    dataset_sample: List[Dict[str, Any]] = Field(default_factory=list)111    schema_info: Optional[DatasetSchema] = None112    column_stats: Optional[Dict[str, Any]] = None113    current_report: AuditReport = Field(default_factory=lambda: AuditReport(task_id=""))114    steps_remaining: int = 20115    message: str = ""116    done: bool = False117 118 119class EnvironmentState(BaseModel):120    """Internal state of the environment (not directly exposed to agent)."""121    task_id: str122    difficulty: TaskDifficulty123    dataset: List[Dict[str, Any]]124    schema: DatasetSchema125    ground_truth_issues: List[DataIssue]126    current_report: AuditReport127    steps_taken: int = 0128    max_steps: int = 20129    done: bool = False130    total_reward: float = 0.0131 132 133# ---------------------------------------------------------------------------134# Task definition model135# ---------------------------------------------------------------------------136 137class Task(BaseModel):138    """A task loaded from a JSON task file."""139    task_id: str140    difficulty: TaskDifficulty141    description: str142    dataset: List[Dict[str, Any]]143    schema: DatasetSchema144    ground_truth_issues: List[DataIssue]145    max_steps: int = 20146    hints: Optional[List[str]] = None147