kumar6591/data-quality-env

1from __future__ import annotations2 3from typing import Any, Literal4 5from pydantic import BaseModel, Field6 7 8class FindingConfidence(BaseModel):9    """A single audit finding with agent-reported confidence."""10 11    value: Any12    confidence: float = Field(ge=0.0, le=1.0)13 14 15class AuditReport(BaseModel):16    """Structured audit report submitted by the agent."""17 18    null_issues: dict[str, FindingConfidence]19    duplicate_row_count: FindingConfidence20    schema_violations: list[dict[str, Any]]21    drifted_columns: list[str]22    drift_details: dict[str, FindingConfidence]23    relational_issues: list[dict[str, Any]]24    recommended_fixes: list[str]25 26 27class Action(BaseModel):28    action_type: Literal["query", "submit_report", "fix_sql"]29    sql: str | None = None30    report: AuditReport | None = None31 32 33class Observation(BaseModel):34    task_id: int35    task_description: str36    tables: dict[str, dict[str, str]]37    row_counts: dict[str, int]38    step: int39    max_steps: int40    query_credits_remaining: int41    phase: Literal["audit", "fix"]42    last_query_result: list[dict] | None43    last_action_error: str | None44    last_fix_score: float | None45 46 47class RewardBreakdown(BaseModel):48    base_audit_score: float49    confidence_brier_adjustment: float50    budget_efficiency_bonus: float51    fix_verification_bonus: float52    total: float53 54 55class Reward(BaseModel):56    value: float = Field(ge=-0.5, le=1.25)57    breakdown: RewardBreakdown58    done: bool59    info: dict[str, Any]60 61 62class EpisodeState(BaseModel):63    task_id: int64    seed: int65    step: int = 066    max_steps: int = 1267    query_credits: int = 1068    phase: Literal["audit", "fix"] = "audit"69    fix_steps_remaining: int = 370    report_submitted: bool = False71    done: bool = False72    gold_faults: dict[str, Any] = {}73    audit_score: float = 0.074    fix_bonus: float = 0.075