dataforge-07 0.1.0__py3-none-any.whl

dataforge/repair_contract.py

@@ -0,0 +1,468 @@
+"""Canonical prompt, parsing, and scoring contract for DataForge repairs."""
+
+from __future__ import annotations
+
+import json
+import re
+from collections import Counter, OrderedDict
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any, Literal, Protocol
+
+from pydantic import BaseModel, Field, ValidationError, model_validator
+
+CONTRACT_VERSION_V1 = "repair_contract_v1"
+CONTRACT_VERSION_V2 = "repair_contract_v2"
+CONTRACT_VERSION = CONTRACT_VERSION_V2
+
+SYSTEM_PROMPT = (
+    "You repair tabular data by proposing exact cell replacements. "
+    "Rows must be absolute row ids from valid_rows and columns must exactly match one of "
+    "the allowed_columns values. "
+    "Use only the provided dirty target rows and optional context rows. "
+    "Return strict JSON only in this object shape: "
+    '{"action":"submit_repairs","repairs":[{"row":0,"column":"Column",'
+    '"new_value":"value","reason":"why"}]}. '
+    'Use {"action":"finish","repairs":[]} when no cells should be changed. '
+    "Do not wrap the JSON in markdown code fences."
+)
+
+_JSON_FENCE_RE = re.compile(r"```(?:json)?\s*\n?(.*?)\n?\s*```", re.DOTALL)
+
+
+class RepairLike(Protocol):
+    """Minimal shape shared by repair objects across DataForge packages."""
+
+    @property
+    def row(self) -> int: ...
+
+    @property
+    def column(self) -> str: ...
+
+    @property
+    def new_value(self) -> str: ...
+
+    @property
+    def reason(self) -> str: ...
+
+
+class TruthLike(Protocol):
+    """Minimal shape shared by ground-truth cell objects across packages."""
+
+    @property
+    def row(self) -> int: ...
+
+    @property
+    def column(self) -> str: ...
+
+    @property
+    def clean_value(self) -> str: ...
+
+
+class RepairFix(BaseModel):
+    """One exact cell replacement proposed by a repair agent."""
+
+    row: int = Field(ge=0)
+    column: str = Field(min_length=1)
+    new_value: str
+    reason: str = Field(default="repair proposal", min_length=1)
+
+    model_config = {"frozen": True}
+
+
+class RepairAction(BaseModel):
+    """The only JSON action shape accepted by the repair contract."""
+
+    action: Literal["submit_repairs", "finish"]
+    repairs: list[RepairFix] = Field(default_factory=list)
+
+    model_config = {"frozen": True}
+
+    @model_validator(mode="after")
+    def _finish_must_be_empty(self) -> RepairAction:
+        if self.action == "finish" and self.repairs:
+            raise ValueError("finish actions must not include repairs")
+        return self
+
+
+class RepairParseResult(BaseModel):
+    """Parsed repair action plus diagnostics suitable for release gates."""
+
+    ok: bool
+    action: RepairAction | None = None
+    error_kind: (
+        Literal[
+            "parse_failure",
+            "truncated_json",
+            "schema_error",
+            "invalid_column",
+            "invalid_row",
+        ]
+        | None
+    ) = None
+    error_message: str | None = None
+    diagnostics: dict[str, int | str | bool] = Field(default_factory=dict)
+
+    model_config = {"frozen": True}
+
+
+class RepairScore(BaseModel):
+    """Exact-match cell repair metrics."""
+
+    tp: int = Field(ge=0)
+    fp: int = Field(ge=0)
+    fn: int = Field(ge=0)
+    precision: float = Field(ge=0.0, le=1.0)
+    recall: float = Field(ge=0.0, le=1.0)
+    f1: float = Field(ge=0.0, le=1.0)
+
+    model_config = {"frozen": True}
+
+
+def _as_jsonable_rows(rows: Sequence[Mapping[str, Any]]) -> list[dict[str, str]]:
+    """Return rows as stable string-valued mappings while preserving ``_row``."""
+    rendered: list[dict[str, str]] = []
+    for row in rows:
+        rendered_row: dict[str, str] = {}
+        for key, value in row.items():
+            rendered_row[str(key)] = str(value)
+        rendered.append(rendered_row)
+    return rendered
+
+
+def _valid_rows_from_target_rows(target_rows: Sequence[Mapping[str, Any]]) -> list[int]:
+    """Return absolute row ids from target rows, falling back to local ids for legacy rows."""
+    valid_rows: list[int] = []
+    for fallback_row, row in enumerate(target_rows):
+        raw_row = row.get("_row", fallback_row)
+        valid_rows.append(int(str(raw_row)))
+    return valid_rows
+
+
+def build_repair_user_payload(
+    *,
+    schema_summary: Mapping[str, Any],
+    target_rows: Sequence[Mapping[str, Any]],
+    context_rows: Sequence[Mapping[str, Any]] = (),
+    allowed_columns: Sequence[str],
+    valid_rows: Sequence[int] | None = None,
+    label_source: str | None = None,
+    dataset_note: str | None = None,
+    metadata: Mapping[str, Any] | None = None,
+    contract_version: str = CONTRACT_VERSION,
+) -> dict[str, Any]:
+    """Build the canonical user payload for repair SFT and evaluation."""
+    payload: dict[str, Any] = {
+        "contract_version": contract_version,
+        "schema_summary": dict(schema_summary),
+        "allowed_columns": list(allowed_columns),
+        "valid_rows": list(valid_rows)
+        if valid_rows is not None
+        else _valid_rows_from_target_rows(target_rows),
+        "target_rows": _as_jsonable_rows(target_rows),
+        "context_rows": _as_jsonable_rows(context_rows),
+    }
+    if label_source is not None:
+        payload["label_source"] = label_source
+    if dataset_note is not None:
+        payload["dataset_note"] = dataset_note
+    if metadata is not None:
+        payload["metadata"] = dict(metadata)
+    return payload
+
+
+def render_repair_messages(
+    *,
+    schema_summary: Mapping[str, Any],
+    target_rows: Sequence[Mapping[str, Any]],
+    allowed_columns: Sequence[str],
+    valid_rows: Sequence[int] | None = None,
+    context_rows: Sequence[Mapping[str, Any]] = (),
+    label_source: str | None = None,
+    dataset_note: str | None = None,
+    metadata: Mapping[str, Any] | None = None,
+    repairs: Sequence[RepairLike] | None = None,
+    contract_version: str = CONTRACT_VERSION,
+) -> list[dict[str, str]]:
+    """Render canonical chat messages for a repair task.
+
+    When ``repairs`` is ``None``, only system and user messages are returned.
+    When repairs are provided, an assistant message is appended for SFT.
+    """
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {
+            "role": "user",
+            "content": json.dumps(
+                build_repair_user_payload(
+                    schema_summary=schema_summary,
+                    target_rows=target_rows,
+                    context_rows=context_rows,
+                    allowed_columns=allowed_columns,
+                    valid_rows=valid_rows,
+                    label_source=label_source,
+                    dataset_note=dataset_note,
+                    metadata=metadata,
+                    contract_version=contract_version,
+                ),
+                sort_keys=True,
+                separators=(",", ":"),
+            ),
+        },
+    ]
+    if repairs is not None:
+        repair_fixes = [
+            RepairFix(
+                row=repair.row,
+                column=repair.column,
+                new_value=repair.new_value,
+                reason=repair.reason,
+            )
+            for repair in repairs
+        ]
+        messages.append(
+            {
+                "role": "assistant",
+                "content": json.dumps(
+                    RepairAction(
+                        action="submit_repairs" if repair_fixes else "finish",
+                        repairs=repair_fixes,
+                    ).model_dump(mode="json"),
+                    sort_keys=True,
+                    separators=(",", ":"),
+                ),
+            }
+        )
+    return messages
+
+
+def _strip_fence(text: str) -> str:
+    """Strip a single markdown JSON fence if the model returned one."""
+    stripped = text.strip()
+    match = _JSON_FENCE_RE.search(stripped)
+    return match.group(1).strip() if match else stripped
+
+
+def extract_json_payload(text: str) -> object:
+    """Extract the first complete JSON object or array from model text."""
+    clean_text = _strip_fence(text)
+    decoder = json.JSONDecoder()
+    saw_start = False
+    for offset, char in enumerate(clean_text):
+        if char not in "[{":
+            continue
+        saw_start = True
+        try:
+            payload, _end = decoder.raw_decode(clean_text[offset:])
+        except json.JSONDecodeError:
+            continue
+        if isinstance(payload, dict | list):
+            return payload
+    if saw_start:
+        raise ValueError("truncated_json")
+    raise ValueError("parse_failure")
+
+
+def _schema_case_error(column: str, allowed_columns: set[str]) -> bool:
+    """Return whether ``column`` only differs from an allowed column by case."""
+    return column.lower() in {allowed.lower() for allowed in allowed_columns}
+
+
+def parse_repair_action(
+    text: str,
+    *,
+    allowed_columns: Iterable[str] | None = None,
+    valid_rows: Iterable[int] | None = None,
+    require_explicit_action: bool = False,
+) -> RepairParseResult:
+    """Parse model text into a canonical repair action without raising.
+
+    By default this remains permissive enough to read legacy v1 artifacts. Pass
+    ``allowed_columns``, ``valid_rows``, and ``require_explicit_action=True`` for
+    the v2 release-gate contract.
+    """
+    try:
+        payload = extract_json_payload(text)
+    except ValueError as exc:
+        if str(exc) == "truncated_json":
+            return RepairParseResult(
+                ok=False,
+                error_kind="truncated_json",
+                error_message=str(exc),
+            )
+        return RepairParseResult(ok=False, error_kind="parse_failure", error_message=str(exc))
+
+    diagnostics: dict[str, int | str | bool] = {}
+    if isinstance(payload, list):
+        if require_explicit_action:
+            return RepairParseResult(
+                ok=False,
+                error_kind="schema_error",
+                error_message="repair payload must include an explicit action",
+            )
+        payload = {"action": "submit_repairs", "repairs": payload}
+    if not isinstance(payload, dict):
+        return RepairParseResult(
+            ok=False,
+            error_kind="schema_error",
+            error_message="repair payload must be a JSON object or array",
+        )
+    if "repairs" in payload and "action" not in payload:
+        if require_explicit_action:
+            return RepairParseResult(
+                ok=False,
+                error_kind="schema_error",
+                error_message="repair payload must include an explicit action",
+            )
+        payload = {**payload, "action": "submit_repairs"}
+    try:
+        action = RepairAction.model_validate(payload)
+    except ValidationError as exc:
+        return RepairParseResult(ok=False, error_kind="schema_error", error_message=str(exc))
+
+    normalized_repairs = normalize_fixes(action.repairs)
+    duplicate_count = len(action.repairs) - len(normalized_repairs)
+    if duplicate_count:
+        diagnostics["duplicate_cell_count"] = duplicate_count
+        action = RepairAction(action=action.action, repairs=normalized_repairs)
+
+    if allowed_columns is not None:
+        allowed = set(allowed_columns)
+        for repair in action.repairs:
+            if repair.column in allowed:
+                continue
+            diagnostics["invalid_column"] = repair.column
+            diagnostics["schema_case_error"] = _schema_case_error(repair.column, allowed)
+            return RepairParseResult(
+                ok=False,
+                error_kind="invalid_column",
+                error_message=f"column {repair.column!r} is not in allowed_columns",
+                diagnostics=diagnostics,
+            )
+
+    if valid_rows is not None:
+        rows = {int(row) for row in valid_rows}
+        for repair in action.repairs:
+            if repair.row in rows:
+                continue
+            diagnostics["invalid_row"] = repair.row
+            return RepairParseResult(
+                ok=False,
+                error_kind="invalid_row",
+                error_message=f"row {repair.row} is not in valid_rows",
+                diagnostics=diagnostics,
+            )
+
+    return RepairParseResult(ok=True, action=action, diagnostics=diagnostics)
+
+
+def normalize_fixes(fixes: Iterable[RepairLike]) -> list[RepairFix]:
+    """Collapse repairs to one final prediction per cell using last-write-wins."""
+    by_cell: OrderedDict[tuple[int, str], RepairFix] = OrderedDict()
+    for fix in fixes:
+        normalized = RepairFix(
+            row=fix.row,
+            column=fix.column,
+            new_value=fix.new_value,
+            reason=fix.reason,
+        )
+        key = (normalized.row, normalized.column)
+        if key in by_cell:
+            del by_cell[key]
+        by_cell[key] = normalized
+    return list(by_cell.values())
+
+
+def canonicalize_cell_value(value: str) -> str:
+    """Return a diagnostics-only canonical value for fuzzy F1 reporting."""
+    return " ".join(str(value).strip().casefold().split())
+
+
+def _strict_cell_value(value: str) -> str:
+    """Return the official exact-match value normalization."""
+    return str(value).rstrip()
+
+
+def score_repair_fixes(
+    ground_truth: Iterable[TruthLike],
+    fixes: Iterable[RepairLike],
+    *,
+    canonicalize_values: bool = False,
+) -> RepairScore:
+    """Score repairs by exact row, column, and string value match."""
+    normalized = normalize_fixes(fixes)
+    value_fn = canonicalize_cell_value if canonicalize_values else _strict_cell_value
+    expected = {(cell.row, cell.column): value_fn(str(cell.clean_value)) for cell in ground_truth}
+    matched: set[tuple[int, str]] = set()
+    tp = 0
+    fp = 0
+    for fix in normalized:
+        key = (fix.row, fix.column)
+        expected_value = expected.get(key)
+        if expected_value is not None and value_fn(fix.new_value) == expected_value:
+            tp += 1
+            matched.add(key)
+        else:
+            fp += 1
+    fn = len(expected) - len(matched)
+    precision = tp / (tp + fp) if (tp + fp) else 0.0
+    recall = tp / (tp + fn) if (tp + fn) else 0.0
+    f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 0.0
+    return RepairScore(
+        tp=tp,
+        fp=fp,
+        fn=fn,
+        precision=round(precision, 4),
+        recall=round(recall, 4),
+        f1=round(f1, 4),
+    )
+
+
+def score_repair_fixes_canonicalized(
+    ground_truth: Iterable[TruthLike],
+    fixes: Iterable[RepairLike],
+) -> RepairScore:
+    """Diagnostics-only F1 after conservative value canonicalization."""
+    return score_repair_fixes(ground_truth, fixes, canonicalize_values=True)
+
+
+def repair_failure_taxonomy(
+    *,
+    ground_truth: Iterable[TruthLike],
+    fixes: Iterable[RepairLike],
+    allowed_columns: Iterable[str],
+    valid_rows: Iterable[int],
+) -> dict[str, int]:
+    """Classify exact-match failures without changing official scoring."""
+    columns = set(allowed_columns)
+    lowercase_columns = {column.lower(): column for column in columns}
+    rows = set(valid_rows)
+    truth_map = {(cell.row, cell.column): str(cell.clean_value) for cell in ground_truth}
+    raw_fixes = list(fixes)
+    normalized_fixes = normalize_fixes(raw_fixes)
+    predictions = {(fix.row, fix.column): fix.new_value for fix in normalized_fixes}
+    counts: Counter[str] = Counter()
+    duplicate_count = len(raw_fixes) - len(normalized_fixes)
+    if duplicate_count:
+        counts["duplicate_cell"] += duplicate_count
+
+    for fix in normalized_fixes:
+        key = (fix.row, fix.column)
+        if fix.column not in columns:
+            if fix.column.lower() in lowercase_columns:
+                counts["schema_case_error"] += 1
+            else:
+                counts["wrong_cell"] += 1
+            continue
+        if fix.row not in rows:
+            counts["wrong_cell"] += 1
+            continue
+        if key not in truth_map:
+            counts["overrepair"] += 1
+            continue
+        if truth_map[key] != fix.new_value:
+            counts["wrong_value"] += 1
+
+    for key in truth_map:
+        if key not in predictions:
+            counts["missed_repair"] += 1
+    return {kind: count for kind, count in sorted(counts.items()) if count}

dataforge/repairers/__init__.py

@@ -0,0 +1,88 @@
+"""Repairer package for turning detected issues into proposed fixes."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from dataforge.detectors.base import Issue, Schema
+from dataforge.repairers.base import ProposedFix, RepairAttempt, Repairer, RetryContext
+from dataforge.repairers.decimal_shift import DecimalShiftRepairer
+from dataforge.repairers.fd_violation import FDViolationRepairer
+from dataforge.repairers.type_mismatch import TypeMismatchRepairer
+from dataforge.table import TableLike
+
+__all__ = [
+    "DecimalShiftRepairer",
+    "FDViolationRepairer",
+    "ProposedFix",
+    "RepairAttempt",
+    "Repairer",
+    "RetryContext",
+    "TypeMismatchRepairer",
+    "build_repairers",
+    "propose_fixes",
+]
+
+
+def build_repairers(
+    *,
+    cache_dir: Path | None,
+    allow_llm: bool,
+    model: str,
+) -> dict[str, Repairer]:
+    """Construct the default repairer registry."""
+    return {
+        "type_mismatch": TypeMismatchRepairer(),
+        "decimal_shift": DecimalShiftRepairer(),
+        "fd_violation": FDViolationRepairer(
+            cache_dir=cache_dir,
+            allow_llm=allow_llm,
+            model=model,
+        ),
+    }
+
+
+def propose_fixes(
+    issues: list[Issue],
+    df: TableLike,
+    schema: Schema | None,
+    *,
+    cache_dir: Path | None,
+    allow_llm: bool = False,
+    model: str = "gemini-2.0-flash",
+) -> list[ProposedFix]:
+    """Run all Week 2 repairers and return proposed fixes.
+
+    Args:
+        issues: Detected issues from the detector layer.
+        df: The input DataFrame being repaired.
+        schema: Optional declared schema.
+        cache_dir: Cache directory for any LLM-backed repair decisions.
+        allow_llm: Whether fd-violation repair may call the LLM provider.
+        model: The provider model name for fd-violation fallback.
+
+    Returns:
+        A deduplicated list of proposed fixes.
+    """
+    registry = build_repairers(
+        cache_dir=cache_dir,
+        allow_llm=allow_llm,
+        model=model,
+    )
+    proposed: list[ProposedFix] = []
+    seen_cells: set[tuple[int, str]] = set()
+
+    for issue in issues:
+        repairer = registry.get(issue.issue_type)
+        if repairer is None:
+            continue
+        fix = repairer.propose(issue, df, schema, retry_context=None)
+        if fix is None:
+            continue
+        key = (fix.fix.row, fix.fix.column)
+        if key in seen_cells:
+            continue
+        seen_cells.add(key)
+        proposed.append(fix)
+
+    return proposed

dataforge/repairers/base.py

@@ -0,0 +1,77 @@
+"""Core models and protocol for DataForge repairers."""
+
+from __future__ import annotations
+
+from typing import Literal, Protocol
+
+from pydantic import BaseModel, Field
+
+from dataforge.detectors.base import Issue, Schema
+from dataforge.table import TableLike
+from dataforge.transactions.txn import CellFix
+
+ProvenanceLiteral = Literal["deterministic", "llm_cache", "llm_live"]
+AttemptStatusLiteral = Literal[
+    "accepted",
+    "denied",
+    "escalated",
+    "rejected",
+    "unknown",
+    "attempted_not_fixed",
+]
+
+
+class ProposedFix(BaseModel):
+    """A repair proposal emitted by a repairer.
+
+    Args:
+        fix: The cell mutation to apply.
+        reason: Human-readable explanation of why this repair is proposed.
+        confidence: Repair confidence in the range [0.0, 1.0].
+        provenance: Where the proposed value came from.
+    """
+
+    fix: CellFix
+    reason: str = Field(min_length=1)
+    confidence: float = Field(ge=0.0, le=1.0)
+    provenance: ProvenanceLiteral
+
+    model_config = {"frozen": True}
+
+
+class RepairAttempt(BaseModel):
+    """Recorded outcome for one issue-repair attempt."""
+
+    issue: Issue
+    attempt_number: int = Field(ge=1)
+    fix: ProposedFix | None = None
+    status: AttemptStatusLiteral
+    reason: str = Field(min_length=1)
+    unsat_core: tuple[str, ...] = Field(default_factory=tuple)
+
+    model_config = {"frozen": True}
+
+
+class RetryContext(BaseModel):
+    """Hints passed back to a repairer after a failed attempt."""
+
+    issue: Issue
+    previous_attempts: tuple[RepairAttempt, ...] = Field(default_factory=tuple)
+    rejected_values: frozenset[str] = Field(default_factory=frozenset)
+    hints: tuple[str, ...] = Field(default_factory=tuple)
+
+    model_config = {"frozen": True}
+
+
+class Repairer(Protocol):
+    """Structural protocol implemented by every repairer."""
+
+    def propose(
+        self,
+        issue: Issue,
+        df: TableLike,
+        schema: Schema | None,
+        retry_context: RetryContext | None = None,
+    ) -> ProposedFix | None:
+        """Return a repair proposal for an issue, or ``None`` if unavailable."""
+        ...  # pragma: no cover

dataforge/repairers/decimal_shift.py

@@ -0,0 +1,43 @@
+"""Deterministic repairer for decimal-shift issues."""
+
+from __future__ import annotations
+
+from dataforge.detectors.base import Issue, Schema
+from dataforge.repairers.base import ProposedFix, RetryContext
+from dataforge.table import TableLike, cell_value, column_names, row_count
+from dataforge.transactions.txn import CellFix
+
+
+class DecimalShiftRepairer:
+    """Repair decimal-shift issues using the detector's expected value."""
+
+    def propose(
+        self,
+        issue: Issue,
+        df: TableLike,
+        schema: Schema | None,
+        retry_context: RetryContext | None = None,
+    ) -> ProposedFix | None:
+        """Return a deterministic fix for a decimal-shift issue."""
+        del schema, retry_context
+        if issue.issue_type != "decimal_shift" or issue.expected is None:
+            return None
+        if issue.row >= row_count(df) or issue.column not in column_names(df):
+            return None
+
+        old_value = cell_value(df, issue.row, issue.column)
+        if old_value == issue.expected:
+            return None
+
+        return ProposedFix(
+            fix=CellFix(
+                row=issue.row,
+                column=issue.column,
+                old_value=old_value,
+                new_value=issue.expected,
+                detector_id="decimal_shift",
+            ),
+            reason=issue.reason,
+            confidence=issue.confidence,
+            provenance="deterministic",
+        )