diff --git a/src/pragmata/api/annotation_import.py b/src/pragmata/api/annotation_import.py index 169bba4e..209b727c 100644 --- a/src/pragmata/api/annotation_import.py +++ b/src/pragmata/api/annotation_import.py @@ -49,10 +49,12 @@ class ImportResult: production dataset. calibration_fraction: Per-task configured fraction. May differ from ``realised_calibration_fraction`` on re-imports because prior - assignments are locked by the manifest. + assignments are locked by the manifest, and under a binding cap. realised_calibration_fraction: Per-task actual share of items routed to calibration this run (calibration / (calibration + production)). Zero when no items were assigned for that task. + calibration_max_items: Per-task configured absolute cap. ``None`` = + uncapped for that task. errors: Per-record validation failures (index + detail). """ @@ -62,6 +64,7 @@ class ImportResult: production_count: dict[Task, int] = field(default_factory=dict) calibration_fraction: dict[Task, float] = field(default_factory=dict) realised_calibration_fraction: dict[Task, float] = field(default_factory=dict) + calibration_max_items: dict[Task, int | None] = field(default_factory=dict) errors: list[RecordError] = field(default_factory=list) @@ -75,6 +78,7 @@ def import_records( base_dir: str | Path | Unset = UNSET, config_path: str | Path | Unset = UNSET, calibration_fraction: float | Unset = UNSET, + calibration_max_items: int | None | Unset = UNSET, calibration_min_submitted: int | None | Unset = UNSET, calibration_partition_seed: int | Unset = UNSET, locale: Locale | Unset = UNSET, @@ -124,6 +128,11 @@ def import_records( unless overridden in YAML config). Falls through to YAML config and the built-in default (0.1) when omitted. Set to 0.0 for production-only batches. + calibration_max_items: Deployment-level absolute cap on calibration + annotation items per task. ``None`` is uncapped (just the + fractional knob). Inherited by workspaces/tasks unless overridden + in YAML config. Cap unit is the annotation item: chunks for + retrieval, records for grounding / generation. calibration_min_submitted: Deployment-level overlap requirement for the calibration dataset. ``None`` disables calibration entirely (must be paired with ``calibration_fraction=0.0``). Inherits to @@ -153,6 +162,7 @@ def import_records( "dataset_id": dataset_id, "base_dir": base_dir, "calibration_fraction": calibration_fraction, + "calibration_max_items": calibration_max_items, "calibration_min_submitted": calibration_min_submitted, "calibration_partition_seed": calibration_partition_seed, "locale": locale, @@ -199,19 +209,26 @@ def import_records( # Manifest is written only after fan-out succeeds. On failure, the # in-memory assignments are dropped; a retry with the same corpus + seed - # re-derives identical buckets (deterministic per-unit hashing). + # re-derives identical buckets under the fraction-only path. Under a + # binding cap, retries are order-dependent (see + # ``test_cap_under_split_imports_is_order_dependent_by_design``). dataset_counts = fan_out_records(client, settings, partition=partition) write_partition_manifest(import_paths.partition_manifest, manifest) + def _per_task_summary(task: Task) -> str: + cap = partition.calibration_max_items[task] + cap_suffix = f", cap={cap}" if cap is not None else "" + return ( + f"{task.value}={summary.calibration_count[task]} " + f"(realised={summary.realised_fraction[task]:.3f}, " + f"configured={summary.configured_fraction[task]:.3f}{cap_suffix})" + ) + logger.info( "Import complete: %d records across %d datasets. Per-task calibration: %s", len(raw), len(dataset_counts), - ", ".join( - f"{task.value}={summary.calibration_count[task]} " - f"(realised={summary.realised_fraction[task]:.3f}, configured={summary.configured_fraction[task]:.3f})" - for task in Task - ), + ", ".join(_per_task_summary(task) for task in Task), ) return ImportResult( total_records=len(raw), @@ -220,5 +237,6 @@ def import_records( production_count=summary.production_count, calibration_fraction=summary.configured_fraction, realised_calibration_fraction=summary.realised_fraction, + calibration_max_items=partition.calibration_max_items, errors=validation.errors, ) diff --git a/src/pragmata/cli/commands/annotation.py b/src/pragmata/cli/commands/annotation.py index 17be0d34..9753a72b 100644 --- a/src/pragmata/cli/commands/annotation.py +++ b/src/pragmata/cli/commands/annotation.py @@ -106,11 +106,21 @@ def import_command( "YAML config). Falls through to YAML config and built-in default (0.1) when " "omitted; set to 0.0 for production-only batches.", ), + calibration_max_items: int | None = typer.Option( + None, + "--calibration-max-items", + help="Deployment-level absolute cap on calibration annotation items per task " + "(inherited by workspaces/tasks unless overridden in YAML config). Smaller of " + "(fraction × N_items, cap) wins. Existing assignments are never demoted. " + "Cap unit is the annotation item: chunks for retrieval, records for grounding " + "and generation. Omit to leave uncapped.", + ), no_calibration: bool = typer.Option( False, "--no-calibration", help="Disable calibration entirely for this batch: sets calibration_min_submitted=None " - "and calibration_fraction=0.0. Cannot be combined with --calibration-fraction > 0.", + "and calibration_fraction=0.0. Cannot be combined with --calibration-fraction > 0 " + "or --calibration-max-items.", ), calibration_partition_seed: int | None = typer.Option( None, @@ -147,6 +157,12 @@ def import_command( err=True, ) raise typer.Exit(code=2) + if no_calibration and calibration_max_items is not None: + typer.echo( + "Error: --no-calibration cannot be combined with --calibration-max-items.", + err=True, + ) + raise typer.Exit(code=2) result = annotation.import_records( records, @@ -159,6 +175,7 @@ def import_command( calibration_fraction=( 0.0 if no_calibration else UNSET if calibration_fraction is None else calibration_fraction ), + calibration_max_items=UNSET if calibration_max_items is None else calibration_max_items, calibration_min_submitted=None if no_calibration else UNSET, calibration_partition_seed=UNSET if calibration_partition_seed is None else calibration_partition_seed, locale=parse_locale(locale) or UNSET, @@ -169,9 +186,11 @@ def import_command( prod_n = result.production_count.get(task, 0) cfg_fraction = result.calibration_fraction.get(task, 0.0) realised = result.realised_calibration_fraction.get(task, 0.0) + cap = result.calibration_max_items.get(task) + cap_str = f", cap={cap}" if cap is not None else "" typer.echo( f" {task.value}: calibration={cal_n}, production={prod_n} " - f"(configured={cfg_fraction:.3f}, realised={realised:.3f})" + f"(configured={cfg_fraction:.3f}, realised={realised:.3f}{cap_str})" ) for ds, count in result.dataset_counts.items(): typer.echo(f" {ds}: {count}") diff --git a/src/pragmata/core/annotation/record_builder.py b/src/pragmata/core/annotation/record_builder.py index 0b49c0da..8f643484 100644 --- a/src/pragmata/core/annotation/record_builder.py +++ b/src/pragmata/core/annotation/record_builder.py @@ -9,7 +9,7 @@ import hashlib import json import logging -from collections.abc import Callable, Iterable +from collections.abc import Iterable from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -235,7 +235,7 @@ def write_partition_manifest(path: Path, manifest: PartitionManifest) -> None: class PartitionResult: """Outcome of ``assign_partitions()``. - Bundles the per-record manifest entries with the per-task fraction + Bundles the per-record manifest entries with the per-task fraction / cap resolved this run and the rid → pair map, so downstream callers don't redo the workspace walk or recompute UUIDs. @@ -245,11 +245,14 @@ class PartitionResult: pairs_by_rid: record_uuid → QueryResponsePair for every input pair. calibration_fraction: Per-task resolved fraction in force this run. ``0.0`` for tasks absent from the workspaces topology. + calibration_max_items: Per-task resolved absolute cap. ``None`` = + uncapped (or task absent from topology). """ assignments: dict[str, PartitionManifestEntry] pairs_by_rid: dict[str, QueryResponsePair] calibration_fraction: dict[Task, float] + calibration_max_items: dict[Task, int | None] def assign_partitions( @@ -263,77 +266,132 @@ def assign_partitions( The annotation item differs by task: for grounding and generation, one item per ``record_uuid``; for retrieval, one item per ``(record_uuid, chunk_id)``. - Per-task fraction is resolved via ``settings.resolved_task`` so workspace / - task overrides for ``calibration_fraction`` are honoured. - - Existing assignments are locked for re-import stability and never rewritten. - A new record is bucketed across all active tasks by per-(task, unit) digest - ``hash(seed || task || unit_id)`` against ``fraction * 2^32``. An existing - record is *backfilled* for any task the topology has since gained (e.g. - retrieval-only at first import, grounding added later) using the same - deterministic digest, so late-activated tasks partition pre-existing records - identically to fresh ones rather than silently defaulting to production. + Per-task fraction and cap are resolved via ``settings.resolved_task`` so + workspace / task overrides for ``calibration_fraction`` and + ``calibration_max_items`` are honoured. + + Already-recorded units are locked for re-import stability and never + rewritten. Every other unit is an *unassigned candidate* this run: all units + of a brand-new record, plus units of any task an existing record has gained + since its last import (*backfill*, e.g. retrieval-only at first import, + grounding added later). Per task, eligible candidates (``fraction >= 1.0`` or + digest ``hash(seed || task || unit_id) < fraction * 2^32``) are ranked by + digest ascending and, when a per-task cap is in force, only the first + ``remaining = max(0, cap - existing_calibration)`` win calibration; the rest + demote to production. New and backfilled candidates compete for the same + budget, so the cap binds on all calibration assigned this run. + + Note on order-dependence under binding cap: when the cap binds across + multiple imports, the final calibration set is a function of + ``(corpus, seed, import_order)``, not ``(corpus, seed)`` alone — once an + existing unit is in calibration, a tightened cap on a later import cannot + demote it (manifest lock). Args: pairs: Validated pairs to partition. manifest: Mutated in place to record new and backfilled assignments. - settings: Resolves per-task fraction via ``resolved_task``. + settings: Resolves per-task fraction and cap via ``resolved_task``. import_id: Stamped on new entries for provenance. Returns: A ``PartitionResult`` bundling the assignments, the rid → pair map, and - the per-task fraction resolved this run. + the per-task fraction and cap resolved this run. """ now = datetime.now(timezone.utc) seed = manifest.partition_seed workspace_for_task = settings.task_to_workspace() pairs_by_rid: dict[str, QueryResponsePair] = {derive_record_uuid(pair): pair for pair in pairs} + active_tasks = {task for task in Task if workspace_for_task.get(task) is not None} + per_task_fraction: dict[Task, float] = {} + per_task_cap: dict[Task, int | None] = {} threshold_for_task: dict[Task, int | None] = {} for task in Task: ws_base = workspace_for_task.get(task) if ws_base is None: per_task_fraction[task] = 0.0 + per_task_cap[task] = None threshold_for_task[task] = None continue - fraction = settings.resolved_task(ws_base, task).calibration_fraction + resolved = settings.resolved_task(ws_base, task) + fraction = resolved.calibration_fraction per_task_fraction[task] = fraction + per_task_cap[task] = resolved.calibration_max_items threshold_for_task[task] = int(fraction * (2**32)) if 0.0 < fraction < 1.0 else None - def is_cal(task: Task, unit_id: str) -> bool: + existing_cal_counts = count_units_per_task(manifest.assignments.values()).calibration + + # Units assigned this run (new records + backfilled tasks), keyed by rid. + new_grnd_gen: dict[str, dict[Task, bool]] = {} + new_retrieval: dict[str, dict[str, bool]] = {} + + def _record(rid: str, chunk_id: str | None, task: Task, *, is_cal: bool) -> None: + new_grnd_gen.setdefault(rid, {}) + new_retrieval.setdefault(rid, {}) + _write_unit(new_grnd_gen[rid], new_retrieval[rid], chunk_id, task, is_cal=is_cal) + + for task in active_tasks: fraction = per_task_fraction[task] threshold = threshold_for_task[task] - return fraction >= 1.0 or (threshold is not None and _calibration_digest(unit_id, task, seed) < threshold) + cap = per_task_cap[task] + existing_cal = existing_cal_counts[task] + if cap is not None and existing_cal > cap: + logger.warning( + "Task %s: existing calibration count %d exceeds cap %d; " + "no new calibration items promoted this run (manifest-lock invariant).", + task.value, + existing_cal, + cap, + ) + remaining: int | None = 0 + else: + remaining = None if cap is None else max(0, cap - existing_cal) + + candidates: list[tuple[int, str, str | None]] = [] # (digest, rid, chunk_id_or_none) + for rid, pair in pairs_by_rid.items(): + existing = manifest.assignments.get(rid) + for unit_id, chunk_id in _enumerate_units(rid, pair, task): + if existing is not None and _unit_assigned(existing, task, chunk_id): + continue # locked - never rewritten + digest = _calibration_digest(unit_id, task, seed) + eligible = fraction >= 1.0 or (threshold is not None and digest < threshold) + if eligible: + candidates.append((digest, rid, chunk_id)) + else: + _record(rid, chunk_id, task, is_cal=False) + + candidates.sort() + promoted = candidates if remaining is None else candidates[:remaining] + demoted = [] if remaining is None else candidates[remaining:] + for _, rid, chunk_id in promoted: + _record(rid, chunk_id, task, is_cal=True) + for _, rid, chunk_id in demoted: + _record(rid, chunk_id, task, is_cal=False) - active_tasks = {task for task in Task if workspace_for_task.get(task) is not None} assignments: dict[str, PartitionManifestEntry] = {} changed = False for rid, pair in pairs_by_rid.items(): existing = manifest.assignments.get(rid) + add_gg = new_grnd_gen.get(rid, {}) + add_ret = new_retrieval.get(rid, {}) if existing is None: - grnd_gen: dict[Task, bool] = {} - retrieval: dict[str, bool] = {} - for task in active_tasks: - for unit_id, chunk_id in _enumerate_units(rid, pair, task): - _write_unit(grnd_gen, retrieval, chunk_id, task, is_cal=is_cal(task, unit_id)) entry = PartitionManifestEntry( - grounding_generation_calibration=grnd_gen, - retrieval_chunk_calibration=retrieval, + grounding_generation_calibration=add_gg, + retrieval_chunk_calibration=add_ret, import_id=import_id, calibration_fraction_at_import=dict(per_task_fraction), + calibration_max_items_at_import=dict(per_task_cap), assigned_at=now, ) - manifest.assignments[rid] = entry - assignments[rid] = entry - changed = True + elif add_gg or add_ret: + entry = _merge_backfill(existing, add_gg, add_ret, per_task_fraction, per_task_cap) + else: + assignments[rid] = existing # nothing new - reuse untouched continue - - entry = _backfill_entry(existing, rid, pair, active_tasks, per_task_fraction, is_cal) - if entry is not existing: - manifest.assignments[rid] = entry - changed = True + manifest.assignments[rid] = entry assignments[rid] = entry + changed = True if changed: manifest.updated_at = now @@ -341,49 +399,42 @@ def is_cal(task: Task, unit_id: str) -> bool: assignments=assignments, pairs_by_rid=pairs_by_rid, calibration_fraction=per_task_fraction, + calibration_max_items=per_task_cap, ) -def _backfill_entry( +def _unit_assigned(entry: PartitionManifestEntry, task: Task, chunk_id: str | None) -> bool: + """Whether this unit already carries a locked assignment on the entry.""" + if task == Task.RETRIEVAL: + return chunk_id in entry.retrieval_chunk_calibration + return task in entry.grounding_generation_calibration + + +def _merge_backfill( entry: PartitionManifestEntry, - rid: str, - pair: QueryResponsePair, - active_tasks: set[Task], + add_grnd_gen: dict[Task, bool], + add_retrieval: dict[str, bool], per_task_fraction: dict[Task, float], - is_cal: Callable[[Task, str], bool], + per_task_cap: dict[Task, int | None], ) -> PartitionManifestEntry: - """Assign active tasks/chunks absent from an existing entry; never rewrite. + """Merge newly-assigned (backfilled) units into an existing entry. - Returns ``entry`` unchanged when nothing is missing, else a new entry with - the newly-active units assigned via the same deterministic digest a fresh - import uses. Provenance (``import_id``, ``assigned_at``) is preserved from the - original assignment; ``calibration_fraction_at_import`` records the resolved - fraction for each backfilled task. + Existing assignments are never rewritten - only tasks/chunks absent from the + entry are added. Provenance (``import_id``, ``assigned_at``) is preserved; + fraction/cap provenance is recorded for each backfilled task. """ - grnd_gen = dict(entry.grounding_generation_calibration) - retrieval = dict(entry.retrieval_chunk_calibration) fraction_at_import = dict(entry.calibration_fraction_at_import) - added = False - for task in active_tasks: - for unit_id, chunk_id in _enumerate_units(rid, pair, task): - if task == Task.RETRIEVAL: - assert chunk_id is not None # retrieval always has chunk_id - if chunk_id in retrieval: - continue - retrieval[chunk_id] = is_cal(task, unit_id) - elif task not in grnd_gen: - grnd_gen[task] = is_cal(task, unit_id) - else: - continue - fraction_at_import[task] = per_task_fraction[task] - added = True - if not added: - return entry + cap_at_import = dict(entry.calibration_max_items_at_import) + backfilled_tasks = set(add_grnd_gen) | ({Task.RETRIEVAL} if add_retrieval else set()) + for task in backfilled_tasks: + fraction_at_import[task] = per_task_fraction[task] + cap_at_import[task] = per_task_cap[task] return entry.model_copy( update={ - "grounding_generation_calibration": grnd_gen, - "retrieval_chunk_calibration": retrieval, + "grounding_generation_calibration": {**entry.grounding_generation_calibration, **add_grnd_gen}, + "retrieval_chunk_calibration": {**entry.retrieval_chunk_calibration, **add_retrieval}, "calibration_fraction_at_import": fraction_at_import, + "calibration_max_items_at_import": cap_at_import, } ) @@ -634,8 +685,15 @@ def fan_out_records( continue ws_base = task_to_ws.get(task) if ws_base is None: - logger.warning("Task %r not in workspaces topology - skipping", task) - continue + # Manifest entries exist for a task absent from the current topology - + # silently skipping would drop data on disk. Force the operator to + # re-add the task or clear the affected partition manifest. + raise RuntimeError( + f"Task {task.value!r} has {len(rg_records)} records assigned via the " + f"partition manifest, but the current workspaces topology does not " + f"include this task. Re-add the task to the topology, or clear the " + f"partition manifest at the affected dataset scope." + ) resolved = settings.resolved_task(ws_base, task) if calibration: if resolved.calibration_min_submitted is None: diff --git a/src/pragmata/core/schemas/annotation_import.py b/src/pragmata/core/schemas/annotation_import.py index ff485b27..94adaead 100644 --- a/src/pragmata/core/schemas/annotation_import.py +++ b/src/pragmata/core/schemas/annotation_import.py @@ -78,6 +78,8 @@ class PartitionManifestEntry(BaseModel): import_id: Identifier of the import call that produced this assignment. calibration_fraction_at_import: Per-task fraction in force at that import call. + calibration_max_items_at_import: Per-task absolute cap in force at + that import call (``None`` = uncapped). assigned_at: When the assignment was made. """ @@ -87,6 +89,7 @@ class PartitionManifestEntry(BaseModel): retrieval_chunk_calibration: dict[str, bool] = Field(default_factory=dict) import_id: NonEmptyStr calibration_fraction_at_import: dict[Task, Annotated[float, Field(ge=0.0, le=1.0)]] + calibration_max_items_at_import: dict[Task, Annotated[int, Field(ge=1)] | None] = Field(default_factory=dict) assigned_at: datetime diff --git a/src/pragmata/core/settings/annotation_settings.py b/src/pragmata/core/settings/annotation_settings.py index 9951e0dc..37cd8529 100644 --- a/src/pragmata/core/settings/annotation_settings.py +++ b/src/pragmata/core/settings/annotation_settings.py @@ -3,12 +3,12 @@ Settings are organised into three scopes: deployment (``AnnotationSettings``), workspace (``WorkspaceSettings``), task (``TaskSettings``). The inheritable fields (``production_min_submitted``, ``calibration_min_submitted``, ``locale``, -``calibration_fraction``) may be set at any scope; child scopes default to -``INHERIT``. Models hold the **specified** values exactly as given -(``INHERIT`` survives validation, raw inputs round-trip losslessly through -``model_dump()``). ``resolved_task(workspace_name, task)`` returns the -**computed** values after walking task → workspace → deployment — the CSS -"computed value" analogy: first non-``INHERIT`` ancestor wins. +``calibration_fraction``, ``calibration_max_items``) may be set at any scope; +child scopes default to ``INHERIT``. Models hold the **specified** values +exactly as given (``INHERIT`` survives validation, raw inputs round-trip +losslessly through ``model_dump()``). ``resolved_task(workspace_name, task)`` +returns the **computed** values after walking task → workspace → deployment — +the CSS "computed value" analogy: first non-``INHERIT`` ancestor wins. Multi-key maps (e.g. ``constraint_severity: dict[str, Severity]``) use **sparse-dict overlay** rather than the ``INHERIT`` sentinel: user-supplied keys @@ -61,6 +61,7 @@ class TaskSettings(BaseModel): calibration_min_submitted: PositiveInt | None | Inherit = INHERIT locale: Locale | Inherit = INHERIT calibration_fraction: CalibrationFractionOverride = INHERIT + calibration_max_items: PositiveInt | None | Inherit = INHERIT class WorkspaceSettings(BaseModel): @@ -79,6 +80,7 @@ class WorkspaceSettings(BaseModel): calibration_min_submitted: PositiveInt | None | Inherit = INHERIT locale: Locale | Inherit = INHERIT calibration_fraction: CalibrationFractionOverride = INHERIT + calibration_max_items: PositiveInt | None | Inherit = INHERIT constraint_severity: dict[str, Severity] = Field(default_factory=dict) tasks: dict[Task, TaskSettings] @@ -91,6 +93,7 @@ class ResolvedTaskSettings: calibration_min_submitted: int | None locale: Locale calibration_fraction: float + calibration_max_items: int | None def _inherit(*candidates: T | Inherit) -> T: @@ -125,6 +128,10 @@ class AnnotationSettings(ResolveSettings): calibration dataset for IAA (0.0 disables for that scope). Inherited by workspaces/tasks. The annotation item is a chunk for retrieval and a record_uuid for grounding / generation. + calibration_max_items: Optional absolute cap on calibration + annotation items per task. ``None`` is uncapped (just the + fractional knob). Smaller of (fraction × N_items, cap) wins. + Inherited by workspaces/tasks. """ argilla: ArgillaSettings = Field(default_factory=ArgillaSettings) @@ -135,6 +142,7 @@ class AnnotationSettings(ResolveSettings): locale: Locale = "en" locale_catalog_dir: Path | None = None calibration_fraction: float = Field(0.1, ge=0.0, le=1.0) + calibration_max_items: PositiveInt | None = None calibration_partition_seed: NonNegativeInt = 0 include_discarded: bool = False constraint_severity: dict[str, Severity] = Field(default_factory=lambda: dict(_DEFAULT_CONSTRAINT_SEVERITY)) @@ -182,6 +190,9 @@ def resolved_task(self, workspace_name: str, task: Task) -> ResolvedTaskSettings ), locale=_inherit(ts.locale, ws.locale, self.locale), calibration_fraction=_inherit(ts.calibration_fraction, ws.calibration_fraction, self.calibration_fraction), + calibration_max_items=_inherit( + ts.calibration_max_items, ws.calibration_max_items, self.calibration_max_items + ), ) @model_validator(mode="before") diff --git a/tests/integration/test_annotation_calibration.py b/tests/integration/test_annotation_calibration.py index 2a208b03..c1826f50 100644 --- a/tests/integration/test_annotation_calibration.py +++ b/tests/integration/test_annotation_calibration.py @@ -41,6 +41,19 @@ def _make_raw(i: int, *, n_chunks: int = 1) -> dict: } +def _make_raw_multi_chunk(i: int, *, n_chunks: int) -> dict: + return { + "query": f"Question {i}?", + "answer": f"Answer {i}.", + "chunks": [ + {"chunk_id": f"c{i}-{k}", "doc_id": "d1", "chunk_rank": k + 1, "text": f"Chunk {i}-{k}."} + for k in range(n_chunks) + ], + "context_set": "ctx-001", + "language": "en", + } + + def _manifest_path(base_dir: Path, dataset_id: str) -> Path: """Resolve the partition manifest path the same way the import API does.""" workspace = WorkspacePaths.from_base_dir(base_dir) @@ -242,6 +255,97 @@ def test_records_carry_calibration_metadata_to_argilla(client: rg.Argilla, base_ teardown_resources(client, auto_settings) +def test_calibration_max_items_caps_grounding(client: rg.Argilla, base_dir: Path) -> None: + """Per-task cap binds: realised calibration count never exceeds the absolute cap.""" + auto_id = "testcap" + auto_settings = AnnotationSettings(dataset_id=auto_id) + teardown_resources(client, auto_settings) + setup_workspaces(client, auto_settings) + + try: + records = [_make_raw(i) for i in range(30)] + result = import_records( + records, + dataset_id=auto_id, + base_dir=base_dir, + calibration_fraction=1.0, + calibration_max_items=3, + **_CREDS, + ) + + assert result.calibration_count[Task.GROUNDING] == 3 + assert result.production_count[Task.GROUNDING] == 27 + assert result.calibration_max_items[Task.GROUNDING] == 3 + + cal_ds = client.datasets( + dataset_name(Task.GROUNDING, calibration=True, dataset_id=auto_id), + workspace="grounding", + ) + prod_ds = client.datasets( + dataset_name(Task.GROUNDING, calibration=False, dataset_id=auto_id), + workspace="grounding", + ) + assert cal_ds is not None and prod_ds is not None + assert len(list(cal_ds.records)) == 3 + assert len(list(prod_ds.records)) == 27 + finally: + teardown_resources(client, auto_settings) + + +def test_per_chunk_retrieval_routing_to_argilla(client: rg.Argilla, base_dir: Path) -> None: + """Different chunks of one record can land in different retrieval datasets on Argilla.""" + auto_id = "testperchunk" + auto_settings = AnnotationSettings(dataset_id=auto_id) + teardown_resources(client, auto_settings) + setup_workspaces(client, auto_settings) + + try: + # 30 records × 4 chunks = 120 retrieval units at fraction=0.5; per-chunk + # independence means some pairs end up with a mixed retrieval bucket set. + records = [_make_raw_multi_chunk(i, n_chunks=4) for i in range(30)] + import_records(records, dataset_id=auto_id, base_dir=base_dir, calibration_fraction=0.5, **_CREDS) + + manifest_path = _manifest_path(base_dir, auto_id) + manifest = PartitionManifest.model_validate_json(manifest_path.read_text()) + + mixed = sum( + 1 + for entry in manifest.assignments.values() + if 0 < sum(entry.retrieval_chunk_calibration.values()) < len(entry.retrieval_chunk_calibration) + ) + assert mixed > 0, "per-chunk independence not realised - no pair had a mixed retrieval bucket" + + cal_ds = client.datasets( + dataset_name(Task.RETRIEVAL, calibration=True, dataset_id=auto_id), + workspace="retrieval", + ) + prod_ds = client.datasets( + dataset_name(Task.RETRIEVAL, calibration=False, dataset_id=auto_id), + workspace="retrieval", + ) + assert cal_ds is not None and prod_ds is not None + + expected_cal_chunks = { + chunk_id + for entry in manifest.assignments.values() + for chunk_id, is_cal in entry.retrieval_chunk_calibration.items() + if is_cal + } + expected_prod_chunks = { + chunk_id + for entry in manifest.assignments.values() + for chunk_id, is_cal in entry.retrieval_chunk_calibration.items() + if not is_cal + } + cal_chunk_ids = {r.metadata["chunk_id"] for r in cal_ds.records} + prod_chunk_ids = {r.metadata["chunk_id"] for r in prod_ds.records} + assert cal_chunk_ids == expected_cal_chunks + assert prod_chunk_ids == expected_prod_chunks + assert cal_chunk_ids.isdisjoint(prod_chunk_ids) + finally: + teardown_resources(client, auto_settings) + + class TestPerWorkspaceMinSubmitted: """Inherited carve-out reaches the live Argilla dataset's distribution config.""" diff --git a/tests/unit/cli/test_cli_annotation.py b/tests/unit/cli/test_cli_annotation.py index c78212e4..9c52ea89 100644 --- a/tests/unit/cli/test_cli_annotation.py +++ b/tests/unit/cli/test_cli_annotation.py @@ -283,6 +283,36 @@ def test_partition_seed_default_is_unset(self, mock_import, tmp_path): assert kwargs["calibration_partition_seed"] is UNSET assert kwargs["calibration_min_submitted"] is UNSET assert kwargs["calibration_fraction"] is UNSET + assert kwargs["calibration_max_items"] is UNSET + + @patch("pragmata.annotation.import_records") + def test_calibration_max_items_threaded_through(self, mock_import, tmp_path): + mock_import.return_value = _empty_import_result() + records_file = tmp_path / "records.jsonl" + records_file.write_text("", encoding="utf-8") + + result = runner.invoke( + app, + ["annotation", "import", str(records_file), "--calibration-max-items", "200"], + ) + + assert result.exit_code == 0 + kwargs = mock_import.call_args.kwargs + assert kwargs["calibration_max_items"] == 200 + + @patch("pragmata.annotation.import_records") + def test_calibration_max_items_conflicts_with_no_calibration(self, mock_import, tmp_path): + records_file = tmp_path / "records.jsonl" + records_file.write_text("", encoding="utf-8") + + result = runner.invoke( + app, + ["annotation", "import", str(records_file), "--no-calibration", "--calibration-max-items", "50"], + ) + + assert result.exit_code == 2 + assert "cannot be combined" in result.output + mock_import.assert_not_called() @patch("pragmata.annotation.import_records") def test_locale_catalog_dir_threaded_through(self, mock_import, tmp_path): diff --git a/tests/unit/core/annotation/test_fan_out_records.py b/tests/unit/core/annotation/test_fan_out_records.py index c215e194..62dffdeb 100644 --- a/tests/unit/core/annotation/test_fan_out_records.py +++ b/tests/unit/core/annotation/test_fan_out_records.py @@ -33,6 +33,7 @@ def _partition( assignments=assignments, pairs_by_rid={derive_record_uuid(p): p for p in records}, calibration_fraction={t: 0.5 for t in Task}, + calibration_max_items={t: None for t in Task}, ) @@ -191,11 +192,11 @@ def test_dataset_counts_reflect_per_dataset_record_counts(self, patched_create_d assert cal_total == 3 * 3 # 3 cal records * 3 tasks assert prod_total == 2 * 3 # 2 prod records * 3 tasks - def test_skips_tasks_not_in_workspaces_topology(self, patched_create_dataset, caplog) -> None: + def test_records_for_task_absent_from_topology_raises(self, patched_create_dataset) -> None: client = MagicMock() - _, created = patched_create_dataset - # Topology only declares retrieval; grounding/generation are absent. + # Topology only declares retrieval; grounding / generation are absent + # but the partition manifest has assignments for them. partial_settings = AnnotationSettings( dataset_id="run1", workspaces={"retrieval": WorkspaceSettings(tasks={Task.RETRIEVAL: TaskSettings()})}, @@ -203,13 +204,8 @@ def test_skips_tasks_not_in_workspaces_topology(self, patched_create_dataset, ca records = [_make_pair("q0")] assignments = _assignments_with_uniform_calibration(records, is_cal=False) - with caplog.at_level("WARNING"): - result = fan_out_records(client, partial_settings, partition=_partition(records, assignments)) - - # Only retrieval gets a dataset; grounding and generation are skipped with a warning. - assert len(result) == 1 - assert any(ds_name.startswith("retrieval_") for (_, ds_name) in created) - assert "not in workspaces topology" in caplog.text + with pytest.raises(RuntimeError, match="not include this task"): + fan_out_records(client, partial_settings, partition=_partition(records, assignments)) def test_per_chunk_retrieval_routing(self, patched_create_dataset) -> None: """Different chunks of one record can route to different retrieval buckets.""" diff --git a/tests/unit/core/annotation/test_partition.py b/tests/unit/core/annotation/test_partition.py index 98af92a1..58e45b25 100644 --- a/tests/unit/core/annotation/test_partition.py +++ b/tests/unit/core/annotation/test_partition.py @@ -132,6 +132,109 @@ def test_retrieval_chunks_partitioned_independently(self, empty_manifest: Partit ) assert mixed > 0, "no pairs had per-chunk mixing - partition might still be per-record" + def test_cap_limits_calibration_count(self, empty_manifest: PartitionManifest) -> None: + """Per-task cap binds: realised count never exceeds cap.""" + # Grounding has its own cap of 5 out of 50 pairs at fraction 1.0. + settings = AnnotationSettings( + calibration_fraction=1.0, + workspaces={ + "g": WorkspaceSettings( + tasks={Task.GROUNDING: TaskSettings(calibration_max_items=5)}, + ), + "x": WorkspaceSettings(tasks={Task.GENERATION: TaskSettings()}), + "r": WorkspaceSettings(tasks={Task.RETRIEVAL: TaskSettings()}), + }, + ) + pairs = [_make_pair(f"q{i}") for i in range(50)] + partition = assign_partitions(pairs, manifest=empty_manifest, settings=settings, import_id="imp1") + + cal_count = sum(1 for e in partition.assignments.values() if e.grounding_generation_calibration[Task.GROUNDING]) + assert cal_count == 5 + + def test_cap_preserves_existing_assignments_when_over_cap(self, empty_manifest: PartitionManifest, caplog) -> None: + """If existing manifest count exceeds new cap, warn and don't demote.""" + now = datetime.now(timezone.utc) + for i in range(10): + rid = f"existing-{i}" + empty_manifest.assignments[rid] = PartitionManifestEntry( + grounding_generation_calibration={Task.GROUNDING: True, Task.GENERATION: False}, + retrieval_chunk_calibration={}, + import_id="prior", + calibration_fraction_at_import={t: 1.0 for t in Task}, + calibration_max_items_at_import={t: None for t in Task}, + assigned_at=now, + ) + # New cap is 5 - below the existing 10. + settings = AnnotationSettings( + calibration_fraction=1.0, + calibration_max_items=5, + workspaces={ + "g": WorkspaceSettings(tasks={Task.GROUNDING: TaskSettings()}), + "x": WorkspaceSettings(tasks={Task.GENERATION: TaskSettings()}), + "r": WorkspaceSettings(tasks={Task.RETRIEVAL: TaskSettings()}), + }, + ) + new_pairs = [_make_pair(f"new-{i}") for i in range(5)] + + with caplog.at_level("WARNING"): + assign_partitions(new_pairs, manifest=empty_manifest, settings=settings, import_id="imp2") + + existing_cal = sum( + 1 + for rid, e in empty_manifest.assignments.items() + if rid.startswith("existing-") and e.grounding_generation_calibration[Task.GROUNDING] + ) + assert existing_cal == 10 + new_cal = sum( + 1 + for rid, e in empty_manifest.assignments.items() + if not rid.startswith("existing-") and e.grounding_generation_calibration[Task.GROUNDING] + ) + assert new_cal == 0 + assert any("exceeds cap" in m for m in caplog.messages) + + def test_cap_under_split_imports_is_order_dependent_by_design(self, empty_manifest: PartitionManifest) -> None: + """Documented property: cap-binding split imports depend on order. + + Because the manifest is append-only, the final calibration set depends + on (corpus, seed, import_order) rather than (corpus, seed) alone. This + test pins the property so a future regression to order-independence is + a deliberate design change requiring this test's update. + """ + settings = AnnotationSettings( + calibration_fraction=1.0, + calibration_max_items=3, + workspaces={ + "g": WorkspaceSettings(tasks={Task.GROUNDING: TaskSettings()}), + "x": WorkspaceSettings(tasks={Task.GENERATION: TaskSettings()}), + "r": WorkspaceSettings(tasks={Task.RETRIEVAL: TaskSettings()}), + }, + ) + pairs = [_make_pair(f"q{i}") for i in range(10)] + + manifest_a = PartitionManifest( + dataset_id="test", + created_at=empty_manifest.created_at, + updated_at=empty_manifest.updated_at, + partition_seed=0, + ) + assign_partitions(pairs[:5], manifest=manifest_a, settings=settings, import_id="imp_a1") + assign_partitions(pairs[5:], manifest=manifest_a, settings=settings, import_id="imp_a2") + + manifest_b = PartitionManifest( + dataset_id="test", + created_at=empty_manifest.created_at, + updated_at=empty_manifest.updated_at, + partition_seed=0, + ) + assign_partitions(pairs, manifest=manifest_b, settings=settings, import_id="imp_b") + + def _cal_count(m: PartitionManifest) -> int: + return sum(1 for e in m.assignments.values() if e.grounding_generation_calibration[Task.GROUNDING]) + + assert _cal_count(manifest_a) == 3 + assert _cal_count(manifest_b) == 3 + def test_existing_assignments_preserved_across_reimports(self, empty_manifest: PartitionManifest) -> None: settings = _default_settings(calibration_fraction=1.0) pair = _make_pair("q0") @@ -143,6 +246,7 @@ def test_existing_assignments_preserved_across_reimports(self, empty_manifest: P retrieval_chunk_calibration={}, import_id="prior", calibration_fraction_at_import={t: 0.0 for t in Task}, + calibration_max_items_at_import={t: None for t in Task}, assigned_at=datetime.now(timezone.utc), ) @@ -152,13 +256,15 @@ def test_existing_assignments_preserved_across_reimports(self, empty_manifest: P assert partition.assignments[rid].grounding_generation_calibration[Task.GROUNDING] is False assert partition.assignments[rid].import_id == "prior" - def test_new_records_stamp_per_task_fraction(self, empty_manifest: PartitionManifest) -> None: - """Per-task fractions are stamped on each new entry.""" + def test_new_records_stamp_per_task_provenance(self, empty_manifest: PartitionManifest) -> None: + """Per-task fractions and caps are stamped on each new entry.""" settings = AnnotationSettings( calibration_fraction=0.5, + calibration_max_items=100, workspaces={ "r": WorkspaceSettings( calibration_fraction=0.1, + calibration_max_items=20, tasks={Task.RETRIEVAL: TaskSettings()}, ), "g": WorkspaceSettings(tasks={Task.GROUNDING: TaskSettings()}), @@ -172,6 +278,12 @@ def test_new_records_stamp_per_task_fraction(self, empty_manifest: PartitionMani assert entry.calibration_fraction_at_import[Task.RETRIEVAL] == 0.1 assert entry.calibration_fraction_at_import[Task.GROUNDING] == 0.5 + assert entry.calibration_max_items_at_import[Task.RETRIEVAL] == 20 + assert entry.calibration_max_items_at_import[Task.GROUNDING] == 100 + + # PartitionResult exposes the per-task resolution directly so api/ doesn't redo the walk. + assert partition.calibration_fraction[Task.RETRIEVAL] == 0.1 + assert partition.calibration_max_items[Task.RETRIEVAL] == 20 def test_calibration_fraction_resolves_per_task(self, empty_manifest: PartitionManifest) -> None: """PartitionResult.calibration_fraction resolves task → workspace → deployment.""" @@ -257,6 +369,32 @@ def test_backfills_task_added_to_topology(self, empty_manifest: PartitionManifes fresh = assign_partitions([pair], manifest=fresh_manifest, settings=expanded, import_id="imp3").assignments[rid] assert backfilled.grounding_generation_calibration == fresh.grounding_generation_calibration + def test_cap_binds_on_backfilled_units(self, empty_manifest: PartitionManifest) -> None: + """Backfilled calibration units compete for the cap budget like new ones.""" + pairs = [_make_pair(f"q{i}") for i in range(10)] + + # First import: retrieval-only, so grounding is unassigned on all 10 records. + retrieval_only = AnnotationSettings( + workspaces={"r": WorkspaceSettings(tasks={Task.RETRIEVAL: TaskSettings()})}, + ) + assign_partitions(pairs, manifest=empty_manifest, settings=retrieval_only, import_id="imp1") + + # Reimport with grounding added at fraction 1.0 but capped at 3. + expanded = AnnotationSettings( + calibration_fraction=1.0, + workspaces={ + "r": WorkspaceSettings(tasks={Task.RETRIEVAL: TaskSettings()}), + "g": WorkspaceSettings(tasks={Task.GROUNDING: TaskSettings(calibration_max_items=3)}), + }, + ) + result = assign_partitions(pairs, manifest=empty_manifest, settings=expanded, import_id="imp2") + + # Only 3 of the 10 backfilled grounding units win calibration; the cap binds. + cal = sum(1 for e in result.assignments.values() if e.grounding_generation_calibration[Task.GROUNDING]) + assert cal == 3 + # All 10 now carry a grounding flag (backfilled), so none re-backfill next run. + assert all(Task.GROUNDING in e.grounding_generation_calibration for e in result.assignments.values()) + class TestManifestIO: def test_load_empty_when_missing(self, tmp_path: Path) -> None: @@ -275,6 +413,7 @@ def test_round_trip(self, tmp_path: Path) -> None: retrieval_chunk_calibration={"chunk-a": True, "chunk-b": False}, import_id="imp1", calibration_fraction_at_import={t: 0.1 for t in Task}, + calibration_max_items_at_import={t: None for t in Task}, assigned_at=datetime(2026, 4, 22, tzinfo=timezone.utc), ) diff --git a/tests/unit/core/schemas/test_annotation_import.py b/tests/unit/core/schemas/test_annotation_import.py index 7bcddc1b..4b3be9c3 100644 --- a/tests/unit/core/schemas/test_annotation_import.py +++ b/tests/unit/core/schemas/test_annotation_import.py @@ -152,6 +152,11 @@ def valid_entry_kwargs(): Task.GENERATION: 0.1, Task.RETRIEVAL: 0.1, }, + "calibration_max_items_at_import": { + Task.GROUNDING: None, + Task.GENERATION: None, + Task.RETRIEVAL: None, + }, "assigned_at": _ENTRY_NOW, } @@ -183,6 +188,11 @@ def test_fraction_negative_rejected(self, valid_entry_kwargs): with pytest.raises(ValidationError): PartitionManifestEntry(**valid_entry_kwargs) + def test_cap_zero_rejected(self, valid_entry_kwargs): + valid_entry_kwargs["calibration_max_items_at_import"][Task.GROUNDING] = 0 + with pytest.raises(ValidationError): + PartitionManifestEntry(**valid_entry_kwargs) + def test_extra_fields_rejected(self, valid_entry_kwargs): valid_entry_kwargs["unknown"] = "x" with pytest.raises(ValidationError): diff --git a/tests/unit/core/settings/test_annotation_settings.py b/tests/unit/core/settings/test_annotation_settings.py index 5ca926ad..a0313216 100644 --- a/tests/unit/core/settings/test_annotation_settings.py +++ b/tests/unit/core/settings/test_annotation_settings.py @@ -420,6 +420,55 @@ def test_zero_at_task_disables_just_that_task(self): assert s.resolved_task("g", Task.GROUNDING).calibration_fraction == 0.1 +class TestCalibrationMaxRecordsInheritance: + """``calibration_max_items`` is inheritable across deployment / workspace / task.""" + + def test_deployment_default_is_uncapped(self): + s = AnnotationSettings( + workspaces={"r": WorkspaceSettings(tasks={Task.RETRIEVAL: TaskSettings()})}, + ) + assert s.resolved_task("r", Task.RETRIEVAL).calibration_max_items is None + + def test_deployment_cap_flows_to_task(self): + s = AnnotationSettings( + calibration_max_items=200, + workspaces={"r": WorkspaceSettings(tasks={Task.RETRIEVAL: TaskSettings()})}, + ) + assert s.resolved_task("r", Task.RETRIEVAL).calibration_max_items == 200 + + def test_workspace_cap_overrides_deployment(self): + s = AnnotationSettings( + calibration_max_items=200, + workspaces={ + "r": WorkspaceSettings( + calibration_max_items=50, + tasks={Task.RETRIEVAL: TaskSettings()}, + ), + }, + ) + assert s.resolved_task("r", Task.RETRIEVAL).calibration_max_items == 50 + + def test_task_cap_overrides_workspace_and_deployment(self): + s = AnnotationSettings( + calibration_max_items=200, + workspaces={ + "r": WorkspaceSettings( + calibration_max_items=100, + tasks={Task.RETRIEVAL: TaskSettings(calibration_max_items=20)}, + ), + }, + ) + assert s.resolved_task("r", Task.RETRIEVAL).calibration_max_items == 20 + + def test_zero_cap_rejected(self): + with pytest.raises(ValidationError): + AnnotationSettings(calibration_max_items=0) + + def test_negative_cap_rejected(self): + with pytest.raises(ValidationError): + AnnotationSettings(calibration_max_items=-1) + + class TestPerTaskTopologyValidator: """``_check_calibration_topology`` walks per-(workspace, task) using resolved values."""