diff --git a/.gitignore b/.gitignore index 570b25a9..a40fa7c5 100644 --- a/.gitignore +++ b/.gitignore @@ -44,8 +44,14 @@ cache/ **/extraction_metadata_*.json **/phase1_candidates.json -# Database files +# Database files. The WAL sidecars hold uncommitted pages of the same private +# content — quotes, source text, hashed authors — while a run is in flight. *.db +*.db-wal +*.db-shm + +# Private per-probe databases (raw provider responses, source text, quotes) +data/probes/ # Local-only agent config + session notes (not code, not versioned in this repo) .cursorrules diff --git a/docs/second_pass_design.md b/docs/second_pass_design.md new file mode 100644 index 00000000..6a0d7f7c --- /dev/null +++ b/docs/second_pass_design.md @@ -0,0 +1,348 @@ +# Second-pass extraction — design + +Status: design agreed, not yet built. Supersedes the prototype in +`studies/psychedelics/extract_pharmacology.py`, which is a POC to be discarded. + +--- + +## 1. The feature + +Ask a follow-up question of a *filtered subset* of the corpus. + +> For posters who mention psilocybin — what dose, what effect, how long, what +> adverse events? + +> For posters who report POTS and tried LDN — what did they try first? + +The first pass (`variable_extraction/`, `src/run_sentiment_pipeline.py`) reads the +whole corpus and produces per-author fields and per-post drug sentiment. A second +pass starts from a *cohort* — defined by the results of the first pass, or by an +earlier second pass — re-reads the raw source text for those authors, and produces +structured, source-anchored **claims**. + +The unit of work is a **probe**: a declared question, scoped to a cohort, answered +against a schema, with its provenance and cost accounted for. + +--- + +## 2. What the POC proved + +`extract_pharmacology.py` is 2,360 lines. Roughly 300 are about psychedelics; the +rest is a hand-rolled second-pass engine. It ran a 25-pair pilot: 38 units, 271 +events, $0.17. + +It also stalled. Four sessions, 43 tests, a working pipeline — and zero validated +output, because the analyst-coding step ran through a CSV whose cells contained +embedded source text and a hand-authored `analyst_events_json` blob. The +extraction was never the bottleneck. The human loop was. + +Both halves of that are inputs to this design. + +--- + +## 3. Learnings ledger + +### Carried forward + +| # | Learning | Evidence | Lands as | +|---|---|---|---| +| L1 | Regex does keyword recall only. Whose experience it is, and whether the use happened, are model-labelled in the same LLM call. | A first-person regex gate silently destroyed 208 patient-drug pairs of recall. Three rounds of pattern narrowing moved groundability 29.8% → 36.5% and never resolved "I haven't tried ECT, but I have tried ketamine infusions." | Engine invariant | +| L2 | Extract the denominator. Plans, third-party reports, and unclear cases are emitted and labelled, not filtered at retrieval. | If nothing lands in `planned_or_considered`, the model isn't classifying — a signal you only get if those events exist. | Every claim schema declares an `included` predicate over its own labels | +| L3 | Per-field evidence anchors; paraphrase allowed. | Exact-substring validation was implemented and abandoned — models legitimately normalize. | Engine requires a non-empty quote per field, tied to a source window. No substring check. | +| L4 | Silence ≠ absence. | `reported` / `explicit_none` / `not_stated` — a boolean turns silence into a negative finding. | Provided base type probes compose | +| L5 | Split identity: *what work* (unit) vs *what work, judged how* (run). | A unit is the same work regardless of how its output is scored; the cache key is prompt-based, so the validator belongs to run identity and not to `unit_key`. | Engine | +| L6 | The cache key must cover everything that changes the answer. | A `**_ignored` kwarg swallow meant reasoning effort never reached the provider, and a max-effort call could be served a no-reasoning cached response. | Engine + a test that a dropped kwarg fails loudly | +| L7 | Log the provider response *before* validating it. | Otherwise truncations, malformed responses, and crash-before-commit vanish, and cost accounting under-reports. | Engine | +| L8 | `billing_uncertain` is a state. | A transport failure with no usage block is not free. | Engine | +| L9 | Sample at the analysis unit, never the compute unit. | Sampling batches made a reviewer code half of one person's history. | Engine samples cohort rows, returns all their units | +| L10 | Under reasoning models, cost is unknowable before a live call. | Pre-pilot projection was $1.18–$11.26; measured was 12,279 reasoning tokens/unit → ~$5.40. The estimate table was obsolete within a session of being written. | Lifecycle forces a *measured* re-projection. No cost tables in code. | +| L11 | Quotes plus hashed authors plus Reddit IDs are re-identifiable. | Study constraint. | Structural DB boundary, not a gitignore rule | +| L12 | Carrying a first-person anchor across sentences is a prompt problem, not a retrieval problem. | The POC's system prompt solves it in prose. | Prompt text survives the rewrite verbatim | +| L13 | Extraction is unstable: 36.5% field-set and 58.8% value agreement across two identical passes. | Corpus README. | Replication is a built-in measurement, not a checklist step. Cohort aggregates only; no per-person stories. | + +### Discarded + +| Prototype | Why | +|---|---| +| `EXPECTED_COHORT` hardcoded counts + drift assertion | A cohort hash detects drift better and needs no maintenance | +| Pinned price constants, three-row cost projection table | Obsolete within one session. Read pricing from provider metadata, record it, report measured spend | +| Provider routing pins, 65,536-token escalation ceiling | Chasing byte-determinism from a cheap reasoning endpoint that cannot provide it. Record what happened instead of pinning what should have | +| `source_units.json` *and* `.jsonl`; `cohort_status` *and* `cohort_status_final`; `finalize` rebuilding JSON from append-only ledgers | All of it is SQLite's job | +| Five CLI verbs with divergent guards | Three | +| `pilot_review.csv` — one row per unit, embedded source text, hand-authored JSON in a cell | The direct cause of the stall | + +--- + +## 4. Architecture + +### 4.1 A probe declares four things + +``` +probes/psychedelic_pharmacology/ + cohort.sql SELECT author_hash, :target AS target FROM ... WHERE ... + evidence.py anchored(fts=..., term=...) | author_window(budget=...) + claim.py Pydantic model + `included` predicate + prompt body + gates.toml thresholds, review sample size, replication rate +``` + +Everything else belongs to the engine: identity, unit batching, cache, transport, +mechanical validation, ledger, cost accounting, sampling, gate scoring, promotion. + +Target size: ~500 lines, against the POC's 2,360. Most of the difference is +SQLite absorbing the ledger, the artifact reconciliation, and `finalize`. + +### 4.2 Cohort — SQL over `patientpunk.db` + +The filters this feature exists to serve ("posters who mention X and report Y") +are joins across `treatment_reports`, `variables`, `conditions`, and `unified`. +A filter DSL would reimplement a fraction of SQL, worse. + +Contract: a single `SELECT`, executed read-only, returning `author_hash` and +optionally `target`. The resolved row set is hashed into run identity, so cohort +drift is detected as a changed hash rather than a hand-maintained count. + +This makes `load_db.py` (README Step 4) a prerequisite for every probe. The POC +filtered a records JSON directly; carrying that forward would mean maintaining two +filter surfaces permanently. + +**Chaining is the point.** A cohort may select from a prior probe's claims: + +```sql +SELECT author_hash FROM claim +WHERE probe = 'psychedelic_pharmacology' AND included = 1 AND ... +``` + +Upstream `run_id`s are recorded as `derived_from`, so a chain of passes is +traceable end to end. + +### 4.3 Evidence — two retrieval modes, declared + +- **`anchored`** — FTS candidates → term regex → paragraph ±1 windows around each + mention, deduped, with a stable `source_window_id`. For questions keyed on a + term. +- **`author_window`** — the author's text within a character budget, no keyword + gate. For questions that aren't ("what was their exercise history?"). + +Declaring the mode is mandatory. The two produce different denominators, and +conflating them yields "share of reports mentioning X" figures that silently mean +different things. + +Filtering before the LLM is limited to keyword recall and bot detection (L1). + +### 4.4 Claim — Pydantic per probe + +Cross-field invariants are the load-bearing part of a claim schema and cannot be +expressed in JSON Schema. In the POC these were +`outcomes_require_actual_self_use` and `adverse_fields_agree` — a plan or someone +else's report structurally cannot carry doses or outcomes. + +Division of responsibility: + +- **Engine — mechanical only.** JSON/schema validity, unknown keys, enum and range + errors, duplicate claims, `source_window_id` belongs to the unit and its + type/ID agree, every evidence quote non-empty, placeholder strings rejected + outside quotes. +- **Probe — semantic.** Its own `model_validator`s, over its own labels. + +The engine never second-guesses a model label. If the model says `subject=self`, +that stands, and the review pass measures whether it was right. + +**Prompt composition.** The engine supplies an invariant preamble: source text is +untrusted data inside ``, evidence-quote discipline, no placeholder +values, emit the denominator. The probe supplies the domain body. The POC's +system prompt — anchor carrying, stacked-attribution rules, modal verbs not making +a use hypothetical — is the most valuable artifact it produced and is copied +forward verbatim into the psychedelics probe body. + +### 4.5 Storage — SQLite, with a privacy boundary + +Private, gitignored, one file per probe: `data/probes/.db` + +```sql +unit(unit_key, run_id, author_hash, target, windows_json, character_count, ...) +attempt(run_id, unit_key, variant, transport_attempt, response_sha256, + input_tokens, output_tokens, reasoning_tokens, + estimated_cost, billing_uncertain, error, recorded_at) +claim(run_id, probe, author_hash, target, source_id, source_window_id, + claim_type, included, payload_json, evidence_json) +review(run_id, claim_rowid, agree, note, coded_at) +``` + +Shared, committed: `patientpunk.db` gets `probe_run` (run identity, spec hashes, +cohort hash, `derived_from`) and de-quoted aggregates that `unified` can roll up. + +Quotes never cross into `patientpunk.db`. L11 becomes a schema property rather +than a matter of discipline — a future analysis physically cannot join a quote +into a committed artifact. + +The four indexed columns on `claim` (`author_hash`, `target`, +`source_window_id`, `included`) are what the engine and gate scoring need. +Per-probe nested structure stays in `payload_json`; there is no attempt at one +relational schema across all probes. + +Using SQLite as the ledger removes `finalize` entirely: writes are transactional, +resume is a query, and there is no rebuild-from-JSONL step to reconcile. + +### 4.6 Lifecycle — three verbs + +```bash +uv run python -m probes plan +uv run python -m probes run --pilot --confirm-paid-run +uv run python -m probes review --export | --import | --score +``` + +- **`plan`** — resolve the cohort, build units, compute identity, project cost. + Never calls an LLM. +- **`run`** — one code path; `--pilot` is a sampling flag, not a separate command. + Always requires `--confirm-paid-run`. After a pilot it prints the *measured* + cost re-projection, replacing the estimate rather than sitting beside it. + `--replicate ` re-runs a stratified subset under a fresh cache key and + reports instability inline (L13). +- **`review`** — export a coding sample, import coded rows, score gates, and + promote to `patientpunk.db` only on a pass. + +### 4.7 Two tiers + +The full apparatus is correct for a study you intend to publish and fatal to "I +just want to ask a follow-up question about these posters." + +- **`probes ask --limit N`** — cached, no ledger, no gates, rows written + with `provisional = 1`. +- **`probes run `** — full run identity, ledger, pilot, gates. + +Hard rule: provisional rows are never promoted, and gate scoring refuses to read +them. Publishing means re-running under full identity. Exploration stays cheap and +the provenance guarantee stays absolute. + +--- + +## 5. Review — scope decision + +**An interactive review tool is out of scope.** The engine owns the export +format, the import format, gate scoring, and promotion; producing the coded data +is the analyst's own workflow. + +The export must not reproduce the POC's failure: + +- one row per **claim**, not per unit; +- source window text on the row; +- coding is `agree` / `disagree` columns filled with y/n; +- no JSON authoring in a spreadsheet cell. + +**Consequence — recall degrades, deliberately.** Precision gates survive at +spreadsheet scale: y/n per emitted claim. Recall requires the analyst to +enumerate what the model *missed*, which is exactly what forced +`analyst_events_json` into a CSV cell. Two options, unresolved: + +1. Measure recall on a much smaller window-level sample with a free-text "missed" + column, and report it as a documented estimate rather than a hard gate. + *(preferred)* +2. Gate on precision only, and state the limitation in the study. + +Gate scoring stays bag-of-claims within a source window rather than +claim-to-claim alignment — model and analyst legitimately split a window into +different numbers of claims. + +`--replicate` still supplies instability measurement without any human coding, +and for a corpus with 36.5% pass-to-pass field agreement that is the number that +matters most. + +--- + +## 6. Build order + +1. Probe DB schema + `plan` — cohort SQL → units → identity → cost projection. No + LLM anywhere in this step. +2. `run` — transport, cache, mechanical validation, attempt/claim ledger, cost + accounting. L6, L7, and L8 as explicit tests. +3. Review export/import + gate scoring + promotion to `patientpunk.db`. +4. Psychedelics rebuilt as a probe spec (~150 lines; prompt carried verbatim), + pilot re-run, coded, gated. This is the first time the POC's question actually + gets answered. +5. `ask` tier. +6. Chained cohorts — `derived_from` recorded, prior-probe claims selectable. + +--- + +## 7. Open decisions + +**Do the existing 271 pilot claims survive?** They live in +`data/psychedelics_pharmacology_pilot_coreweave_fp8_provider_workers10_evidence_quotes_20260807/`. +Under a new engine their run identity is meaningless, but they remain model output +over known source windows. Importing them as `provisional` gives step 3 a free +test corpus for the import path. The alternative is a clean $0.17 re-run. + +**Are gates a house default with per-probe override, or fully per-probe?** The POC +declared 100% grounding, ≥95% attribution and self-report, ≥90% dose/duration/AE. +A probe author setting their own passing grade is a conflict of interest, which +argues for a house default — but then one probe can be blocked by a bar it never +needed. + +**Recall gating** — option 1 or 2 in §5. + +--- + +## 8. Constraints that carry over unchanged + +From `studies/psychedelics/_handoff.txt` §2, and binding on every probe: + +- Self-selected Reddit reporting cohort, not a clinical cohort. +- No efficacy, causal, incidence, or dose-response-over-time claims. +- No chronology inferred from timestamps; duration must be explicitly stated. +- A rate is "share of extractable reports mentioning X," never incidence. Always + state the denominator. +- Silence is never "none." +- One person may report repeated or contradictory exposures; do not collapse them + into one summary. +- No collective "stack" outcome assigned to a single component unless the source + attributes it specifically. +- Quote-bearing artifacts stay private (§4.5) and are never committed. + +--- + +## 9. V1 data-model and storage boundary + +§4 describes the design in full. V1 builds a strict subset of it, so this section +records what `probes/models.py` and `probes/store.py` actually contain — the +sketch in §4.5 is the target, not the shipped schema. + +### In V1 + +Six tables in the gitignored private `data/probes/.db`: + +| Table | Holds | +|---|---| +| `probe_run` | immutable run identity: spec/cohort/source/unit-set hashes + `config_json` | +| `cohort_member` | the resolved SQL row set, one row per ordinal | +| `unit` | one bounded provider input, with its lifecycle status | +| `source_window` | private quote-bearing source text, checksummed | +| `attempt` | one provider response or transport failure, written before validation (L7) | +| `claim` | opaque `values_json` + `evidence_json`, keyed to a source window | + +The models are deliberately domain-blind: `Claim.values` is a `dict` the engine +never interprets. A probe validates its own semantics and normalizes into that +field, which is what keeps one engine serving unrelated questions (§4.4). + +Three differences from the §4.5 sketch, all intentional: + +- **`probe_run` lives in the private DB, not `patientpunk.db`.** Nothing is + promoted in V1, so a run identity in the shared database would have no reader. +- **`source_window` is its own table**, not a `windows_json` blob on `unit`. It is + the FK target for `claim`, which is how "this claim's window belongs to this + unit" becomes a database constraint instead of an engine check. +- **No `variant` column on `attempt`, and no `review` table.** Both belong to + features V1 excludes. + +`patientpunk.db` is an input to a probe and never an output. Quotes, raw +responses, source text, and hashed author IDs exist only behind +`data/probes/`, which satisfies L11 structurally. + +### Not in V1 + +Review/export, gate scoring, promotion to `patientpunk.db`, `provisional` mode +and the `ask` tier, `--replicate`, chained cohorts and `derived_from`, public +de-quoted aggregates, and migration of the POC's 271 pilot claims. §7's open +decisions stay open — V1 does not answer them. + +The engine (§4.6 `plan`/`run`) and the psychedelics probe (§4.1) follow in +later PRs on top of these contracts. diff --git a/probes/__init__.py b/probes/__init__.py new file mode 100644 index 00000000..62f699b3 --- /dev/null +++ b/probes/__init__.py @@ -0,0 +1,33 @@ +"""Reusable second-pass probe data contracts and private storage.""" + +from .models import ( + Attempt, + AttemptStatus, + Claim, + CohortMember, + EvidenceAnchor, + ProbeRun, + RunConfig, + SourceWindow, + StrictModel, + Unit, + UnitStatus, + Usage, +) +from .store import ProbeStore + +__all__ = [ + "Attempt", + "AttemptStatus", + "Claim", + "CohortMember", + "EvidenceAnchor", + "ProbeRun", + "ProbeStore", + "RunConfig", + "SourceWindow", + "StrictModel", + "Unit", + "UnitStatus", + "Usage", +] diff --git a/probes/models.py b/probes/models.py new file mode 100644 index 00000000..08d8b8dd --- /dev/null +++ b/probes/models.py @@ -0,0 +1,203 @@ +"""Shared data contracts for second-pass probes. + +The models in this module deliberately know nothing about a probe's domain. +Probe-specific values live in :class:`Claim.values`; the engine owns only +identity, provenance, and mechanical integrity checks. +""" + +from __future__ import annotations + +from datetime import datetime +from enum import StrEnum +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, model_validator + + +class StrictModel(BaseModel): + """Reject undeclared fields so provider output cannot silently drift.""" + + model_config = ConfigDict(extra="forbid", str_strip_whitespace=True) + + +class RawModel(BaseModel): + """Strict, but never rewrites strings. + + Models carrying private payloads must store exactly the bytes they were + given: a stripped body no longer matches its checksum, and stripped source + text no longer matches the corpus it was selected from. + """ + + model_config = ConfigDict(extra="forbid") + + +class UnitStatus(StrEnum): + """Lifecycle state of one bounded provider input.""" + + PLANNED = "planned" + RUNNING = "running" + COMPLETE = "complete" + FAILED = "failed" + + +class AttemptStatus(StrEnum): + """State of one provider response or transport attempt.""" + + # A response is written as received before JSON or claim validation. + RECEIVED = "received" + ACCEPTED = "accepted" + VALIDATION_FAILED = "validation_failed" + TRANSPORT_FAILED = "transport_failed" + + +class CohortMember(StrictModel): + """One ordered row returned by a probe's read-only cohort query.""" + + # This is a join key, not a Reddit username. Raw identity never enters the + # probe database. + author_hash: str = Field(min_length=1) + target: str | None = None + + +class SourceWindow(RawModel): + """Private source text supplied to a provider for one unit. + + ``source_window_id`` is stable for the same source identity and normalized + text. The text and source IDs are quote-bearing private data and must stay + in the per-probe database. + """ + + source_window_id: str = Field(min_length=1) + source_type: str = Field(min_length=1) + source_id: str = Field(min_length=1) + text: str = Field(min_length=1) + + +class RunConfig(StrictModel): + """Every request setting that can change a run's answer or cost.""" + + provider: str = Field(min_length=1) + base_url: str | None = None + model: str = Field(min_length=1) + temperature: float = Field(ge=0) + max_tokens: int = Field(gt=0) + reasoning_effort: str | None = None + service_tier: str | None = None + provider_routing: dict[str, Any] = Field(default_factory=dict) + evidence_config: dict[str, Any] = Field(default_factory=dict) + + +class ProbeRun(StrictModel): + """Immutable identity for one private probe database run. + + ``unit_key`` deliberately does not include validator identity; a unit is + the same requested work even when acceptance rules change. The enclosing + run identity records those rules through the probe/spec hashes. + """ + + run_id: str = Field(min_length=1) + probe: str = Field(min_length=1) + spec_hash: str = Field(min_length=1) + cohort_hash: str = Field(min_length=1) + source_fingerprint: str = Field(min_length=1) + unit_set_hash: str = Field(min_length=1) + config: RunConfig + created_at: datetime + + +class Unit(StrictModel): + """One bounded LLM input assembled from a cohort member's windows.""" + + unit_key: str = Field(min_length=1) + author_hash: str = Field(min_length=1) + target: str | None = None + windows: list[SourceWindow] = Field(min_length=1) + character_count: int = Field(gt=0) + status: UnitStatus = UnitStatus.PLANNED + + @model_validator(mode="after") + def windows_are_unique_and_counted(self) -> "Unit": + """Keep source windows distinct and account for exact input size.""" + window_ids = [window.source_window_id for window in self.windows] + if len(window_ids) != len(set(window_ids)): + raise ValueError("source_window_id must be unique within a unit") + if self.character_count != sum(len(window.text) for window in self.windows): + raise ValueError("character_count must equal selected window text") + return self + + +class Usage(StrictModel): + """Provider usage and known cost reported for one attempt.""" + + input_tokens: int = Field(default=0, ge=0) + output_tokens: int = Field(default=0, ge=0) + reasoning_tokens: int = Field(default=0, ge=0) + provider_cost: float | None = Field(default=None, ge=0) + + +class Attempt(RawModel): + """Auditable provider result recorded before claim validation. + + ``response_body`` is private raw provider output, including malformed JSON. + A transport failure may have no response body, but its billing uncertainty + must be explicit rather than being mistaken for a free request. + """ + + unit_key: str = Field(min_length=1) + attempt_no: int = Field(gt=0) + status: AttemptStatus + response_body: str | None = None + response_sha256: str | None = Field(default=None, min_length=1) + cache_key: str | None = None + usage: Usage | None = None + cache_hit: bool = False + billing_uncertain: bool = False + error: str | None = None + recorded_at: datetime + + @model_validator(mode="after") + def attempt_state_is_auditable(self) -> "Attempt": + """Require enough metadata to account for every response or failure.""" + if self.status != AttemptStatus.TRANSPORT_FAILED: + if self.response_body is None or self.response_sha256 is None: + raise ValueError("a provider response requires body and checksum") + if self.status == AttemptStatus.TRANSPORT_FAILED and not self.error: + raise ValueError("a transport failure requires an error") + if ( + self.status == AttemptStatus.TRANSPORT_FAILED + and self.usage is None + and not self.billing_uncertain + ): + raise ValueError("unknown transport billing must be marked uncertain") + return self + + +class EvidenceAnchor(StrictModel): + """One non-empty quote supporting a field in a probe claim.""" + + field_path: str = Field(min_length=1) + quote: str = Field(min_length=1) + + +class Claim(StrictModel): + """Generic source-anchored result emitted by a probe-specific schema. + + ``values`` is intentionally opaque to the engine. The probe validates its + semantic payload before normalizing it here; the engine only checks source + membership, duplicate claims, anchors, and placeholders. + """ + + claim_id: str = Field(min_length=1) + unit_key: str = Field(min_length=1) + source_window_id: str = Field(min_length=1) + included: bool + values: dict[str, Any] = Field(min_length=1) + evidence: list[EvidenceAnchor] = Field(min_length=1) + + @model_validator(mode="after") + def evidence_paths_are_unique(self) -> "Claim": + """Avoid competing evidence anchors for one declared field path.""" + paths = [anchor.field_path for anchor in self.evidence] + if len(paths) != len(set(paths)): + raise ValueError("each field_path may have only one evidence anchor") + return self diff --git a/probes/store.py b/probes/store.py new file mode 100644 index 00000000..e73be7bd --- /dev/null +++ b/probes/store.py @@ -0,0 +1,411 @@ +"""Private SQLite persistence for second-pass probe runs. + +``patientpunk.db`` is an input to a probe, never its output. This store keeps +raw responses, source text, quotes, and hashed author IDs behind the private +``data/probes/`` boundary so they cannot accidentally enter a committed +analysis database. +""" + +from __future__ import annotations + +import json +import sqlite3 +import threading +from contextlib import contextmanager +from datetime import datetime, timezone +from hashlib import sha256 +from pathlib import Path +from typing import Iterator + +from .models import ( + Attempt, + AttemptStatus, + Claim, + CohortMember, + ProbeRun, + Unit, + UnitStatus, + Usage, +) + + +SCHEMA = """ +CREATE TABLE IF NOT EXISTS probe_run ( + run_id TEXT PRIMARY KEY, + probe TEXT NOT NULL, + spec_hash TEXT NOT NULL, + cohort_hash TEXT NOT NULL, + source_fingerprint TEXT NOT NULL, + unit_set_hash TEXT NOT NULL, + config_json TEXT NOT NULL, + created_at TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS cohort_member ( + run_id TEXT NOT NULL REFERENCES probe_run(run_id), + ordinal INTEGER NOT NULL, + author_hash TEXT NOT NULL, + target TEXT, + PRIMARY KEY (run_id, ordinal) +); + +CREATE TABLE IF NOT EXISTS unit ( + run_id TEXT NOT NULL REFERENCES probe_run(run_id), + unit_key TEXT NOT NULL, + author_hash TEXT NOT NULL, + target TEXT, + character_count INTEGER NOT NULL CHECK (character_count > 0), + status TEXT NOT NULL, + PRIMARY KEY (run_id, unit_key) +); + +CREATE TABLE IF NOT EXISTS source_window ( + run_id TEXT NOT NULL, + unit_key TEXT NOT NULL, + source_window_id TEXT NOT NULL, + source_type TEXT NOT NULL, + source_id TEXT NOT NULL, + text TEXT NOT NULL, + text_sha256 TEXT NOT NULL, + PRIMARY KEY (run_id, unit_key, source_window_id), + FOREIGN KEY (run_id, unit_key) + REFERENCES unit(run_id, unit_key) +); + +CREATE TABLE IF NOT EXISTS attempt ( + run_id TEXT NOT NULL, + unit_key TEXT NOT NULL, + attempt_no INTEGER NOT NULL CHECK (attempt_no > 0), + status TEXT NOT NULL, + response_body TEXT, + response_sha256 TEXT, + cache_key TEXT, + usage_json TEXT, + cache_hit INTEGER NOT NULL DEFAULT 0, + billing_uncertain INTEGER NOT NULL DEFAULT 0, + error TEXT, + recorded_at TEXT NOT NULL, + -- ISO-8601 text does not sort by time across offsets, so cache lookup + -- orders on this normalized UTC value instead. + recorded_at_epoch REAL NOT NULL, + PRIMARY KEY (run_id, unit_key, attempt_no), + FOREIGN KEY (run_id, unit_key) + REFERENCES unit(run_id, unit_key) +); + +CREATE TABLE IF NOT EXISTS claim ( + run_id TEXT NOT NULL, + claim_id TEXT NOT NULL, + unit_key TEXT NOT NULL, + source_window_id TEXT NOT NULL, + included INTEGER NOT NULL CHECK (included IN (0, 1)), + values_json TEXT NOT NULL, + evidence_json TEXT NOT NULL, + PRIMARY KEY (run_id, claim_id), + FOREIGN KEY (run_id, unit_key, source_window_id) + REFERENCES source_window(run_id, unit_key, source_window_id) +); + +-- cohort_member and source_window are already covered by their primary keys. +CREATE INDEX IF NOT EXISTS idx_unit_run_status + ON unit(run_id, status); +CREATE INDEX IF NOT EXISTS idx_attempt_cache_status + ON attempt(cache_key, status); +CREATE INDEX IF NOT EXISTS idx_claim_lookup + ON claim(run_id, unit_key, source_window_id, included); +""" + + +def canonical_json(value: object) -> str: + """Serialize data consistently before storing or hashing it.""" + + return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + + +def text_sha256(text: str) -> str: + """Return the checksum used to detect changed private source text.""" + + return sha256(text.encode("utf-8")).hexdigest() + + +def _epoch(moment: datetime) -> float: + """Return a comparable UTC timestamp, reading naive values as UTC.""" + + if moment.tzinfo is None: + moment = moment.replace(tzinfo=timezone.utc) + return moment.timestamp() + + +class ProbeStore: + """Transactional store for one probe's private runs. + + The database path is intentionally probe-scoped rather than global. A + future public analysis can consume de-quoted aggregates without receiving + the source text, provider response, or evidence quotes stored here. + """ + + def __init__(self, path: Path) -> None: + self.path = Path(path) + self.path.parent.mkdir(parents=True, exist_ok=True) + # The engine dispatches provider calls concurrently, so the connection + # is shared across threads and serialized by `_lock` instead. + self.connection = sqlite3.connect(self.path, check_same_thread=False) + self._lock = threading.RLock() + self.connection.row_factory = sqlite3.Row + self.connection.execute("PRAGMA foreign_keys = ON") + self.connection.execute("PRAGMA journal_mode = WAL") + self.connection.executescript(SCHEMA) + self.connection.commit() + + @classmethod + def default_path(cls, probe: str, root: Path | None = None) -> Path: + """Return the gitignored default path for a probe database. + + The default is anchored to the repository root, not the working + directory: a probe run from a subdirectory would otherwise write its + WAL sidecars outside the ignored `data/` tree, where quote-bearing + uncommitted pages could be staged by `git add -A`. + """ + + base = root or Path(__file__).resolve().parent.parent + return base / "data" / "probes" / f"{probe}.db" + + @contextmanager + def transaction(self) -> Iterator[sqlite3.Connection]: + """Commit a group of writes atomically, rolling back on failure.""" + + with self._lock: + try: + yield self.connection + except Exception: + self.connection.rollback() + raise + else: + self.connection.commit() + + def close(self) -> None: + """Close the connection and release SQLite resources.""" + + self.connection.close() + + def __enter__(self) -> "ProbeStore": + return self + + def __exit__(self, *_: object) -> None: + self.close() + + def save_probe_run(self, run: ProbeRun) -> None: + """Insert an immutable run identity; reject accidental replacement.""" + + with self.transaction() as connection: + connection.execute( + """ + INSERT INTO probe_run ( + run_id, probe, spec_hash, cohort_hash, source_fingerprint, + unit_set_hash, config_json, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + run.run_id, + run.probe, + run.spec_hash, + run.cohort_hash, + run.source_fingerprint, + run.unit_set_hash, + canonical_json(run.config.model_dump(mode="json")), + run.created_at.isoformat(), + ), + ) + + def save_cohort_members( + self, run_id: str, members: list[CohortMember] + ) -> None: + """Persist the resolved SQL result in the order the caller supplies.""" + + with self.transaction() as connection: + connection.executemany( + """ + INSERT INTO cohort_member (run_id, ordinal, author_hash, target) + VALUES (?, ?, ?, ?) + """, + [ + (run_id, ordinal, member.author_hash, member.target) + for ordinal, member in enumerate(members) + ], + ) + + def save_units(self, run_id: str, units: list[Unit]) -> None: + """Persist units and their private source windows in one transaction.""" + + with self.transaction() as connection: + for unit in units: + connection.execute( + """ + INSERT INTO unit ( + run_id, unit_key, author_hash, target, + character_count, status + ) VALUES (?, ?, ?, ?, ?, ?) + """, + ( + run_id, + unit.unit_key, + unit.author_hash, + unit.target, + unit.character_count, + unit.status.value, + ), + ) + connection.executemany( + """ + INSERT INTO source_window ( + run_id, unit_key, source_window_id, source_type, + source_id, text, text_sha256 + ) VALUES (?, ?, ?, ?, ?, ?, ?) + """, + [ + ( + run_id, + unit.unit_key, + window.source_window_id, + window.source_type, + window.source_id, + window.text, + text_sha256(window.text), + ) + for window in unit.windows + ], + ) + + def set_unit_status( + self, run_id: str, unit_key: str, status: UnitStatus + ) -> None: + """Update one unit lifecycle state without rewriting its inputs.""" + + with self.transaction() as connection: + cursor = connection.execute( + """ + UPDATE unit SET status = ? + WHERE run_id = ? AND unit_key = ? + """, + (status.value, run_id, unit_key), + ) + # A silent no-op would leave the unit `running` forever, and resume + # is a query over unit status: the engine would re-dispatch it on + # every resume and pay for it again. + if cursor.rowcount != 1: + raise KeyError(f"no unit {unit_key!r} in run {run_id!r}") + + def record_attempt(self, run_id: str, attempt: Attempt) -> None: + """Write a response or transport failure before validating it.""" + + usage_json = ( + canonical_json(attempt.usage.model_dump(mode="json")) + if attempt.usage is not None + else None + ) + with self.transaction() as connection: + connection.execute( + """ + INSERT INTO attempt ( + run_id, unit_key, attempt_no, status, response_body, + response_sha256, cache_key, usage_json, cache_hit, + billing_uncertain, error, recorded_at, recorded_at_epoch + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + run_id, + attempt.unit_key, + attempt.attempt_no, + attempt.status.value, + attempt.response_body, + attempt.response_sha256, + attempt.cache_key, + usage_json, + int(attempt.cache_hit), + int(attempt.billing_uncertain), + attempt.error, + attempt.recorded_at.isoformat(), + _epoch(attempt.recorded_at), + ), + ) + + def update_attempt_status( + self, + run_id: str, + unit_key: str, + attempt_no: int, + status: AttemptStatus, + *, + error: str | None = None, + ) -> None: + """Advance a received attempt after validation succeeds or fails.""" + + with self.transaction() as connection: + cursor = connection.execute( + """ + UPDATE attempt SET status = ?, error = COALESCE(?, error) + WHERE run_id = ? AND unit_key = ? AND attempt_no = ? + """, + (status.value, error, run_id, unit_key, attempt_no), + ) + if cursor.rowcount != 1: + raise KeyError( + f"no attempt {attempt_no} for unit {unit_key!r} in run {run_id!r}" + ) + + def cached_attempt(self, cache_key: str) -> Attempt | None: + """Return the latest accepted response for a request cache key.""" + + with self._lock: + row = self.connection.execute( + """ + SELECT unit_key, attempt_no, status, response_body, + response_sha256, cache_key, usage_json, cache_hit, + billing_uncertain, error, recorded_at + FROM attempt + WHERE cache_key = ? AND status = ? + ORDER BY recorded_at_epoch DESC, rowid DESC + LIMIT 1 + """, + (cache_key, AttemptStatus.ACCEPTED.value), + ).fetchone() + if row is None: + return None + usage = Usage.model_validate(json.loads(row["usage_json"])) if row["usage_json"] else None + return Attempt( + unit_key=row["unit_key"], + attempt_no=row["attempt_no"], + status=AttemptStatus(row["status"]), + response_body=row["response_body"], + response_sha256=row["response_sha256"], + cache_key=row["cache_key"], + usage=usage, + cache_hit=bool(row["cache_hit"]), + billing_uncertain=bool(row["billing_uncertain"]), + error=row["error"], + recorded_at=datetime.fromisoformat(row["recorded_at"]), + ) + + def save_claim(self, run_id: str, claim: Claim) -> None: + """Persist one normalized claim after engine and probe validation.""" + + with self.transaction() as connection: + connection.execute( + """ + INSERT INTO claim ( + run_id, claim_id, unit_key, source_window_id, included, + values_json, evidence_json + ) VALUES (?, ?, ?, ?, ?, ?, ?) + """, + ( + run_id, + claim.claim_id, + claim.unit_key, + claim.source_window_id, + int(claim.included), + canonical_json(claim.values), + canonical_json( + [anchor.model_dump(mode="json") for anchor in claim.evidence] + ), + ), + )