File size: 5,592 Bytes
eed1cab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | """In-episode hypothesis and issue tracker for the DataForge RL agent.
The scratchpad is a mutable, episode-scoped data structure that the agent
uses to record hypotheses, confirmed issues, and dead ends. The environment
exposes a compact summary of the scratchpad in each observation, enabling
the agent to reason about its investigation history without direct access
to the underlying data structure.
Example::
>>> from dataforge.agent.scratchpad import Scratchpad
>>> pad = Scratchpad()
>>> pad.add_hypothesis("Rating column has decimal shift", [5], ["rating"], "decimal_shift")
>>> pad.confirm_issue(5, "rating", "decimal_shift")
>>> pad.summary()
'Hypotheses: 1 (0 pending). Confirmed: 1. Dead ends: 0.'
"""
from __future__ import annotations
from dataclasses import dataclass, field
__all__ = [
"ConfirmedIssue",
"DeadEnd",
"HypothesisRecord",
"Scratchpad",
]
@dataclass(frozen=True)
class HypothesisRecord:
"""A recorded hypothesis about a data-quality root cause.
Args:
claim: Textual description of the hypothesis.
affected_rows: Row indices the hypothesis covers.
affected_columns: Column names the hypothesis covers.
root_cause_type: Detector-vocabulary root cause type.
confirmed: Whether the hypothesis was confirmed by ground truth.
"""
claim: str
affected_rows: tuple[int, ...]
affected_columns: tuple[str, ...]
root_cause_type: str
confirmed: bool = False
@dataclass(frozen=True)
class ConfirmedIssue:
"""A confirmed data-quality issue at a specific location.
Args:
row: Zero-indexed row number.
column: Column name.
issue_type: Issue type classification.
"""
row: int
column: str
issue_type: str
@dataclass(frozen=True)
class DeadEnd:
"""A recorded dead end — an investigation path that yielded nothing.
Args:
description: What was tried and why it failed.
step_number: Step at which the dead end was recorded.
"""
description: str
step_number: int
@dataclass
class Scratchpad:
"""Mutable in-episode tracker for hypotheses, confirmed issues, and dead ends.
Reset at the start of each episode. The ``summary()`` method produces a
compact string for inclusion in agent observations.
Example::
>>> pad = Scratchpad()
>>> pad.add_hypothesis("Decimal shift in rating", [5], ["rating"], "decimal_shift")
>>> len(pad.hypotheses)
1
"""
hypotheses: list[HypothesisRecord] = field(default_factory=list)
confirmed_issues: list[ConfirmedIssue] = field(default_factory=list)
dead_ends: list[DeadEnd] = field(default_factory=list)
def add_hypothesis(
self,
claim: str,
affected_rows: list[int],
affected_columns: list[str],
root_cause_type: str,
) -> HypothesisRecord:
"""Record a new hypothesis.
Args:
claim: Textual description of the hypothesis.
affected_rows: Row indices the hypothesis covers.
affected_columns: Column names the hypothesis covers.
root_cause_type: Detector-vocabulary root cause type.
Returns:
The recorded hypothesis.
"""
record = HypothesisRecord(
claim=claim,
affected_rows=tuple(affected_rows),
affected_columns=tuple(affected_columns),
root_cause_type=root_cause_type,
)
self.hypotheses.append(record)
return record
def confirm_hypothesis(self, index: int) -> None:
"""Mark a hypothesis as confirmed.
Args:
index: Index into the ``hypotheses`` list.
Raises:
IndexError: If the index is out of range.
"""
old = self.hypotheses[index]
self.hypotheses[index] = HypothesisRecord(
claim=old.claim,
affected_rows=old.affected_rows,
affected_columns=old.affected_columns,
root_cause_type=old.root_cause_type,
confirmed=True,
)
def confirm_issue(self, row: int, column: str, issue_type: str) -> None:
"""Record a confirmed issue.
Args:
row: Zero-indexed row number.
column: Column name.
issue_type: Issue type classification.
"""
self.confirmed_issues.append(ConfirmedIssue(row=row, column=column, issue_type=issue_type))
def add_dead_end(self, description: str, step_number: int) -> None:
"""Record a dead end.
Args:
description: What was tried and why it failed.
step_number: Step at which the dead end was recorded.
"""
self.dead_ends.append(DeadEnd(description=description, step_number=step_number))
def reset(self) -> None:
"""Clear all tracked state for a new episode."""
self.hypotheses.clear()
self.confirmed_issues.clear()
self.dead_ends.clear()
def summary(self) -> str:
"""Produce a compact summary string for observation embedding.
Returns:
A one-line summary of scratchpad state.
Example::
>>> Scratchpad().summary()
'Hypotheses: 0 (0 pending). Confirmed: 0. Dead ends: 0.'
"""
pending = sum(1 for h in self.hypotheses if not h.confirmed)
return (
f"Hypotheses: {len(self.hypotheses)} ({pending} pending). "
f"Confirmed: {len(self.confirmed_issues)}. "
f"Dead ends: {len(self.dead_ends)}."
)
|