dataforge-07 0.1.0__py3-none-any.whl

dataforge/agent/tool_actions.py

@@ -0,0 +1,343 @@
+"""Typed tool-use action models for the DataForge RL environment.
+
+This module defines a discriminated union of 8 action types that an RL agent
+can submit to the DataForge environment. Each action is a standalone Pydantic
+model with its own validation rules, preventing cross-model field pollution.
+
+The ``parse_action`` function is the single entry point for HTTP handlers
+and tests to validate raw action dicts into typed models.
+
+Action Types:
+    INSPECT_ROWS  — View a slice of the dataset.
+    SQL_QUERY     — Execute read-only SQL against the episode DataFrame.
+    STAT_TEST     — Run a statistical test on a column.
+    PATTERN_MATCH — Evaluate a regex pattern against column values.
+    HYPOTHESIS    — Record a causal-root claim for credit.
+    ROOT_CAUSE    — Analyze selected detected errors for minimal roots.
+    DIAGNOSE      — Flag a suspected issue at (row, column).
+    FIX           — Propose a corrected value for a diagnosed issue.
+
+Example::
+
+    >>> from dataforge.agent.tool_actions import parse_action
+    >>> action = parse_action({"action_type": "INSPECT_ROWS", "row_indices": [0, 1]})
+    >>> action.action_type
+    'INSPECT_ROWS'
+"""
+
+from __future__ import annotations
+
+from typing import Annotated, Any, Literal
+
+from pydantic import BaseModel, Field, field_validator
+
+__all__ = [
+    "Action",
+    "Diagnose",
+    "Fix",
+    "Hypothesis",
+    "InspectRows",
+    "PatternMatch",
+    "RootCause",
+    "SqlQuery",
+    "StatTest",
+    "parse_action",
+]
+
+
+class InspectRows(BaseModel):
+    """View a slice of dataset rows.
+
+    Args:
+        action_type: Must be ``"INSPECT_ROWS"``.
+        row_indices: Zero-indexed row indices to retrieve. At least 1 required.
+        column_names: Optional column filter. If omitted, all columns returned.
+
+    Example::
+
+        >>> InspectRows(action_type="INSPECT_ROWS", row_indices=[0, 1, 2])
+    """
+
+    action_type: Literal["INSPECT_ROWS"]
+    row_indices: list[int] = Field(min_length=1, description="Row indices to inspect (0-indexed).")
+    column_names: list[str] | None = Field(default=None, description="Optional column filter.")
+
+    @field_validator("row_indices")
+    @classmethod
+    def _validate_row_indices(cls, v: list[int]) -> list[int]:
+        """Validate that all row indices are non-negative."""
+        if any(i < 0 for i in v):
+            raise ValueError("All row indices must be >= 0")
+        return v
+
+    model_config = {"frozen": True}
+
+
+class SqlQuery(BaseModel):
+    """Execute read-only SQL against the episode DataFrame via DuckDB.
+
+    Args:
+        action_type: Must be ``"SQL_QUERY"``.
+        query: SQL query string. Must be read-only (SELECT only).
+
+    Example::
+
+        >>> SqlQuery(action_type="SQL_QUERY", query="SELECT * FROM data LIMIT 5")
+    """
+
+    action_type: Literal["SQL_QUERY"]
+    query: str = Field(min_length=1, description="Read-only SQL query.")
+
+    model_config = {"frozen": True}
+
+
+class StatTest(BaseModel):
+    """Run a statistical test on a dataset column.
+
+    Args:
+        action_type: Must be ``"STAT_TEST"``.
+        test_type: One of ``"zscore"``, ``"iqr"``, ``"ks"``.
+        column: Column name to test.
+        threshold: Optional threshold override. Defaults vary by test type.
+
+    Example::
+
+        >>> StatTest(action_type="STAT_TEST", test_type="zscore", column="rating")
+    """
+
+    action_type: Literal["STAT_TEST"]
+    test_type: Literal["zscore", "iqr", "ks"] = Field(description="Statistical test to run.")
+    column: str = Field(min_length=1, description="Column name to test.")
+    threshold: float | None = Field(default=None, description="Optional threshold override.")
+
+    model_config = {"frozen": True}
+
+
+class PatternMatch(BaseModel):
+    """Evaluate a regex pattern against column values.
+
+    Args:
+        action_type: Must be ``"PATTERN_MATCH"``.
+        pattern: Regular expression string.
+        column: Column name to evaluate.
+        expect_match: If True, report rows that match. If False, report non-matches.
+
+    Example::
+
+        >>> PatternMatch(
+        ...     action_type="PATTERN_MATCH",
+        ...     pattern=r"^\\d{5}$",
+        ...     column="zip_code",
+        ... )
+    """
+
+    action_type: Literal["PATTERN_MATCH"]
+    pattern: str = Field(min_length=1, description="Regex pattern.")
+    column: str = Field(min_length=1, description="Column name to evaluate.")
+    expect_match: bool = Field(
+        default=True,
+        description="True to report matches, False to report non-matches.",
+    )
+
+    model_config = {"frozen": True}
+
+
+class Hypothesis(BaseModel):
+    """Record a causal-root claim for root-cause credit.
+
+    Args:
+        action_type: Must be ``"HYPOTHESIS"``.
+        claim: Textual description of the hypothesized root cause.
+        affected_rows: Row indices believed to be affected.
+        affected_columns: Column names believed to be affected.
+        root_cause_type: Detector-vocabulary root cause type
+            (e.g., ``"decimal_shift"``, ``"type_mismatch"``).
+
+    Example::
+
+        >>> Hypothesis(
+        ...     action_type="HYPOTHESIS",
+        ...     claim="Column 'rating' has a decimal shift at row 5",
+        ...     affected_rows=[5],
+        ...     affected_columns=["rating"],
+        ...     root_cause_type="decimal_shift",
+        ... )
+    """
+
+    action_type: Literal["HYPOTHESIS"]
+    claim: str = Field(min_length=1, description="Root-cause claim.")
+    affected_rows: list[int] = Field(min_length=1, description="Affected row indices.")
+    affected_columns: list[str] = Field(min_length=1, description="Affected column names.")
+    root_cause_type: str = Field(min_length=1, description="Detector-vocabulary root cause type.")
+
+    @field_validator("affected_rows")
+    @classmethod
+    def _validate_affected_rows(cls, v: list[int]) -> list[int]:
+        """Validate that all affected row indices are non-negative."""
+        if any(i < 0 for i in v):
+            raise ValueError("All affected row indices must be >= 0")
+        return v
+
+    model_config = {"frozen": True}
+
+
+class RootCause(BaseModel):
+    """Analyze selected detected errors for minimal causal roots.
+
+    Args:
+        action_type: Must be ``"ROOT_CAUSE"``.
+        error_indices: Zero-based indices into the episode's detected issue list.
+
+    Example::
+
+        >>> RootCause(action_type="ROOT_CAUSE", error_indices=[0, 1])
+    """
+
+    action_type: Literal["ROOT_CAUSE"]
+    error_indices: list[int] = Field(min_length=1, description="Detected issue indices.")
+
+    @field_validator("error_indices")
+    @classmethod
+    def _validate_error_indices(cls, v: list[int]) -> list[int]:
+        """Validate that all error indices are non-negative."""
+        if any(i < 0 for i in v):
+            raise ValueError("All error indices must be >= 0")
+        return v
+
+    model_config = {"frozen": True}
+
+
+class Diagnose(BaseModel):
+    """Flag a suspected data-quality issue at a specific (row, column).
+
+    Args:
+        action_type: Must be ``"DIAGNOSE"``.
+        row: Zero-indexed row number.
+        column: Column name.
+        issue_type: Issue type from detector vocabulary.
+
+    Example::
+
+        >>> Diagnose(
+        ...     action_type="DIAGNOSE",
+        ...     row=5, column="rating",
+        ...     issue_type="decimal_shift",
+        ... )
+    """
+
+    action_type: Literal["DIAGNOSE"]
+    row: int = Field(ge=0, description="Zero-indexed row number.")
+    column: str = Field(min_length=1, description="Column name.")
+    issue_type: str = Field(min_length=1, description="Issue type classification.")
+
+    model_config = {"frozen": True}
+
+
+class Fix(BaseModel):
+    """Propose a corrected value for a diagnosed issue.
+
+    Args:
+        action_type: Must be ``"FIX"``.
+        row: Zero-indexed row number.
+        column: Column name.
+        new_value: The corrected cell value as a string.
+        justification: Explanation of why this fix is correct.
+        fix_type: How to fix the issue. Defaults to ``"correct_value"``.
+
+    Example::
+
+        >>> Fix(
+        ...     action_type="FIX",
+        ...     row=5, column="rating",
+        ...     new_value="4.5",
+        ...     justification="Decimal shift: 45.0 should be 4.5",
+        ... )
+    """
+
+    action_type: Literal["FIX"]
+    row: int = Field(ge=0, description="Zero-indexed row number.")
+    column: str = Field(min_length=1, description="Column name.")
+    new_value: str = Field(description="Corrected cell value.")
+    justification: str = Field(min_length=1, description="Fix justification.")
+    fix_type: Literal["correct_value", "delete_row", "impute", "standardize"] = Field(
+        default="correct_value", description="Fix operation type."
+    )
+
+    model_config = {"frozen": True}
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Discriminated union and parser
+# ═══════════════════════════════════════════════════════════════════════════
+
+Action = Annotated[
+    InspectRows | SqlQuery | StatTest | PatternMatch | Hypothesis | RootCause | Diagnose | Fix,
+    Field(discriminator="action_type"),
+]
+"""Discriminated union of all valid DataForge environment actions."""
+
+
+def parse_action(raw: dict[str, Any]) -> Action:
+    """Parse and validate a raw action dict into the appropriate typed model.
+
+    This is the single entry point for HTTP handlers and tests to validate
+    actions. The ``action_type`` field is used as the discriminator.
+
+    Args:
+        raw: Dictionary with an ``action_type`` key and action-specific fields.
+
+    Returns:
+        A validated action model instance.
+
+    Raises:
+        pydantic.ValidationError: If the action is malformed or invalid.
+        KeyError: If ``action_type`` is missing.
+        ValueError: If ``action_type`` is not recognized.
+
+    Example::
+
+        >>> action = parse_action({"action_type": "INSPECT_ROWS", "row_indices": [0]})
+        >>> isinstance(action, InspectRows)
+        True
+    """
+    from pydantic import TypeAdapter
+
+    adapter: TypeAdapter[Action] = TypeAdapter(Action)
+    return adapter.validate_python(_normalize_action(raw))
+
+
+def _normalize_action(raw: dict[str, Any]) -> dict[str, Any]:
+    """Return a canonical action dictionary from supported external aliases."""
+    normalized = dict(raw)
+    action_type = normalized.get("action_type")
+    if action_type == "SQL_QUERY" and "sql" in normalized and "query" not in normalized:
+        normalized["query"] = normalized["sql"]
+    if action_type == "STAT_TEST" and "test" in normalized and "test_type" not in normalized:
+        normalized["test_type"] = normalized["test"]
+    if action_type == "PATTERN_MATCH":
+        if "regex" in normalized and "pattern" not in normalized:
+            normalized["pattern"] = normalized["regex"]
+        if "expect" in normalized and "expect_match" not in normalized:
+            normalized["expect_match"] = normalized["expect"] == "match"
+    if action_type == "HYPOTHESIS":
+        root_column = normalized.get("root_column")
+        downstream = normalized.get("downstream")
+        if root_column is not None and "affected_columns" not in normalized:
+            downstream_columns = downstream if isinstance(downstream, list) else []
+            normalized["affected_columns"] = [root_column, *downstream_columns]
+        if "affected_rows" not in normalized:
+            normalized["affected_rows"] = [0]
+        if root_column is not None and "root_cause_type" not in normalized:
+            normalized["root_cause_type"] = root_column
+    if (
+        action_type == "ROOT_CAUSE"
+        and "indices" in normalized
+        and "error_indices" not in normalized
+    ):
+        normalized["error_indices"] = normalized["indices"]
+    if action_type == "FIX":
+        if "proposed_value" in normalized and "new_value" not in normalized:
+            normalized["new_value"] = normalized["proposed_value"]
+        if "justification" not in normalized:
+            normalized["justification"] = "Agent proposed value via FIX."
+    return normalized

dataforge/bench/__init__.py

@@ -0,0 +1,31 @@
+"""Shared benchmark helpers for real-world DataForge evaluation."""
+
+from dataforge.bench.core import (
+    AggregateBenchmarkResult,
+    BenchmarkRepair,
+    BenchmarkRunOutput,
+    SeedBenchmarkResult,
+    chunk_row_indices,
+    estimate_llm_calls,
+    normalize_repairs,
+    quota_units,
+    score_repairs,
+    validate_estimated_calls,
+)
+from dataforge.bench.report import write_benchmark_outputs
+from dataforge.bench.runner import run_agent_comparison
+
+__all__ = [
+    "AggregateBenchmarkResult",
+    "BenchmarkRepair",
+    "BenchmarkRunOutput",
+    "SeedBenchmarkResult",
+    "chunk_row_indices",
+    "estimate_llm_calls",
+    "normalize_repairs",
+    "quota_units",
+    "run_agent_comparison",
+    "score_repairs",
+    "validate_estimated_calls",
+    "write_benchmark_outputs",
+]