Skip to content

feat(annotation): persist n_retrieved_chunks on retrieval records#246

Merged
henrycgbaker merged 3 commits into
mainfrom
feat/annotation-retrieval-completeness/01-persist-k
Jul 6, 2026
Merged

feat(annotation): persist n_retrieved_chunks on retrieval records#246
henrycgbaker merged 3 commits into
mainfrom
feat/annotation-retrieval-completeness/01-persist-k

Conversation

@henrycgbaker

@henrycgbaker henrycgbaker commented Jun 1, 2026

Copy link
Copy Markdown
Collaborator

Stack (6 PRs): #246#247#268#248#269#249

Chain goal: Persist K (number of retrieved chunks per query) on retrieval records, surface it as panel-completeness columns + sidecar at export time, ship a live annotation status CLI on top, and add an opt-in --tag-partial-panels advisory write so annotators can filter partially-done panels in the Argilla UI.

This PR (1/6): Foundation. Stamps n_retrieved_chunks on retrieval record metadata at import and surfaces it on RetrievalAnnotation. Ships nothing user-visible on its own — the export columns (#247/#268) and live status CLI (#248/#269) consume it.


Scope

Record building

  • core/annotation/record_builder.py: add "n_retrieved_chunks": len(pair.chunks) to every chunk-record's metadata dict.

Task definition

  • core/annotation/argilla_task_definitions.py: declare IntegerMetadataProperty("n_retrieved_chunks", min=1, visible_for_annotators=False) on retrieval task settings.

Schema

  • core/schemas/annotation_export.py: RetrievalAnnotation.n_retrieved_chunks: int field (mirrors the existing chunk_rank: int).

Export readback

  • core/annotation/export_fetcher.py: _build_row retrieval branch reads metadata.get("n_retrieved_chunks", 0) (defensive default for records imported before this PR).

Tests

  • New test_n_retrieved_chunks_stamped_on_every_chunk_record in test_import.py locks the persist contract.
  • Metadata-name list in test_argilla_task_definitions.py updated.
  • Fixture refresh across test_annotation_export.py, test_export_runner.py, test_export_constraint_checks.py, test_iaa_runner.py, test_export_api.py.

Why

K varies per query (~15% K<5, 61% K=5, 24% K>5; the retriever is threshold-based, not top-K), so the export side can't infer K from chunk_rank or from record-counting alone — it needs a persisted SSOT to drive true top-K metrics in the consumer pipeline (precision@K, recall@K, NDCG@K, MRR@K).

End-to-end mirror of the existing chunk_rank field: stamped at import, declared on the task, modelled on RetrievalAnnotation, read back by the fetcher.

A one-off backfill script for records imported before this PR (in-flight batches) lives outside the supported CLI/API surface — retrospective ops migration, not packaged functionality.

Base

feat/querygen-resumable-stack — kitchen-sink integration branch carrying the full dep surface this stack assumes (Locale, atomic_write_text, LogicalConstraint, WIDGET_FIELD_PLACEHOLDERS, the export_constraint_checks rename). Rebases cleanly to main once those upstream stacks land.

Test plan

  • uv run python -m pytest tests/unit (863 passing)

K (= len(pair.chunks)) is now stamped as record metadata at import,
declared on the Argilla retrieval task, modelled on RetrievalAnnotation,
and read back by the export fetcher. Mirrors the existing chunk_rank
field end-to-end so retrieval metrics can compute true top-K and detect
incomplete panels.
@henrycgbaker henrycgbaker force-pushed the feat/annotation-retrieval-completeness/01-persist-k branch from b232f54 to 30b5118 Compare July 3, 2026 13:34
@github-actions github-actions Bot added size: XL 1000+ LOC and removed size: XS < 50 LOC labels Jul 3, 2026
@henrycgbaker henrycgbaker changed the base branch from feat/querygen-resumable-stack to main July 3, 2026 13:34
@henrycgbaker henrycgbaker requested a review from ddimmery July 3, 2026 13:45
@henrycgbaker henrycgbaker marked this pull request as ready for review July 3, 2026 13:45

@ddimmery ddimmery left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is already live, right? Otherwise there would need to be some kind of migration to deal with adding columns mid-run.

chunk_id=metadata.get("chunk_id", ""),
doc_id=metadata.get("doc_id", ""),
chunk_rank=metadata.get("chunk_rank", 0),
n_retrieved_chunks=metadata.get("n_retrieved_chunks", 0),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This zero is a sentinel, right? I ask because I assume downstream we're assuming k > 0 (e.g. in the task definition) and then use it as a denominator in metrics, so we'll need to watch it there.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes, 0 is sentinel - real values guaranteed ≥1 by Argilla schema (IntegerMetadataProperty(..., min=1)). I could change to None, only (non-major) issue is that this is close in the pipeline to csv/df export - so this would be adding non-int datatype which i think pandas coerces into NaN. In reality this is a non-issue as it's not reachable/edge case due to IntegerMetadataProperty

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe -1 to keep it an integer but make it very clear it's different?

@saschagobel saschagobel added the annotation Changes affecting the annotation tool label Jul 4, 2026
@github-actions github-actions Bot added size: XS < 50 LOC and removed size: XL 1000+ LOC labels Jul 4, 2026
@henrycgbaker

Copy link
Copy Markdown
Collaborator Author

This is already live, right? Otherwise there would need to be some kind of migration to deal with adding columns mid-run.

yes this is already live (I had a non-git tracked backfill script)

Addresses review feedback on #246: 0 could theoretically be confused
with a real count; -1 falls outside the valid domain (real K is always
>=1) and keeps the field an int rather than switching to None.
@henrycgbaker henrycgbaker merged commit 6c45733 into main Jul 6, 2026
4 checks passed
@henrycgbaker henrycgbaker deleted the feat/annotation-retrieval-completeness/01-persist-k branch July 6, 2026 13:02
henrycgbaker added a commit that referenced this pull request Jul 6, 2026
**Stack (6 PRs):** #246 (merged) → #247#268#248#269#249

**Chain goal:** Persist K (number of retrieved chunks per query) on
retrieval records, surface it as panel-completeness columns + sidecar at
export time, ship a live `annotation status` CLI on top, and add an
opt-in `--tag-partial-panels` advisory write so annotators can filter
partially-done panels in the Argilla UI.

**This PR (2/6):** The standalone, export-agnostic panel-completeness
calculator, plus the shared retrieval-walk primitive it and #268 both
consume. Split out of what was originally a single ~1000-line PR; wiring
completeness into the retrieval export path ships separately in #268.

---

## Scope

### Completeness aggregator
- New `core/annotation/completeness.py`: groups retrieval snapshots by
`record_uuid`, distinct-by-`chunk_id`, derives per-panel facts, emits
`CompletenessReport` (`by_uuid` + `summary`). Aggregation is split into
a typed `_PanelAccumulator` plus `_aggregate_snapshots` (grouping) /
`_summarize_groups` (transform) passes.

### Shared retrieval walk
- `core/annotation/export_fetcher.py`: `RetrievalRecordSnapshot` +
`walk_retrieval_records` — one Argilla scroll, no status filter (so
records lacking terminal responses are still seen for the integrity
check), feeding both this module's completeness pass and #268's row
emission. Plus `resolve_task_purposes`, a topology lookup extracted from
`fetch_task` (no behaviour change).

### Schema
- `core/schemas/annotation_export.py`: `CompletenessSummary` +
`KBucketStat` models for the export sidecar (populated by #268).

### Tests
- New `test_completeness.py`: happy path, distinct-by-chunk_id, orphan
exclusion, integrity warnings (records-vs-K + mixed-backfill within
panel), all-unknown-K warning, K-bucket cross-tab + per-K histogram,
STRICT `panel_complete` (discards do NOT count), `n_discarded_chunks`
subset of `n_annotated_chunks`, calibration walked when topology
declares it.

## Why STRICT `panel_complete`

The computed `panel_complete: bool` is STRICT: True iff every K chunk in
the panel has at least one **submitted** response. Discarded responses
are abstentions, not judgements (pragmata's `DiscardReason` enum is
refusal-only: `INVALID_OR_UNREALISTIC` / `UNCLEAR` /
`OUTSIDE_REVIEWER_EXPERTISE`), so treating them as covered would feed
unjudged chunks into NDCG@K / precision@K denominators as 0-relevance
labels. The bool is derivable from the count columns, so a consumer
wanting the permissive policy computes `n_annotated_chunks ==
n_retrieved_chunks` in one line.

The retriever is threshold-based, so dropping partial panels biases
metrics toward small K (MNAR); this module is the shared calculation
both the export path (#268) and the live status path (#248) build on —
one aggregator, two consumers.

## Test plan
- [x] `uv run python -m pytest tests/unit` (1040 passing)
henrycgbaker added a commit that referenced this pull request Jul 7, 2026
…ry) (#268)

**Stack (6 PRs):** #246#247#268#248#269#249

**This PR (3/6):** Split out of #247 (was 1001 lines / XL). Wires the
panel-completeness calculator (#247) into the retrieval export path.
Stacked on #247.

---

## Scope

### Schema
- `core/schemas/annotation_export.py`:
- `RetrievalExportRow` += `panel_complete: bool`, `n_annotated_chunks:
int`, `n_submitted_chunks: int`, `n_discarded_chunks: int`,
`n_records_seen: int`. All defaulted; appended at the tail so existing
column order is preserved.
- New `KBucketStat` (per-K-bucket panel counts) and
`CompletenessSummary` (aggregates + `by_k_bucket` cross-tab + full per-K
histogram) Pydantic models.
- `AnnotationExportMeta` += `completeness_summary: CompletenessSummary |
None` and `completeness_status: Literal["ok","failed","not_requested"]`.

### Single-walk retrieval pipeline
- `core/annotation/export_fetcher.py`: new `walk_retrieval_records` (one
Argilla scroll per dataset, no status filter — sees discards-only
records too so the integrity check has the full record set) +
`RetrievalRecordSnapshot` dataclass + `fetch_retrieval_from_records`
(status filter applied in Python at typed-row construction).
- `core/annotation/export_runner.py`: when `Task.RETRIEVAL in tasks`,
walks once and feeds both the typed-row projection and the completeness
aggregator. `compute_completeness` failures degrade the sidecar (sets
`completeness_status="failed"`) without losing the already-fetched
export.

### API
- `annotation/__init__.py`: lazy re-exports for `CompletenessSummary`,
`KBucketStat`.

### Tests
- `test_export_runner.py`: `write_export_csv` stamps completeness
columns by `record_uuid`; sidecar carries `completeness_summary`;
`completeness_status` ok/failed/not_requested set explicitly.
- `test_annotation_export.py`: `RetrievalExportRow` field-order locked.

## Why

The retriever is threshold-based, so dropping partial panels biases the
metrics toward small K (MNAR). Surfacing both per-row predicate columns
AND a bucketed sidecar aggregate lets the eval pipeline either filter
complete panels per-K or apply condensed-list scoring without
re-querying the data layer.

The single-walk shape (one retrieval scroll feeds both fetch_task and
compute_completeness) costs one extra in-memory pass over
already-fetched records but spares the Argilla server a second full
scroll per export — meaningful on the in-flight batches (~hundreds of
panels).

The STRICT `panel_complete` policy question (which shape to ship — this
PR assumes (A) from #247) is discussed on #247, since that's where the
predicate is actually computed; this PR just surfaces it as a column.

## Test plan
- [ ] `uv run python -m pytest
tests/unit/core/annotation/test_export_runner.py
tests/unit/core/schemas/test_annotation_export.py`
henrycgbaker added a commit that referenced this pull request Jul 7, 2026
…249)

**Stack (6 PRs):** #246#247#268#248#269#249

**This PR (6/6):** Adds the advisory `needs_completion` tag write onto
`annotation status`.

- `status --tag-partial-panels`: stamps `needs_completion` on the
**unresolved chunks of PARTIAL panels** (`0 < submitted < K`) and clears
stale tags, sharing the single status walk. Annotators filter
`needs_completion=true` in the Argilla UI to focus on the records that
will complete a partially-done panel.
- **Partial-only gate** (not "any incomplete panel"): a fully-unstarted
panel is never tagged, so the tag stays a selective mop-up signal
instead of ≈ everything pending.
- **Multi-domain + prod/cal-split aware**: a panel whose chunks are
split across a workspace's production + calibration datasets (per-item
calibration) is evaluated as one unit; writes route to each chunk's
owning dataset.

Reuses `metadata_ops` for the replace-safe upsert. Stacked on #269.

Also re-exports `TagResult` from the `pragmata.annotation` facade:
`StatusReport.tag_result` is a public field of this type, but the facade
export was pre-declared two PRs ago (before `panel_status.py` existed),
dropped once that PR's version of the module still didn't define the
class, and never re-added once this PR added it for real.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

annotation Changes affecting the annotation tool feature Adds or expands user-facing functionality size: XS < 50 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants