diff --git a/src/pragmata/annotation/__init__.py b/src/pragmata/annotation/__init__.py index b876a588..e28afc55 100644 --- a/src/pragmata/annotation/__init__.py +++ b/src/pragmata/annotation/__init__.py @@ -21,6 +21,7 @@ "Locale", "SetupResult", "StatusReport", + "TagResult", "Task", "UserSpec", "compute_iaa", @@ -39,6 +40,7 @@ "Locale": ("pragmata.core.schemas.annotation_task", "Locale"), "SetupResult": ("pragmata.core.annotation.setup", "SetupResult"), "StatusReport": ("pragmata.core.annotation.panel_status", "StatusReport"), + "TagResult": ("pragmata.core.annotation.panel_status", "TagResult"), "Task": ("pragmata.core.schemas.annotation_task", "Task"), "UserSpec": ("pragmata.core.settings.annotation_settings", "UserSpec"), "compute_iaa": ("pragmata.api.annotation_iaa", "compute_iaa"), @@ -75,6 +77,7 @@ def __dir__() -> list[str]: from pragmata.core.annotation.export_runner import ExportResult as ExportResult from pragmata.core.annotation.panel_status import HeadlineTotals as HeadlineTotals from pragmata.core.annotation.panel_status import StatusReport as StatusReport + from pragmata.core.annotation.panel_status import TagResult as TagResult from pragmata.core.annotation.setup import SetupResult as SetupResult from pragmata.core.schemas.annotation_task import Locale as Locale from pragmata.core.schemas.annotation_task import Task as Task diff --git a/src/pragmata/api/annotation_status.py b/src/pragmata/api/annotation_status.py index 1b1f0b57..3489e4b6 100644 --- a/src/pragmata/api/annotation_status.py +++ b/src/pragmata/api/annotation_status.py @@ -1,16 +1,23 @@ -"""Annotation status API - live per-panel completeness report. +"""Annotation status API - live progress + retrieval panel report + optional tag. -Pure read and config-free: resolves Argilla credentials, then walks the live -retrieval datasets (optionally narrowed by workspace) with no local topology -config. The optional ``--tag-partial-panels`` advisory write ships in a -follow-up PR so the read path carries no Argilla mutation surface. +Config-free: resolves Argilla credentials, then walks the live datasets (all +tasks; optionally narrowed by workspace) with no local topology config. +``tag_partial_panels=True`` additionally stamps the ``needs_completion`` +advisory tag on partial retrieval panels' unresolved chunks - the one live-write +surface, sharing the same retrieval walk as the panel report. """ import logging import os from pragmata.core.annotation.client import resolve_argilla_client -from pragmata.core.annotation.panel_status import StatusReport, compute_panel_status, compute_task_progress +from pragmata.core.annotation.panel_status import ( + StatusReport, + _apply_tags, + _build_report, + _collect_records, + compute_task_progress, +) from pragmata.core.settings.settings_base import UNSET, Unset, resolve_api_key logger = logging.getLogger(__name__) @@ -21,8 +28,9 @@ def report_status( api_url: str | Unset = UNSET, api_key: str | Unset = UNSET, workspace: str | None = None, + tag_partial_panels: bool = False, ) -> StatusReport: - """Fetch live retrieval panel status from Argilla (config-free). + """Fetch live annotation status from Argilla (config-free). Credential resolution (config-free): - ``api_url``: kwarg > ``ARGILLA_API_URL`` env @@ -32,16 +40,25 @@ def report_status( api_url: Argilla server URL. api_key: Argilla API key. workspace: If set, only datasets in this Argilla workspace. + tag_partial_panels: If True, stamp ``needs_completion`` on partial + panels' unresolved chunks (and clear stale tags). Opt-in live write. Returns: - ``StatusReport`` with the all-task ``progress`` summary plus the - retrieval per-panel facts. + ``StatusReport`` with the all-task ``progress`` summary and the + retrieval per-panel facts, plus an optional ``tag_result`` populated + when ``tag_partial_panels=True``. """ url = api_url if isinstance(api_url, str) else os.environ.get("ARGILLA_API_URL") key = api_key if isinstance(api_key, str) else resolve_api_key("argilla") client = resolve_argilla_client(url, key) + progress = compute_task_progress(client, workspace=workspace) - report = compute_panel_status(client, workspace=workspace).with_progress(progress) + # One retrieval walk shared between the panel report and the optional tag write. + collected = _collect_records(client, workspace=workspace) + report = _build_report(collected).with_progress(progress) + if tag_partial_panels: + report = report.with_tag_result(_apply_tags(collected)) + logger.info( "Status: %d panels, %d complete (%.0f%%), %d overlap-satisfied, %d integrity warnings", report.n_panels, diff --git a/src/pragmata/cli/commands/annotation.py b/src/pragmata/cli/commands/annotation.py index 17be0d34..d8ee9cdd 100644 --- a/src/pragmata/cli/commands/annotation.py +++ b/src/pragmata/cli/commands/annotation.py @@ -243,12 +243,23 @@ def status_command( ), by_workspace: bool = typer.Option(False, "--by-workspace", help="Add a per-workspace progress breakdown."), by_dataset: bool = typer.Option(False, "--by-dataset", help="Add a per-dataset progress breakdown."), + tag_partial_panels: bool = typer.Option( + False, + "--tag-partial-panels", + help=( + "Live write: stamp 'needs_completion' on the unresolved chunks of PARTIAL panels " + "(some but not all chunks annotated) and clear stale tags, so annotators can filter " + "straight to them in the Argilla UI. Off by default (read-only)." + ), + ), ) -> None: """Report live annotation progress across all tasks, plus retrieval panel-completeness. Config-free: walks every Argilla dataset. Record progress (total / completed) is shown per task; the retrieval row also carries panel-completeness. Add - --by-workspace / --by-dataset for finer breakdowns. + --by-workspace / --by-dataset for finer breakdowns. With --tag-partial-panels, + also stamps the 'needs_completion' advisory tag on partial panels' unresolved + chunks (a live write). """ from pragmata import annotation @@ -256,6 +267,7 @@ def status_command( api_url=UNSET if api_url is None else api_url, api_key=UNSET if api_key is None else api_key, workspace=workspace, + tag_partial_panels=tag_partial_panels, ) def _num(n: int) -> str: @@ -313,6 +325,11 @@ def _pct(done: int, total: int) -> str: ) if report.n_orphans_skipped: typer.echo(f"orphans skipped: {report.n_orphans_skipped} record(s) with empty record_uuid") + if report.tag_result is not None: + tr = report.tag_result + typer.echo( + f"tag-partial-panels: tagged={tr.n_tagged} cleared={tr.n_cleared} already_tagged={tr.n_already_tagged}" + ) @annotation_app.command("iaa") diff --git a/src/pragmata/core/annotation/metadata_ops.py b/src/pragmata/core/annotation/metadata_ops.py new file mode 100644 index 00000000..9d9a2aa0 --- /dev/null +++ b/src/pragmata/core/annotation/metadata_ops.py @@ -0,0 +1,76 @@ +"""Shared safe metadata operations for live Argilla mutations. + +Used by the ``--tag-partial-panels`` write path in ``panel_status``. +Centralises the two safety invariants that every metadata write must respect +on Argilla v2.8.0: + +1. **Argilla metadata is REPLACE, not merge.** Every ``dataset.records.log`` + call replaces the record's metadata wholesale. To avoid clobbering + existing keys, always fetch the current dict, merge in the update, and + send the FULL resulting dict. +2. **Property declaration is additive and idempotent.** Adding a metadata + property to an existing dataset is non-destructive, but the SDK raises + if the property already exists (with override warning). Skip the add + when the property is already present. + +Writes go via ``rg.Record(id=..., metadata={...})`` rather than a raw dict +payload: the SDK's ``IngestedRecordMapper`` flattens dict keys against the +dataset schema, so a ``{"id": ..., "metadata": {...}}`` shape would treat +"metadata" as an unknown top-level attribute and silently send an empty +metadata dict (wiping the record). Passing an ``rg.Record`` bypasses the +mapper. +""" + +import logging +from collections.abc import Iterable, Mapping + +import argilla as rg + +logger = logging.getLogger(__name__) + + +def ensure_metadata_property(dataset: rg.Dataset, prop: rg.MetadataType) -> bool: + """Idempotently declare ``prop`` on ``dataset``. + + Returns True if the property was newly added (and the dataset settings + pushed to the server), False if it was already present. + """ + existing = dataset.settings.metadata[prop.name] + if existing is not None: + return False + dataset.settings.add(prop) + dataset.settings.update() + logger.info("Declared metadata property %r on dataset %s", prop.name, dataset.name) + return True + + +def build_metadata_upsert( + record: rg.Record, + updates: Mapping[str, object], + *, + remove_keys: Iterable[str] = (), +) -> rg.Record | None: + """Merge ``updates`` into ``record.metadata`` and return an upsert Record. + + Returns ``None`` when the merge produces no change (idempotent no-op). + Mutates ``record``'s metadata in place and returns it, so its fields (and + suggestions) ride along in the upsert payload: Argilla v2.8.0 rejects a + field-less record with 422 "fields cannot be empty" because the required + text fields must be present, so ``id`` + metadata alone is not a valid + upsert. + + Callers batch the returned Records into a single ``dataset.records.log`` + call per dataset to amortise the round-trip. + """ + current = dict(record.metadata) + merged = dict(current) + merged.update(updates) + for key in remove_keys: + merged.pop(key, None) + if merged == current: + return None + for key, value in updates.items(): + record.metadata[key] = value + for key in remove_keys: + record.metadata.pop(key, None) + return record diff --git a/src/pragmata/core/annotation/panel_status.py b/src/pragmata/core/annotation/panel_status.py index 2a8f71d3..d76d4632 100644 --- a/src/pragmata/core/annotation/panel_status.py +++ b/src/pragmata/core/annotation/panel_status.py @@ -34,10 +34,11 @@ metadata; the live K is the ground truth pre-backfill and the metadata is cross-checked for integrity. -Pure read; safe to invoke against live datasets without side effects. The -optional ``--tag-partial-panels`` advisory write that stamps records for -annotator UI filtering ships in a follow-up PR (separates the read path from -any Argilla mutation surface). +The read path (``compute_panel_status``) is side-effect free. The optional +``--tag-partial-panels`` advisory write (``_apply_tags``, reached via +``report_status``) stamps partial panels' unresolved chunks for annotator UI +filtering, sharing this same single walk - it is the only Argilla mutation +surface here. """ import logging @@ -46,6 +47,8 @@ import argilla as rg +from pragmata.core.annotation.metadata_ops import build_metadata_upsert, ensure_metadata_property + logger = logging.getLogger(__name__) # Terminal response statuses = a judgement was recorded (vs pending/draft). @@ -55,6 +58,11 @@ RETRIEVAL_TASK = "retrieval" _TASK_ORDER = {"retrieval": 0, "grounding": 1, "generation": 2} +# Changing either of these orphans old key/value metadata already written to +# Argilla records - nothing here renames or removes it. Manual cleanup of the +# stale metadata would be needed on the live datasets. +NEEDS_COMPLETION_KEY = "needs_completion" +NEEDS_COMPLETION_VALUE = "true" def _task_of(dataset_name: str) -> str: @@ -62,6 +70,22 @@ def _task_of(dataset_name: str) -> str: return dataset_name.split("_", 1)[0] +def _has_needs_completion_tag(record: rg.Record) -> bool: + """Defensive equality check for the needs_completion tag. + + Argilla TermsMetadataProperty may round-trip as the bare value, a + 1-element list, or a normalised-case string. Treat anything string-equal + to NEEDS_COMPLETION_VALUE (after str-coercion + lowercase) as tagged, so + the idempotency check survives SDK encoding differences. + """ + raw = record.metadata.get(NEEDS_COMPLETION_KEY) + if raw is None: + return False + if isinstance(raw, (list, tuple)): + return any(str(v).strip().lower() == NEEDS_COMPLETION_VALUE for v in raw) + return str(raw).strip().lower() == NEEDS_COMPLETION_VALUE + + def _select_datasets(client: rg.Argilla, workspace: str | None, task: str | None) -> Iterator[rg.Dataset]: """Config-free dataset selection: iterate the live server, filter by name. @@ -133,12 +157,22 @@ class ProgressReport: by_dataset: list[ProgressRow] +@dataclass(frozen=True) +class TagResult: + """Counts from one ``--tag-partial-panels`` pass.""" + + n_tagged: int # chunks newly stamped with needs_completion + n_cleared: int # chunks where the stale tag was removed + n_already_tagged: int # already had the tag and still need it (no-op) + + @dataclass(frozen=True) class StatusReport: """Live per-panel status + headline aggregates. ``progress`` (all-task record counts) is attached by ``report_status``; the - panel fields below are retrieval-only. + panel fields below are retrieval-only. ``tag_result`` is None unless + ``--tag-partial-panels`` ran in the same pass. """ panels: dict[tuple[str, str], PanelStatus] @@ -149,11 +183,16 @@ class StatusReport: n_integrity_warnings: int n_orphans_skipped: int progress: "ProgressReport | None" = None + tag_result: "TagResult | None" = None def with_progress(self, progress: ProgressReport) -> "StatusReport": """Return a copy with the all-task ``progress`` summary attached.""" return replace(self, progress=progress) + def with_tag_result(self, tag_result: TagResult) -> "StatusReport": + """Return a copy with ``tag_result`` set (the dataclass is frozen).""" + return replace(self, tag_result=tag_result) + @dataclass class _ChunkRecord: @@ -415,3 +454,89 @@ def _row(label: str, task: str, b: dict[str, int]) -> ProgressRow: ) ] return ProgressReport(grand=HeadlineTotals(**grand), by_task=task_rows, by_workspace=ws_rows, by_dataset=ds_rows) + + +def _apply_tags(collected: _CollectedRecords) -> TagResult: + """Stamp / clear the ``needs_completion`` advisory tag on PARTIAL panels. + + Tag predicate: the panel is PARTIAL (at least one chunk has a submitted + response but NOT all chunks do) AND this chunk is UNRESOLVED (no terminal + response). The tag is cleared on resolved chunks and on non-partial panels + (fully-unstarted or complete), so a fully-unstarted panel is never tagged. + Idempotent: every run re-derives the set. + + Dataset-local and overlap-indifferent: PARTIAL/UNRESOLVED are derived + from ``has_terminal`` (>=1 response, any status) vs ``k_records`` alone - + ``min_submitted`` never enters this predicate. So a calibration chunk + with 1-of-3 submissions already counts as resolved for tagging purposes, + even though it hasn't hit its overlap target; this tag means "nobody has + looked at this yet," not "this hasn't reached ``overlap_satisfied``" (see + the module docstring's PARTIAL/``overlap_satisfied`` distinction). + + Batches one ``dataset.records.log`` per owning dataset. Datasets are keyed + by ``(workspace, name)`` because the same bare name (``retrieval_production``) + recurs across domains, so batching by name alone would misroute payloads. + """ + datasets_by_key: dict[tuple[str, str], rg.Dataset] = {} + for rec in collected.records: + datasets_by_key.setdefault((rec.workspace, rec.dataset.name), rec.dataset) + for dataset in datasets_by_key.values(): + ensure_metadata_property(dataset, rg.TermsMetadataProperty(NEEDS_COMPLETION_KEY, visible_for_annotators=True)) + + batched: dict[tuple[str, str], list[rg.Record]] = {} + n_tagged = n_cleared = n_already = 0 + for (_ws, uuid), group in _group_by_panel(collected.records).items(): + facts = _panel_facts(uuid, group) + # PARTIAL: some but not all chunks have a submitted response. A panel + # that is fully-unstarted (0 submitted) or complete is NOT partial. + panel_partial = 0 < len(facts.chunk_ids_submitted) < facts.k_records + for rec in group: + should_have_tag = panel_partial and not rec.has_terminal + already = _has_needs_completion_tag(rec.record) + if should_have_tag and already: + n_already += 1 + continue + if should_have_tag: + upsert = build_metadata_upsert(rec.record, {NEEDS_COMPLETION_KEY: NEEDS_COMPLETION_VALUE}) + elif already: + upsert = build_metadata_upsert(rec.record, {}, remove_keys=[NEEDS_COMPLETION_KEY]) + else: + continue + if upsert is None: + continue + key = (rec.workspace, rec.dataset.name) + batched.setdefault(key, []).append(upsert) + if should_have_tag: + n_tagged += 1 + else: + n_cleared += 1 + + for key, payloads in batched.items(): + # payloads is never empty: keys are created only on first append. + datasets_by_key[key].records.log(payloads) + + logger.info( + "tag_partial_panels: tagged=%d cleared=%d already_tagged=%d (panels=%d, datasets=%d)", + n_tagged, + n_cleared, + n_already, + len({(r.workspace, r.record_uuid) for r in collected.records}), + len(datasets_by_key), + ) + return TagResult(n_tagged=n_tagged, n_cleared=n_cleared, n_already_tagged=n_already) + + +def tag_partial_panels( + client: rg.Argilla, *, workspace: str | None = None, task: str | None = RETRIEVAL_TASK +) -> TagResult: + """Stamp / clear ``needs_completion`` advisory tags on partial retrieval panels. + + Tag predicate: panel is PARTIAL and this chunk is UNRESOLVED (see + ``_apply_tags``). Config-free; covers every workspace in one pass. + + Self-contained wrapper around ``_apply_tags``. Callers that already ran + ``_collect_records`` (e.g. ``report_status`` with ``tag_partial_panels=True``) + should call ``_apply_tags`` directly with the shared collection to avoid a + second walk. + """ + return _apply_tags(_collect_records(client, workspace=workspace, task=task)) diff --git a/tests/test_facade_lazy.py b/tests/test_facade_lazy.py index db66b4b8..2d67f999 100644 --- a/tests/test_facade_lazy.py +++ b/tests/test_facade_lazy.py @@ -123,6 +123,7 @@ def test_dir_returns_public_surface() -> None: "Locale", "SetupResult", "StatusReport", + "TagResult", "Task", "UserSpec", "compute_iaa", diff --git a/tests/unit/cli/test_cli_annotation.py b/tests/unit/cli/test_cli_annotation.py index c78212e4..82c6ef9f 100644 --- a/tests/unit/cli/test_cli_annotation.py +++ b/tests/unit/cli/test_cli_annotation.py @@ -124,6 +124,19 @@ def test_status_threads_defaults(self, mock_status): kwargs = mock_status.call_args.kwargs assert kwargs["api_url"] is UNSET assert kwargs["workspace"] is None + assert kwargs["tag_partial_panels"] is False + + @patch("pragmata.annotation.report_status") + def test_status_tag_partial_panels_flag_and_echo(self, mock_status): + from pragmata.core.annotation.panel_status import TagResult + + mock_status.return_value = self._stub_report().with_tag_result( + TagResult(n_tagged=5, n_cleared=2, n_already_tagged=1) + ) + result = runner.invoke(app, ["annotation", "status", "--tag-partial-panels"]) + assert result.exit_code == 0 + assert mock_status.call_args.kwargs["tag_partial_panels"] is True + assert "tag-partial-panels: tagged=5 cleared=2 already_tagged=1" in result.output class TestIaaCommand: diff --git a/tests/unit/core/annotation/test_metadata_ops.py b/tests/unit/core/annotation/test_metadata_ops.py new file mode 100644 index 00000000..66bcbf58 --- /dev/null +++ b/tests/unit/core/annotation/test_metadata_ops.py @@ -0,0 +1,104 @@ +"""Unit tests for shared safe metadata ops (used by status --tag-partial-panels + backfill).""" + +from unittest.mock import MagicMock + +import argilla as rg + +from pragmata.core.annotation.metadata_ops import ensure_metadata_property + + +def _mock_dataset(*, existing_metadata_props: list[str] | None = None) -> MagicMock: + dataset = MagicMock() + dataset.name = "ds" + metadata = MagicMock() + metadata.__getitem__ = MagicMock(side_effect=lambda key: key if key in (existing_metadata_props or []) else None) + dataset.settings.metadata = metadata + return dataset + + +def _mock_record(metadata: dict[str, object], record_id: str = "rec-1") -> MagicMock: + # spec=rg.Record so the "returns an rg.Record, not a dict" assertion is + # meaningful (a bare MagicMock never satisfies isinstance(..., rg.Record)). + record = MagicMock(spec=rg.Record) + record.id = record_id + record.metadata = metadata + return record + + +class TestEnsureMetadataProperty: + def test_adds_when_absent(self) -> None: + dataset = _mock_dataset(existing_metadata_props=[]) + # The MagicMock for IntegerMetadataProperty can be unauthenticated; use a real instance + # via a stand-in MagicMock to avoid the Argilla SDK requiring credentials at construction. + prop = MagicMock(spec=rg.IntegerMetadataProperty) + prop.name = "n_retrieved_chunks" + + added = ensure_metadata_property(dataset, prop) + + assert added is True + dataset.settings.add.assert_called_once_with(prop) + dataset.settings.update.assert_called_once() + + def test_skips_when_present(self) -> None: + dataset = _mock_dataset(existing_metadata_props=["n_retrieved_chunks"]) + prop = MagicMock(spec=rg.IntegerMetadataProperty) + prop.name = "n_retrieved_chunks" + + added = ensure_metadata_property(dataset, prop) + + assert added is False + dataset.settings.add.assert_not_called() + dataset.settings.update.assert_not_called() + + +class TestBuildMetadataUpsert: + def test_returns_record_with_full_merged_metadata(self) -> None: + from pragmata.core.annotation.metadata_ops import build_metadata_upsert + + record = _mock_record( + {"record_uuid": "u1", "chunk_id": "c1", "chunk_rank": 3, "doc_id": "d1"}, + record_id="rec-1", + ) + + upsert = build_metadata_upsert(record, {"n_retrieved_chunks": 5}) + + assert upsert is not None + assert upsert.id == "rec-1" + assert dict(upsert.metadata) == { + "record_uuid": "u1", + "chunk_id": "c1", + "chunk_rank": 3, + "doc_id": "d1", + "n_retrieved_chunks": 5, + } + + def test_returns_none_when_unchanged(self) -> None: + from pragmata.core.annotation.metadata_ops import build_metadata_upsert + + record = _mock_record({"chunk_id": "c1", "flag": "yes"}) + assert build_metadata_upsert(record, {"flag": "yes"}) is None + + def test_remove_keys_drops_existing(self) -> None: + from pragmata.core.annotation.metadata_ops import build_metadata_upsert + + record = _mock_record({"chunk_id": "c1", "needs_completion": "true"}) + upsert = build_metadata_upsert(record, {}, remove_keys=["needs_completion"]) + assert upsert is not None + assert "needs_completion" not in dict(upsert.metadata) + assert dict(upsert.metadata)["chunk_id"] == "c1" + + def test_returns_record_as_rg_Record_not_dict(self) -> None: + """The upsert MUST be an rg.Record, not a dict. + + Sending a dict to dataset.records.log runs through Argilla's + IngestedRecordMapper, which flattens against the dataset schema and + silently strips the 'metadata' key (wiping the record). + """ + from pragmata.core.annotation.metadata_ops import build_metadata_upsert + + record = _mock_record({"chunk_id": "c1"}) + upsert = build_metadata_upsert(record, {"flag": "yes"}) + + import argilla as rg + + assert isinstance(upsert, rg.Record) diff --git a/tests/unit/core/annotation/test_panel_status.py b/tests/unit/core/annotation/test_panel_status.py index f4ecf077..239c5285 100644 --- a/tests/unit/core/annotation/test_panel_status.py +++ b/tests/unit/core/annotation/test_panel_status.py @@ -1,11 +1,15 @@ """Unit tests for live panel status (read-only, config-free).""" import logging -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest -from pragmata.core.annotation.panel_status import compute_panel_status, compute_task_progress +from pragmata.core.annotation.panel_status import ( + compute_panel_status, + compute_task_progress, + tag_partial_panels, +) WS = "dom_retrieval" @@ -17,10 +21,13 @@ def _record( chunk_id: str = "c1", n_retrieved_chunks: int | None = None, response_statuses: list[str] | None = None, + needs_completion: str | None = None, ) -> MagicMock: metadata: dict[str, object] = {"record_uuid": record_uuid, "chunk_id": chunk_id} if n_retrieved_chunks is not None: metadata["n_retrieved_chunks"] = n_retrieved_chunks + if needs_completion is not None: + metadata["needs_completion"] = needs_completion rec = MagicMock() rec.id = record_id rec.metadata = metadata @@ -251,3 +258,107 @@ def test_all_tasks_present_even_with_zero_progress(self) -> None: pr = compute_task_progress(client) assert [r.task for r in pr.by_task] == ["grounding"] assert pr.grand.completed == 0 + + +class TestTagPartialPanels: + def test_tags_unresolved_chunks_of_partial_panel(self) -> None: + # K=3: one chunk submitted, two unstarted → partial → tag the two. + recs = [ + _record(chunk_id="c1", response_statuses=["submitted"], record_id="r1"), + _record(chunk_id="c2", response_statuses=[], record_id="r2"), + _record(chunk_id="c3", response_statuses=[], record_id="r3"), + ] + ds = _dataset("retrieval_production", recs) + with patch("argilla.TermsMetadataProperty"): + result = tag_partial_panels(_client([ds])) + assert (result.n_tagged, result.n_cleared, result.n_already_tagged) == (2, 0, 0) + (logged,) = ds.records.log.call_args.args + assert {r.id for r in logged} == {"r2", "r3"} + + def test_fully_unstarted_panel_not_tagged(self) -> None: + recs = [_record(chunk_id=f"c{i}", response_statuses=[], record_id=f"r{i}") for i in range(3)] + ds = _dataset("retrieval_production", recs) + with patch("argilla.TermsMetadataProperty"): + result = tag_partial_panels(_client([ds])) + assert result.n_tagged == 0 + ds.records.log.assert_not_called() + + def test_complete_panel_not_tagged(self) -> None: + recs = [_record(chunk_id=f"c{i}", response_statuses=["submitted"], record_id=f"r{i}") for i in range(3)] + ds = _dataset("retrieval_production", recs) + with patch("argilla.TermsMetadataProperty"): + result = tag_partial_panels(_client([ds])) + assert result.n_tagged == 0 + + def test_clears_stale_tag_on_non_partial_panel(self) -> None: + # Fully-unstarted panel whose chunk carries a stale needs_completion tag + # (e.g. from the old broad predicate) → cleared, never re-tagged. + recs = [ + _record(chunk_id="c1", response_statuses=[], record_id="r1", needs_completion="true"), + _record(chunk_id="c2", response_statuses=[], record_id="r2"), + ] + ds = _dataset("retrieval_production", recs) + with patch("argilla.TermsMetadataProperty"): + result = tag_partial_panels(_client([ds])) + assert (result.n_tagged, result.n_cleared) == (0, 1) + (logged,) = ds.records.log.call_args.args + assert {r.id for r in logged} == {"r1"} + assert "needs_completion" not in logged[0].metadata + + def test_already_tagged_is_idempotent(self) -> None: + recs = [ + _record(chunk_id="c1", response_statuses=["submitted"], record_id="r1"), + _record(chunk_id="c2", response_statuses=[], record_id="r2", needs_completion="true"), + ] + ds = _dataset("retrieval_production", recs) + with patch("argilla.TermsMetadataProperty"): + result = tag_partial_panels(_client([ds])) + assert (result.n_tagged, result.n_cleared, result.n_already_tagged) == (0, 0, 1) + ds.records.log.assert_not_called() + + def test_split_panel_across_prod_and_cal(self) -> None: + # Per-item calibration: the only annotated chunk is in cal; prod chunks + # unstarted. Panel is partial ACROSS datasets → tag the prod chunks. + prod = _dataset( + "retrieval_production", + [ + _record(chunk_id="c2", response_statuses=[], record_id="p2"), + _record(chunk_id="c3", response_statuses=[], record_id="p3"), + ], + ) + cal = _dataset( + "retrieval_calibration", + [_record(chunk_id="c1", response_statuses=["submitted"], record_id="cal1")], + min_submitted=3, + ) + with patch("argilla.TermsMetadataProperty"): + result = tag_partial_panels(_client([prod, cal])) + assert result.n_tagged == 2 + (logged,) = prod.records.log.call_args.args + assert {r.id for r in logged} == {"p2", "p3"} + cal.records.log.assert_not_called() + + def test_writes_routed_per_workspace_when_names_collide(self) -> None: + d1 = _dataset( + "retrieval_production", + [ + _record(record_uuid="uA", chunk_id="c1", response_statuses=["submitted"], record_id="a1"), + _record(record_uuid="uA", chunk_id="c2", response_statuses=[], record_id="a2"), + ], + workspace="dom1_retrieval", + ) + d2 = _dataset( + "retrieval_production", + [ + _record(record_uuid="uB", chunk_id="c1", response_statuses=["submitted"], record_id="b1"), + _record(record_uuid="uB", chunk_id="c2", response_statuses=[], record_id="b2"), + ], + workspace="dom2_retrieval", + ) + with patch("argilla.TermsMetadataProperty"): + result = tag_partial_panels(_client([d1, d2])) + assert result.n_tagged == 2 + (l1,) = d1.records.log.call_args.args + (l2,) = d2.records.log.call_args.args + assert {r.id for r in l1} == {"a2"} + assert {r.id for r in l2} == {"b2"}