models.py

from typing import Any, Dict, Optional
from openenv.core.env_server.types import Action, Observation
from pydantic import Field

class DataCleaningAction(Action):
    """Actions the agent can take to clean a dirty dataset."""

    operation: str = Field(
        ...,
        description=(
            "Cleaning operation to apply. One of: "
            "'impute_mean', 'impute_mode', 'drop_missing_rows', "
            "'remove_duplicates', 'fix_type_errors', "
            "'remove_outliers', 'normalize_text', 'fill_quantity_mean'"
        ),
    )
    column: Optional[str] = Field(
        default=None,
        description="Target column (optional). If omitted the op applies to all relevant columns.",
    )

class DataCleaningObservation(Observation):
    """The dataset state observed after each cleaning step."""

    current_text: str = Field(
        default="",
        description="Human-readable table of the current dataset rows.",
    )
    is_normalized: bool = Field(
        default=False,
        description="True when there are no missing values, duplicates, or outliers.",
    )
    html_found: bool = Field(
        default=False,
        description="Unused field kept for API compatibility (always False).",
    )
    remaining_typos: int = Field(
        default=0,
        description=(
            "Composite count of remaining issues: "
            "missing values + duplicate rows + outlier rows."
        ),
    )
    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description=(
            "Runtime metadata: quality_score, missing_count, has_duplicates, "
            "has_outliers, ops_already_applied, recommended_next, error."
        ),
    )