feat(annotation): persist n_retrieved_chunks on retrieval records#246
Conversation
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.
b232f54 to
30b5118
Compare
ddimmery
left a comment
There was a problem hiding this comment.
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), |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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
There was a problem hiding this comment.
maybe -1 to keep it an integer but make it very clear it's different?
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.
**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)
…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`
…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.
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 statusCLI on top, and add an opt-in--tag-partial-panelsadvisory write so annotators can filter partially-done panels in the Argilla UI.This PR (1/6): Foundation. Stamps
n_retrieved_chunkson retrieval record metadata at import and surfaces it onRetrievalAnnotation. 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: declareIntegerMetadataProperty("n_retrieved_chunks", min=1, visible_for_annotators=False)on retrieval task settings.Schema
core/schemas/annotation_export.py:RetrievalAnnotation.n_retrieved_chunks: intfield (mirrors the existingchunk_rank: int).Export readback
core/annotation/export_fetcher.py:_build_rowretrieval branch readsmetadata.get("n_retrieved_chunks", 0)(defensive default for records imported before this PR).Tests
test_n_retrieved_chunks_stamped_on_every_chunk_recordintest_import.pylocks the persist contract.test_argilla_task_definitions.pyupdated.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_rankor 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_rankfield: stamped at import, declared on the task, modelled onRetrievalAnnotation, 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, theexport_constraint_checksrename). Rebases cleanly tomainonce those upstream stacks land.Test plan
uv run python -m pytest tests/unit(863 passing)