dataforge/agent/controller.py

"""Verified autonomous agent controller for DataForge.

This is the production entry point that makes DataForge *truly agentic* without
weakening its moat. The control flow is:

1. **Deterministic-first seed.** Run detectors and the deterministic repairers
   (the high-accuracy floor). Their fixes are already safety+SMT verified by
   :func:`dataforge.engine.repair.propose_repairs`.
2. **Closed agent loop over the residual.** For issues the rules could not fix,
   an autonomous policy proposes actions. Every ``FIX`` is gated by the same
   safety constitution and SMT verifier; rejections are fed back so the policy
   self-corrects.
3. **Single verified commit.** Floor + agent fixes are committed through the
   existing :func:`dataforge.engine.repair.apply_transaction` — the same atomic,
   journaled, byte-for-byte reversible write path the CLI already uses.

Because the agent only *adds* verified fixes on top of the deterministic floor,
its output can never be worse than the deterministic baseline, and nothing
unverified ever reaches disk — regardless of how weak or adversarial the policy
is.
"""

from __future__ import annotations

from pathlib import Path

from pydantic import BaseModel, ConfigDict, Field

from dataforge.agent.executor import VerifiedActionExecutor
from dataforge.agent.policy import AgentObservation, Policy, ResidualIssue, make_policy
from dataforge.agent.scratchpad import Scratchpad
from dataforge.detectors import run_all_detectors
from dataforge.detectors.base import Issue, Schema
from dataforge.engine.repair import (
    CandidateFix,
    RepairMode,
    VerifiedFix,
    apply_transaction,
    propose_repairs,
)
from dataforge.repairers.base import ProposedFix, RepairAttempt
from dataforge.safety import SafetyContext, SafetyFilter, SafetyVerdict
from dataforge.table import (
    cell_value,
    column_names,
    copy_table,
    read_csv,
    row_count,
    set_cell_value,
)
from dataforge.transactions.log import sha256_bytes, sha256_file

__all__ = [
    "AgentActionRecord",
    "AgentRepairRequest",
    "AgentRepairResult",
    "run_agent_repair",
]

_SAMPLE_WINDOW = 3


class AgentRepairRequest(BaseModel):
    """Input contract for the verified agent repair controller."""

    source_path: Path
    mode: RepairMode = "dry_run"
    repair_schema: Schema | None = Field(default=None, alias="schema")
    policy: str = "hosted"
    provider: str | None = None
    max_steps: int = Field(default=30, ge=1, le=200)
    model: str | None = None
    temperature: float = Field(default=0.1, ge=0.0, le=2.0)
    allow_pii: bool = False
    confirm_pii: bool = False
    confirm_escalations: bool = False

    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid", populate_by_name=True)


class AgentActionRecord(BaseModel):
    """One step in the agent's audit trace."""

    step: int = Field(ge=1)
    action_type: str
    accepted: bool | None = None
    detail: str

    model_config = ConfigDict(frozen=True)


class AgentRepairResult(BaseModel):
    """Output contract for a verified agent repair run."""

    mode: RepairMode
    applied: bool
    reversible: bool = True
    source_path: str
    source_sha256: str = Field(pattern=r"^[0-9a-f]{64}$")
    post_sha256: str | None = Field(default=None, pattern=r"^[0-9a-f]{64}$")
    txn_id: str | None = None
    revert_command: str | None = None
    policy_name: str
    steps_used: int = Field(ge=0)
    max_steps: int = Field(ge=1)
    floor_fix_count: int = Field(ge=0)
    agent_fix_count: int = Field(ge=0)
    fixes_count: int = Field(ge=0)
    residual_count: int = Field(ge=0)
    issues_count: int = Field(ge=0)
    safety_verdict: str
    fixes: list[VerifiedFix] = Field(default_factory=list)
    trace: list[AgentActionRecord] = Field(default_factory=list)
    reason: str

    model_config = ConfigDict(frozen=True)


def _residual_issue(issue: Issue) -> ResidualIssue:
    """Project a detector Issue into the policy-facing residual record."""
    return ResidualIssue(
        row=issue.row,
        column=issue.column,
        issue_type=issue.issue_type,
        severity=issue.severity.value,
        expected=issue.expected,
        actual=issue.actual,
        reason=issue.reason,
    )


def _residual_issues(issues: list[Issue], attempt_groups: list[list[RepairAttempt]]) -> list[Issue]:
    """Return issues the deterministic floor did not accept a fix for."""
    residual: list[Issue] = []
    for issue, attempts in zip(issues, attempt_groups, strict=False):
        if not attempts or attempts[-1].status != "accepted":
            residual.append(issue)
    return residual


def _sample_rows(df: object, focus_row: int | None) -> tuple[dict[str, str], ...]:
    """Return a small window of rows around a focus row for the observation."""
    total = row_count(df)  # type: ignore[arg-type]
    if total == 0:
        return ()
    columns = column_names(df)  # type: ignore[arg-type]
    if focus_row is None:
        start, end = 0, min(total, _SAMPLE_WINDOW)
    else:
        start = max(0, focus_row - _SAMPLE_WINDOW)
        end = min(total, focus_row + _SAMPLE_WINDOW + 1)
    return tuple(
        {c: cell_value(df, i, c) for c in columns}  # type: ignore[arg-type]
        for i in range(start, end)
    )


def _build_observation(
    df: object,
    residual: dict[tuple[int, str], ResidualIssue],
    scratchpad: Scratchpad,
    last_result: str,
    steps_taken: int,
    max_steps: int,
    staged_count: int,
) -> AgentObservation:
    """Assemble the per-turn observation handed to the policy."""
    residual_list = tuple(residual.values())
    focus = residual_list[0].row if residual_list else None
    return AgentObservation(
        columns=tuple(column_names(df)),  # type: ignore[arg-type]
        row_count=row_count(df),  # type: ignore[arg-type]
        residual_issues=residual_list,
        sample_rows=_sample_rows(df, focus),
        scratchpad_summary=scratchpad.summary(),
        last_result=last_result,
        steps_taken=steps_taken,
        max_steps=max_steps,
        staged_fix_count=staged_count,
    )


def _verified_fix_payload(fix: ProposedFix, reason: str) -> VerifiedFix:
    """Build the public verified-fix payload from an accepted proposal."""
    return VerifiedFix(
        **CandidateFix.from_proposed(fix).model_dump(),
        verifier_reason=reason,
    )


def run_agent_repair(
    request: AgentRepairRequest,
    *,
    policy: Policy | None = None,
) -> AgentRepairResult:
    """Run the verified autonomous agent repair pipeline.

    Args:
        request: The repair request contract.
        policy: Optional pre-built policy (used by tests and callers that hold a
            backend). When omitted, a policy is constructed from
            ``request.policy`` with graceful fallback to deterministic.

    Returns:
        A frozen :class:`AgentRepairResult` describing what was verified,
        committed, and reverted-able.
    """
    source_path = request.source_path.resolve()
    source_bytes = source_path.read_bytes()
    source_sha256 = sha256_bytes(source_bytes)
    schema = request.repair_schema

    df = read_csv(source_path)
    issues = run_all_detectors(df, schema)

    # 1. Deterministic floor (already safety+SMT verified, no LLM).
    floor_fixes, attempt_groups = propose_repairs(
        issues,
        source_path,
        copy_table(df),
        schema,
        allow_llm=False,
        model=request.model,
        allow_pii=request.allow_pii,
        confirm_pii=request.confirm_pii,
        confirm_escalations=request.confirm_escalations,
        interactive=False,
    )

    # Rebuild the post-floor working table the agent reasons and verifies against.
    working_df = copy_table(df)
    for fix in floor_fixes:
        set_cell_value(working_df, fix.fix.row, fix.fix.column, fix.fix.new_value)

    residual_issues = _residual_issues(issues, attempt_groups)
    residual: dict[tuple[int, str], ResidualIssue] = {
        (issue.row, issue.column): _residual_issue(issue) for issue in residual_issues
    }

    safety_context = SafetyContext(
        allow_pii=request.allow_pii,
        confirm_pii=request.confirm_pii,
        confirm_escalations=request.confirm_escalations,
    )
    scratchpad = Scratchpad()

    active_policy = policy or make_policy(
        request.policy,
        model=request.model,
        temperature=request.temperature,
        provider=request.provider,
    )

    executor = VerifiedActionExecutor(
        working_df,
        schema,
        safety_context=safety_context,
        scratchpad=scratchpad,
        provenance=active_policy.provenance,
    )
    # Prevent the agent from re-touching cells the floor already fixed.
    for fix in floor_fixes:
        executor.mark_resolved(fix.fix.row, fix.fix.column)

    trace: list[AgentActionRecord] = []
    last_result = ""
    steps_used = 0

    initial_obs = _build_observation(
        working_df, residual, scratchpad, last_result, 0, request.max_steps, 0
    )
    active_policy.reset(initial_obs)

    # 2. Closed agent loop over the residual, with verified writes.
    for step in range(1, request.max_steps + 1):
        if not residual:
            break
        observation = _build_observation(
            working_df,
            residual,
            scratchpad,
            last_result,
            step - 1,
            request.max_steps,
            len(executor.staged_fixes),
        )
        action = active_policy.propose_action(observation)
        if action is None:
            break

        outcome = executor.execute(action)
        steps_used = step
        last_result = outcome.feedback
        trace.append(
            AgentActionRecord(
                step=step,
                action_type=outcome.action_type,
                accepted=outcome.accepted,
                detail=outcome.feedback[:500],
            )
        )
        if outcome.resolved_cell is not None:
            residual.pop(outcome.resolved_cell, None)

    agent_fixes = executor.staged_fixes
    all_fixes = [*floor_fixes, *agent_fixes]

    # 3. Batch safety gate (mirrors the deterministic pipeline: any non-ALLOW
    #    verdict voids the batch rather than shipping an inconsistent set).
    batch_safety = SafetyFilter().evaluate_batch(
        all_fixes, SafetyContext(confirm_escalations=request.confirm_escalations)
    )
    if batch_safety.verdict != SafetyVerdict.ALLOW:
        all_fixes = []
        agent_fixes = []
        floor_fixes = []

    applied = False
    txn_id: str | None = None
    post_sha256: str | None = None
    reason = "No accepted fixes were produced."

    if batch_safety.verdict != SafetyVerdict.ALLOW:
        reason = batch_safety.reason
    elif request.mode == "apply" and all_fixes:
        txn_id = apply_transaction(source_path, all_fixes, source_bytes)
        post_sha256 = sha256_file(source_path)
        applied = True
        reason = (
            f"Applied {len(all_fixes)} verified fix(es) "
            f"({len(floor_fixes)} deterministic, {len(agent_fixes)} agent)."
        )
    elif all_fixes:
        reason = (
            f"Dry run produced {len(all_fixes)} verified fix(es) "
            f"({len(floor_fixes)} deterministic, {len(agent_fixes)} agent); "
            "no source data was mutated."
        )

    fix_payloads = [
        _verified_fix_payload(fix, "Accepted by safety and SMT verifier.") for fix in all_fixes
    ]

    return AgentRepairResult(
        mode=request.mode,
        applied=applied,
        source_path=str(source_path),
        source_sha256=source_sha256,
        post_sha256=post_sha256,
        txn_id=txn_id,
        revert_command=f"dataforge revert {txn_id}" if txn_id is not None else None,
        policy_name=active_policy.name,
        steps_used=steps_used,
        max_steps=request.max_steps,
        floor_fix_count=len(floor_fixes),
        agent_fix_count=len(agent_fixes),
        fixes_count=len(all_fixes),
        residual_count=len(residual),
        issues_count=len(issues),
        safety_verdict=batch_safety.verdict.value,
        fixes=fix_payloads,
        trace=trace,
        reason=reason,
    )