From b3b3d394e558baf694e64e0f53cf17b3a432a0f9 Mon Sep 17 00:00:00 2001 From: cirogam22 Date: Sun, 5 Jul 2026 03:00:26 +0200 Subject: [PATCH 1/3] refactor(output): pin schema_version + int/float wire invariants (closes #152, #153) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two stage-2 follow-ups to PR #149's typed aggregate models. Both pin wire-format invariants the future writer migration will depend on. schema_version survives exclude_unset=True (#152). The current writer path (json.dump(dict)) relies on the orchestrator manually stamping aggregate["schema_version"] = 1 before write. Once the writer migrates to model.model_dump_json(exclude_unset=True), a RunAggregate whose schema_version was left at the model default would silently drop the envelope field — every downstream dashboard that dispatches on schema_version would break. Fix: @model_serializer(mode="wrap") on RunAggregate unconditionally injects schema_version into the dumped dict, regardless of dump options. The field's default stays 1; callers can still pass it explicitly. Idiomatic Pydantic v2 pattern. Numeric fields preserve source type on JSON (#153). The producer in metrics.py emits sum([]) == 0 (int) for empty-aggregate total_* / avg_* fields and float otherwise. A pre-fix model with float fields coerces 0 → 0.0 on validation, and model_dump_json emits 0.0 where the source dict emits 0 — a wire-format break invisible to Python-dict comparison (0 == 0.0) but real when downstream tooling diffs the JSON byte-for-byte. Fix: change numeric fields on PerTaskMetrics and AggregateMetrics from float to int | float unions. Pydantic v2's smart-union picks the more specific type on validation, so a source int round-trips as int and a source float as float. Defaults change from 0.0 to 0 to match the empty-aggregate source. Fields staying float are those that come from a division (success_rate, pass@k, avg_score) — always float-valued from the producer. Writer docstring on RunAggregateWriter.write_aggregate now documents both invariants explicitly so future consumers can rely on them without needing to read the model source. Two new canonical tests in tests/canonical/test_run_aggregate_models_snapshot.py: - test_schema_version_survives_exclude_unset — constructs RunAggregate from a payload that omits schema_version, dumps with exclude_unset=True, asserts the field is present in the output. Also verifies the JSON round-trip preserves the value. - test_int_float_json_string_no_drift — for each of the four writer artifacts, produces a real dict from the metric-calc functions on an empty-aggregate trajectory (the harshest case for int/float drift), feeds it through the model, and compares byte-canonical JSON strings between the dict path and the model path. Divergences fail loud with the exact source/model JSON diff. Verification: - make lint clean. - uv run pytest tests/canonical/test_run_aggregate_models_snapshot.py -v — all 12 tests pass (10 existing + 2 new). - uv run pytest tests/ -m "unit or canonical" -q — 2447 pass. - Grep confirms zero downstream consumers reference the model classes directly today; the int | float widening is safe. --- .../test_run_aggregate_models_snapshot.py | 118 +++++++++++++ tolokaforge/core/output/aggregate_models.py | 161 +++++++++++------- tolokaforge/core/output/aggregates.py | 16 +- 3 files changed, 233 insertions(+), 62 deletions(-) diff --git a/tests/canonical/test_run_aggregate_models_snapshot.py b/tests/canonical/test_run_aggregate_models_snapshot.py index 9696f1db..ecc58699 100644 --- a/tests/canonical/test_run_aggregate_models_snapshot.py +++ b/tests/canonical/test_run_aggregate_models_snapshot.py @@ -345,3 +345,121 @@ def test_failure_attribution_envelope_round_trip() -> None: } _round_trip(FailureAttribution, payload) + + +# --------------------------------------------------------------------------- +# Wire-format invariants (closes #152, #153) +# --------------------------------------------------------------------------- + + +def test_schema_version_survives_exclude_unset() -> None: + """``RunAggregate.schema_version`` must appear in the dumped output + even when the model is constructed without an explicit value AND + dumped with ``exclude_unset=True`` (the mode the migrated writer + will use to preserve wire-format parity with the current dict path). + + The regression this catches: once the writer migrates from + ``json.dump(dict)`` to ``model.model_dump_json(exclude_unset=True)``, + a ``RunAggregate`` whose ``schema_version`` was left at the model + default would silently drop the envelope field — every downstream + dashboard that dispatches on ``schema_version`` would break. The + model's ``@model_serializer`` forces the field into the dump + regardless. Closes #152. + """ + # Payload deliberately omits ``schema_version`` — the model must + # still emit it via the default. + payload = calculate_aggregate_metrics(_sample_task_metrics_list(), weighted=True) + assert "schema_version" not in payload, "guard: payload must not stamp schema_version" + + model = RunAggregate.model_validate(payload) + dumped = model.model_dump(by_alias=True, mode="json", exclude_unset=True) + + assert "schema_version" in dumped, ( + "schema_version must survive exclude_unset=True on RunAggregate — " + f"got dumped keys: {sorted(dumped.keys())}" + ) + assert ( + dumped["schema_version"] == 1 + ), f"schema_version default drifted; expected 1, got {dumped['schema_version']!r}" + + # And it survives a JSON round-trip too. + reloaded = RunAggregate.model_validate_json( + model.model_dump_json(by_alias=True, exclude_unset=True) + ) + assert reloaded.schema_version == 1 + + +def test_int_float_json_string_no_drift() -> None: + """The producer emits ``int`` for empty-sum aggregates + (``sum([]) == 0``) and ``float`` for populated ones. The model's + ``int | float`` unions preserve the source type on validation so + the dumped JSON is byte-identical to the source dict's JSON — no + ``0`` → ``0.0`` drift on the wire. + + A pre-fix version of the model with `float` fields would coerce + ``0`` to ``0.0`` on validation, and the JSON dump would carry + ``0.0`` where the source dict carries ``0`` — a wire-format break + invisible to Python-dict comparison (``0 == 0.0``) but real when + downstream tooling diffs the JSON files byte-for-byte. Closes #153. + """ + import json + + def _canonical(payload: dict[str, Any]) -> str: + return json.dumps(payload, sort_keys=True, default=str) + + # Empty aggregates — the harshest test for int/float drift because + # every ``total_*`` / ``avg_*`` is ``sum([]) == 0`` on the source. + empty_traj = _make_trajectory() + task_payload = calculate_task_metrics([empty_traj]) + _augment_task_metrics(task_payload, task_id="task-int-float") + + # PerTaskMetrics — source dict JSON vs. model-dumped JSON must match + # byte-for-byte. + model = PerTaskMetrics.model_validate(task_payload) + dumped_dict = model.model_dump(by_alias=True, mode="json", exclude_unset=True) + source_json = _canonical(task_payload) + dumped_json = _canonical(dumped_dict) + assert source_json == dumped_json, ( + "PerTaskMetrics JSON drift between dict path and model path.\n" + f"source: {source_json}\n" + f"model: {dumped_json}" + ) + + # AggregateMetrics — same story; empty-aggregate `total_*_tokens` + # start as ``int`` from the producer. + agg_payload = calculate_aggregate_metrics([task_payload], weighted=True) + agg_model = AggregateMetrics.model_validate(agg_payload) + agg_dumped = agg_model.model_dump(by_alias=True, mode="json", exclude_unset=True) + assert _canonical(agg_payload) == _canonical(agg_dumped), ( + "AggregateMetrics JSON drift between dict path and model path.\n" + f"source: {_canonical(agg_payload)}\n" + f"model: {_canonical(agg_dumped)}" + ) + + # RunAggregate — model injects `schema_version`. Compare with the + # source dict stamped to match. + run_payload = dict(agg_payload) + run_payload["schema_version"] = 1 + run_model = RunAggregate.model_validate(run_payload) + run_dumped = run_model.model_dump(by_alias=True, mode="json", exclude_unset=True) + assert _canonical(run_payload) == _canonical(run_dumped), ( + "RunAggregate JSON drift between dict path and model path.\n" + f"source: {_canonical(run_payload)}\n" + f"model: {_canonical(run_dumped)}" + ) + + # FailureAttribution — the envelope isn't a numeric-heavy shape, + # but we exercise it to prove the invariant holds for every + # writer artifact. + attributions = [attribute_failure(empty_traj)] + fa_payload = { + "summary": summarize_failure_attributions(attributions), + "failures": attributions, + } + fa_model = FailureAttribution.model_validate(fa_payload) + fa_dumped = fa_model.model_dump(by_alias=True, mode="json", exclude_unset=True) + assert _canonical(fa_payload) == _canonical(fa_dumped), ( + "FailureAttribution JSON drift between dict path and model path.\n" + f"source: {_canonical(fa_payload)}\n" + f"model: {_canonical(fa_dumped)}" + ) diff --git a/tolokaforge/core/output/aggregate_models.py b/tolokaforge/core/output/aggregate_models.py index 876d39f8..619a1ca4 100644 --- a/tolokaforge/core/output/aggregate_models.py +++ b/tolokaforge/core/output/aggregate_models.py @@ -24,13 +24,24 @@ The ``pass@k`` / ``pass_hat@k`` keys aren't valid Python identifiers, so they carry :class:`~pydantic.Field` ``alias`` values. Dump with ``model_dump(by_alias=True, mode="json")`` to produce the wire shape. + +Two wire-format invariants pinned by the canonical tests: + +1. ``schema_version`` on :class:`RunAggregate` is **always emitted** on + dump, regardless of ``exclude_unset``. See the model-serializer on + the class for details. +2. Numeric fields carry ``int | float`` unions so the wire type + matches what the producer emitted (``sum([]) == 0`` from an empty + aggregate stays ``int`` on the wire; populated fields stay + ``float``). This preserves byte-identical JSON between the + dict-based writer path and any future model-based writer path. """ from __future__ import annotations from typing import Any -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, SerializationInfo, model_serializer __all__ = [ "AggregateMetrics", @@ -76,42 +87,48 @@ class PerTaskMetrics(BaseModel): pass_hat_at_5: float | None = Field(default=None, alias="pass_hat@5") pass_hat_at_10: float | None = Field(default=None, alias="pass_hat@10") - # Averages across trials. - avg_score: float - avg_latency_s: float - avg_turns: float - avg_tool_calls: float + # Averages across trials. Typed as ``int | float`` unions so the + # model preserves the numeric type the producer emitted: e.g. + # ``sum([]) == 0`` (``int``) for empty aggregates, ``float`` + # otherwise. Pydantic's smart-union picks the more specific type on + # validation, so a source ``0`` round-trips to ``0`` (not ``0.0``) + # in the JSON dump — pinning the wire format against the current + # dict-based writer path. See :class:`RunAggregateWriter` docstring. + avg_score: int | float + avg_latency_s: int | float + avg_turns: int | float + avg_tool_calls: int | float # Per-usage-field averages (Metrics.usage — six fields). - avg_prompt_tokens: float = 0.0 - avg_completion_tokens: float = 0.0 - avg_reasoning_tokens: float = 0.0 - avg_cached_tokens: float = 0.0 - avg_cache_creation_input_tokens: float = 0.0 - avg_cache_read_input_tokens: float = 0.0 + avg_prompt_tokens: int | float = 0 + avg_completion_tokens: int | float = 0 + avg_reasoning_tokens: int | float = 0 + avg_cached_tokens: int | float = 0 + avg_cache_creation_input_tokens: int | float = 0 + avg_cache_read_input_tokens: int | float = 0 # Cost is ``None`` when no trial reported a cost (unknown-provider case). - total_cost_usd: float | None = None - avg_cost_usd: float | None = None + total_cost_usd: int | float | None = None + avg_cost_usd: int | float | None = None # Judge-cost split — ``judge_cost_usd`` is the LLM-judge grader's # spend, tracked separately so the agent-vs-judge cost breakdown # survives round-trips. ``total_cost_incl_judge_usd`` = agent + # judge; ``None`` when neither is known. - judge_cost_usd: float | None = None - total_cost_incl_judge_usd: float | None = None + judge_cost_usd: int | float | None = None + total_cost_incl_judge_usd: int | float | None = None # Per-trial wall-time percentiles. - latency_p50_s: float = 0.0 - latency_p90_s: float = 0.0 - latency_p99_s: float = 0.0 + latency_p50_s: int | float = 0 + latency_p90_s: int | float = 0 + latency_p99_s: int | float = 0 # Per-API-call percentiles aggregated across every call in every trial. - api_call_latency_p50_s: float = 0.0 - api_call_latency_p90_s: float = 0.0 - api_call_latency_p99_s: float = 0.0 + api_call_latency_p50_s: int | float = 0 + api_call_latency_p90_s: int | float = 0 + api_call_latency_p99_s: int | float = 0 - stuck_rate: float = 0.0 + stuck_rate: int | float = 0 class AggregateMetrics(BaseModel): @@ -130,66 +147,90 @@ class AggregateMetrics(BaseModel): total_tasks: int total_trials: int - # Weighted (micro) averages — set when ``weighted=True``. - success_rate_micro: float | None = None - avg_score_micro: float | None = None + # Weighted (micro) averages — set when ``weighted=True``. Union of + # ``int | float | None`` so an empty-aggregate ``int 0`` from the + # producer preserves the JSON wire type verbatim (see + # :class:`PerTaskMetrics` for the same convention). + success_rate_micro: int | float | None = None + avg_score_micro: int | float | None = None # Unweighted (macro) averages — set when ``weighted=False``. - success_rate_macro: float | None = None - avg_score_macro: float | None = None + success_rate_macro: int | float | None = None + avg_score_macro: int | float | None = None # pass@k macro averages across tasks, plus their pass_hat aliases. - pass_at_1_macro: float | None = Field(default=None, alias="pass@1_macro") - pass_at_5_macro: float | None = Field(default=None, alias="pass@5_macro") - pass_at_10_macro: float | None = Field(default=None, alias="pass@10_macro") - pass_hat_at_1_macro: float | None = Field(default=None, alias="pass_hat@1_macro") - pass_hat_at_5_macro: float | None = Field(default=None, alias="pass_hat@5_macro") - pass_hat_at_10_macro: float | None = Field(default=None, alias="pass_hat@10_macro") + pass_at_1_macro: int | float | None = Field(default=None, alias="pass@1_macro") + pass_at_5_macro: int | float | None = Field(default=None, alias="pass@5_macro") + pass_at_10_macro: int | float | None = Field(default=None, alias="pass@10_macro") + pass_hat_at_1_macro: int | float | None = Field(default=None, alias="pass_hat@1_macro") + pass_hat_at_5_macro: int | float | None = Field(default=None, alias="pass_hat@5_macro") + pass_hat_at_10_macro: int | float | None = Field(default=None, alias="pass_hat@10_macro") # Simple cross-task averages. - avg_latency_s: float = 0.0 - avg_turns: float = 0.0 - avg_tool_calls: float = 0.0 - stuck_rate: float = 0.0 + avg_latency_s: int | float = 0 + avg_turns: int | float = 0 + avg_tool_calls: int | float = 0 + stuck_rate: int | float = 0 # Per-usage-field totals + macro averages. - total_prompt_tokens: float = 0.0 - total_completion_tokens: float = 0.0 - total_reasoning_tokens: float = 0.0 - total_cached_tokens: float = 0.0 - total_cache_creation_input_tokens: float = 0.0 - total_cache_read_input_tokens: float = 0.0 - avg_prompt_tokens: float = 0.0 - avg_completion_tokens: float = 0.0 - avg_reasoning_tokens: float = 0.0 - avg_cached_tokens: float = 0.0 - avg_cache_creation_input_tokens: float = 0.0 - avg_cache_read_input_tokens: float = 0.0 - - total_cost_usd: float | None = None - avg_cost_usd: float | None = None + total_prompt_tokens: int | float = 0 + total_completion_tokens: int | float = 0 + total_reasoning_tokens: int | float = 0 + total_cached_tokens: int | float = 0 + total_cache_creation_input_tokens: int | float = 0 + total_cache_read_input_tokens: int | float = 0 + avg_prompt_tokens: int | float = 0 + avg_completion_tokens: int | float = 0 + avg_reasoning_tokens: int | float = 0 + avg_cached_tokens: int | float = 0 + avg_cache_creation_input_tokens: int | float = 0 + avg_cache_read_input_tokens: int | float = 0 + + total_cost_usd: int | float | None = None + avg_cost_usd: int | float | None = None # Judge-cost split at the run/slice level — same shape as # ``PerTaskMetrics`` but aggregated across the tasks/slice. - judge_cost_usd: float | None = None - total_cost_incl_judge_usd: float | None = None + judge_cost_usd: int | float | None = None + total_cost_incl_judge_usd: int | float | None = None - latency_p50_s_macro: float = 0.0 - latency_p90_s_macro: float = 0.0 - latency_p99_s_macro: float = 0.0 + latency_p50_s_macro: int | float = 0 + latency_p90_s_macro: int | float = 0 + latency_p99_s_macro: int | float = 0 class RunAggregate(AggregateMetrics): """The top-level ``aggregate.json`` shape — :class:`AggregateMetrics` - plus the ``schema_version`` envelope field the orchestrator stamps - before writing. + plus the ``schema_version`` envelope field every downstream consumer + (dashboards, metric collectors, historical-run readers) reads to + dispatch between wire-format generations. Slice values inside ``metadata_slices.json`` are :class:`AggregateMetrics` directly; only the top-level artifact carries the envelope. + + ``schema_version`` is **always emitted** on dump, regardless of + whether the caller constructs the model with an explicit value or + lets the default apply. Pydantic's ``exclude_unset=True`` normally + drops fields left at their default — for this envelope that would + silently ship an aggregate.json without a version, which would + break every version-dispatching downstream consumer. The + :func:`_always_include_schema_version` model-serializer below + guarantees the field is present regardless of dump options. """ schema_version: int = 1 + @model_serializer(mode="wrap") + def _always_include_schema_version( + self, handler: Any, info: SerializationInfo + ) -> dict[str, Any]: + """Wrap the default serializer so ``schema_version`` is always + present in the dumped dict, even under ``exclude_unset=True``.""" + data = handler(self) + if isinstance(data, dict) and "schema_version" not in data: + data["schema_version"] = self.schema_version + return data + class MetadataSlices(BaseModel): """The ``metadata_slices.json`` envelope — four ``by_*`` dictionaries. diff --git a/tolokaforge/core/output/aggregates.py b/tolokaforge/core/output/aggregates.py index e4d0bd20..5bb7d81d 100644 --- a/tolokaforge/core/output/aggregates.py +++ b/tolokaforge/core/output/aggregates.py @@ -66,8 +66,20 @@ def write_aggregate( ) -> None: """Persist ``output_dir/aggregate.json`` from *aggregate*. - Callers are responsible for injecting any envelope fields the - artifact carries (e.g. ``schema_version``) before calling. + Wire-format invariants (also pinned by the canonical tests on + :class:`~tolokaforge.core.output.aggregate_models.RunAggregate`): + + * **``schema_version`` is always present** on the written JSON. + Callers passing a dict must set it before calling. Callers + passing a :class:`~tolokaforge.core.output.aggregate_models.RunAggregate` + model get this for free — the model's serializer forces the + field into the dump regardless of ``exclude_unset``. + * **Numeric fields preserve source type.** Values the producer + emitted as ``int`` (e.g. ``sum([]) == 0`` for empty + aggregates) stay ``int`` on the wire; values emitted as + ``float`` stay ``float``. The model achieves this via + ``int | float`` unions; the direct dict path preserves it + naturally. """ ... From d7f2e3b23345a4d21dd58c7b5f018bec7a834359 Mon Sep 17 00:00:00 2001 From: cirogam22 Date: Sun, 5 Jul 2026 03:13:42 +0200 Subject: [PATCH 2/3] review: narrow rate-shaped fields + rewrite int/float test to exercise int inputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Local review pointed out three real issues on the previous commit: 1. **Widening was inconsistent.** ``AggregateMetrics.success_rate_micro/ macro``, ``avg_score_micro/macro``, and ``pass_at_*_macro`` were widened to ``int | float | None`` in the previous commit, but every one of them is ``sum(...) / n`` in metrics.py — always float at the producer, matching the ``PerTaskMetrics.success_rate``/``avg_score``/ ``pass@k`` siblings that stayed narrow ``float``. Narrowed the AggregateMetrics rate fields back to ``float | None`` for symmetry. Fields that stay widened (``total_*_tokens``, ``avg_*_tokens``, latency percentiles, ``stuck_rate``, cost sums) are the ones with naturally-integer semantics. 2. **The `int/float` test didn't exercise int inputs.** The previous ``test_int_float_json_string_no_drift`` fed ``calculate_task_metrics([single_trajectory])`` and compared JSON strings. But every widened field in that path resolves to ``float`` via division — the pre-fix ``float``-only model produces byte-identical output. The test guarded nothing. Replaced with three sharper tests: - ``test_int_valued_numeric_fields_preserve_int_type`` — hand-injects ``int`` values into every widened field of ``PerTaskMetrics`` and asserts the model preserves ``int`` type at both the Python level (``isinstance(value, int)``) and the JSON wire level (dump keys are ``int``, not ``float``). Would fail if any widened field is narrowed back to ``float``. - ``test_int_valued_aggregate_fields_preserve_int_type`` — same invariant for ``AggregateMetrics``. Also asserts that ``success_rate_micro`` STAYS narrow ``float`` — pinning the symmetry with ``PerTaskMetrics``. - ``test_current_producer_output_matches_model_dump_byte_for_byte`` — round-trip guard for the current live producer output across weighted and unweighted aggregates. Catches accidental coercion in the narrow rate fields. 3. **Unused ``info`` parameter** on ``_always_include_schema_version``. Dropped along with the ``SerializationInfo`` import. Clarified the docstring: the serializer injects only when the wrapped handler omits the field (i.e. under ``exclude_unset=True`` with the default), never overwrites a caller-set value. Also updated the module docstring to spell out which fields ARE widened (``int``-natural: tokens, latency, stuck) vs. which stay narrow (rate-shaped divisions) and acknowledge that this is defensive future-proofing — today's producer path always emits float for every widened field because the metric-calc functions short-circuit on empty input. Verification: 14 tests pass on the snapshot file (up from 12); full unit + canonical suite: 2449 pass. --- .../test_run_aggregate_models_snapshot.py | 225 +++++++++++++----- tolokaforge/core/output/aggregate_models.py | 71 ++++-- 2 files changed, 206 insertions(+), 90 deletions(-) diff --git a/tests/canonical/test_run_aggregate_models_snapshot.py b/tests/canonical/test_run_aggregate_models_snapshot.py index ecc58699..29d84ad6 100644 --- a/tests/canonical/test_run_aggregate_models_snapshot.py +++ b/tests/canonical/test_run_aggregate_models_snapshot.py @@ -389,77 +389,172 @@ def test_schema_version_survives_exclude_unset() -> None: assert reloaded.schema_version == 1 -def test_int_float_json_string_no_drift() -> None: - """The producer emits ``int`` for empty-sum aggregates - (``sum([]) == 0``) and ``float`` for populated ones. The model's - ``int | float`` unions preserve the source type on validation so - the dumped JSON is byte-identical to the source dict's JSON — no - ``0`` → ``0.0`` drift on the wire. - - A pre-fix version of the model with `float` fields would coerce - ``0`` to ``0.0`` on validation, and the JSON dump would carry - ``0.0`` where the source dict carries ``0`` — a wire-format break - invisible to Python-dict comparison (``0 == 0.0``) but real when - downstream tooling diffs the JSON files byte-for-byte. Closes #153. +def test_int_valued_numeric_fields_preserve_int_type() -> None: + """The ``int | float`` unions on token-count / latency-percentile / + cost-sum fields must preserve ``int`` inputs verbatim through + validation AND JSON dump. A regression that narrows any of these + fields back to ``float`` would coerce ``int 42`` → ``42.0`` on + validation and emit ``42.0`` in the JSON — a byte-level wire drift + invisible to Python-dict comparison (``42 == 42.0``) but real + when downstream tooling diffs the JSON files byte-for-byte. + + Today's producers in ``metrics.py`` happen to always emit + ``float`` for these fields (division short-circuits on empty + input), so no live regression exists. This test guards against a + future producer refactor introducing ``int`` output — the model + must handle it correctly without a coordinated widening. Closes #153. """ import json - def _canonical(payload: dict[str, Any]) -> str: - return json.dumps(payload, sort_keys=True, default=str) - - # Empty aggregates — the harshest test for int/float drift because - # every ``total_*`` / ``avg_*`` is ``sum([]) == 0`` on the source. - empty_traj = _make_trajectory() - task_payload = calculate_task_metrics([empty_traj]) - _augment_task_metrics(task_payload, task_id="task-int-float") - - # PerTaskMetrics — source dict JSON vs. model-dumped JSON must match - # byte-for-byte. + # Hand-construct a payload with int-valued token counts + latency + # percentiles. The producer path today never emits these as int, + # so this is the only way to exercise the int branch of the union. + task_payload: dict[str, Any] = { + "task_id": "int-preserve-guard", + "benchmark_type": "airline", + "complexity": "simple", + "tags": [], + "expected_failure_modes": [], + "total_trials": 1, + "successful_trials": 1, + "success_rate": 1.0, # rate — stays float + "avg_score": 1.0, # rate — stays float + "avg_latency_s": 5, # int — union widened field + "avg_turns": 3, # int — union widened field + "avg_tool_calls": 2, # int — union widened field + # Token counts — natural integers. + "avg_prompt_tokens": 100, + "avg_completion_tokens": 40, + "avg_reasoning_tokens": 10, + "avg_cached_tokens": 5, + "avg_cache_creation_input_tokens": 3, + "avg_cache_read_input_tokens": 2, + # Latency percentiles — could be int seconds. + "latency_p50_s": 4, + "latency_p90_s": 6, + "latency_p99_s": 8, + "api_call_latency_p50_s": 1, + "api_call_latency_p90_s": 2, + "api_call_latency_p99_s": 3, + "stuck_rate": 0, + } model = PerTaskMetrics.model_validate(task_payload) - dumped_dict = model.model_dump(by_alias=True, mode="json", exclude_unset=True) - source_json = _canonical(task_payload) - dumped_json = _canonical(dumped_dict) - assert source_json == dumped_json, ( - "PerTaskMetrics JSON drift between dict path and model path.\n" - f"source: {source_json}\n" - f"model: {dumped_json}" - ) - # AggregateMetrics — same story; empty-aggregate `total_*_tokens` - # start as ``int`` from the producer. - agg_payload = calculate_aggregate_metrics([task_payload], weighted=True) - agg_model = AggregateMetrics.model_validate(agg_payload) - agg_dumped = agg_model.model_dump(by_alias=True, mode="json", exclude_unset=True) - assert _canonical(agg_payload) == _canonical(agg_dumped), ( - "AggregateMetrics JSON drift between dict path and model path.\n" - f"source: {_canonical(agg_payload)}\n" - f"model: {_canonical(agg_dumped)}" - ) + # Type preservation at the Python level — union picks int on validation. + for field in ( + "avg_latency_s", + "avg_turns", + "avg_tool_calls", + "avg_prompt_tokens", + "avg_completion_tokens", + "avg_reasoning_tokens", + "avg_cached_tokens", + "avg_cache_creation_input_tokens", + "avg_cache_read_input_tokens", + "latency_p50_s", + "latency_p90_s", + "latency_p99_s", + "api_call_latency_p50_s", + "api_call_latency_p90_s", + "api_call_latency_p99_s", + "stuck_rate", + ): + value = getattr(model, field) + assert isinstance(value, int) and not isinstance(value, bool), ( + f"PerTaskMetrics.{field}: int input coerced to {type(value).__name__} — " + f"the int|float union is broken. Got {value!r}." + ) - # RunAggregate — model injects `schema_version`. Compare with the - # source dict stamped to match. - run_payload = dict(agg_payload) - run_payload["schema_version"] = 1 - run_model = RunAggregate.model_validate(run_payload) - run_dumped = run_model.model_dump(by_alias=True, mode="json", exclude_unset=True) - assert _canonical(run_payload) == _canonical(run_dumped), ( - "RunAggregate JSON drift between dict path and model path.\n" - f"source: {_canonical(run_payload)}\n" - f"model: {_canonical(run_dumped)}" - ) + # Type preservation on the JSON wire — dump keys must be int, not float. + dumped = model.model_dump(by_alias=True, mode="json") + dumped_json = json.dumps(dumped, sort_keys=True) + # An int field dumped as int has no trailing ``.0`` in the JSON string. + for field in ("avg_prompt_tokens", "latency_p50_s", "stuck_rate"): + assert isinstance(dumped[field], int) and not isinstance(dumped[field], bool), ( + f"PerTaskMetrics.{field} dumped as {type(dumped[field]).__name__}, " + f"expected int. Full dump: {dumped_json}" + ) - # FailureAttribution — the envelope isn't a numeric-heavy shape, - # but we exercise it to prove the invariant holds for every - # writer artifact. - attributions = [attribute_failure(empty_traj)] - fa_payload = { - "summary": summarize_failure_attributions(attributions), - "failures": attributions, + +def test_int_valued_aggregate_fields_preserve_int_type() -> None: + """Same invariant as + :func:`test_int_valued_numeric_fields_preserve_int_type` but for the + ``AggregateMetrics`` shape. Divison-produced rate fields on + ``AggregateMetrics`` (``success_rate_micro/macro``, ``avg_score_*``, + ``pass_at_*_macro``) stay narrow ``float`` — a caller passing ``int`` + there is a producer bug, not a wire invariant to preserve, so we + do NOT assert int preservation for those fields.""" + agg_payload: dict[str, Any] = { + "total_tasks": 2, + "total_trials": 5, + "success_rate_micro": 1.0, # rate — stays float + "avg_score_micro": 1.0, # rate — stays float + "avg_latency_s": 5, # int — widened + "avg_turns": 3, # int — widened + "avg_tool_calls": 2, # int — widened + "stuck_rate": 0, # int — widened + "total_prompt_tokens": 500, + "total_completion_tokens": 200, + "total_reasoning_tokens": 50, + "total_cached_tokens": 25, + "total_cache_creation_input_tokens": 15, + "total_cache_read_input_tokens": 10, + "avg_prompt_tokens": 100, + "avg_completion_tokens": 40, + "avg_reasoning_tokens": 10, + "avg_cached_tokens": 5, + "avg_cache_creation_input_tokens": 3, + "avg_cache_read_input_tokens": 2, + "latency_p50_s_macro": 4, + "latency_p90_s_macro": 6, + "latency_p99_s_macro": 8, } - fa_model = FailureAttribution.model_validate(fa_payload) - fa_dumped = fa_model.model_dump(by_alias=True, mode="json", exclude_unset=True) - assert _canonical(fa_payload) == _canonical(fa_dumped), ( - "FailureAttribution JSON drift between dict path and model path.\n" - f"source: {_canonical(fa_payload)}\n" - f"model: {_canonical(fa_dumped)}" + model = AggregateMetrics.model_validate(agg_payload) + for field in ( + "avg_latency_s", + "avg_turns", + "avg_tool_calls", + "stuck_rate", + "total_prompt_tokens", + "avg_prompt_tokens", + "latency_p50_s_macro", + ): + value = getattr(model, field) + assert isinstance(value, int) and not isinstance(value, bool), ( + f"AggregateMetrics.{field}: int input coerced to {type(value).__name__}. " + f"Got {value!r}." + ) + + # And rate fields DID stay narrow float — a regression that widens them + # would be caught by mypy at consumer sites, but a runtime guard here + # documents the invariant. + assert isinstance(model.success_rate_micro, float), ( + "AggregateMetrics.success_rate_micro widened to int|float unexpectedly — " + "rate-shaped fields must stay narrow float for consumer type contracts." ) + + +def test_current_producer_output_matches_model_dump_byte_for_byte() -> None: + """Round-trip guard against the current live producer output. + + Even though the ``int | float`` union covers a case the producer + doesn't emit today, the by-product must still hold: every dict the + production metric-calc functions actually produce today survives a + model round-trip byte-for-byte in the JSON representation. This + catches any accidental type coercion in the round trip for the + fields we DID narrow (``success_rate_*``, ``avg_score_*``, + ``pass_at_*_macro``).""" + import json + + def _canonical(payload: dict[str, Any]) -> str: + return json.dumps(payload, sort_keys=True, default=str) + + task_metrics_list = _sample_task_metrics_list() + for weighted in (True, False): + agg_payload = calculate_aggregate_metrics(task_metrics_list, weighted=weighted) + model = AggregateMetrics.model_validate(agg_payload) + dumped = model.model_dump(by_alias=True, mode="json", exclude_unset=True) + assert _canonical(agg_payload) == _canonical(dumped), ( + f"AggregateMetrics(weighted={weighted}) JSON drift between dict and " + f"model path.\nsource: {_canonical(agg_payload)}\nmodel: {_canonical(dumped)}" + ) diff --git a/tolokaforge/core/output/aggregate_models.py b/tolokaforge/core/output/aggregate_models.py index 619a1ca4..055943d5 100644 --- a/tolokaforge/core/output/aggregate_models.py +++ b/tolokaforge/core/output/aggregate_models.py @@ -30,18 +30,32 @@ 1. ``schema_version`` on :class:`RunAggregate` is **always emitted** on dump, regardless of ``exclude_unset``. See the model-serializer on the class for details. -2. Numeric fields carry ``int | float`` unions so the wire type - matches what the producer emitted (``sum([]) == 0`` from an empty - aggregate stays ``int`` on the wire; populated fields stay - ``float``). This preserves byte-identical JSON between the - dict-based writer path and any future model-based writer path. +2. **Type-preserving numeric fields.** Fields that can naturally be + ``int`` at the producer (token counts, latency percentiles, + stuck-rate counts, cost sums) carry ``int | float`` unions so the + wire type matches whatever the producer emitted. Pydantic v2's + smart-union picks the more specific type on validation, so a + source ``int 42`` round-trips to ``42`` (not ``42.0``) in the + JSON dump. Rate-shaped fields (``success_rate``, ``pass@k``, + ``avg_score`` and their macro/micro aggregates) stay narrow + ``float`` because they're always ``sum(...) / n`` divisions — + never ``int`` at the producer. + + The current producers in ``metrics.py`` and + ``failure_attribution.py`` happen to emit ``float`` for every + widened field today (both metric-calc functions short-circuit on + empty input rather than dividing by zero). The union is + defensive future-proofing: if a future refactor changes a + producer to emit ``int`` (e.g. a counting aggregator that returns + ``0`` for empty), the model preserves the wire type without a + coordinated model change. """ from __future__ import annotations from typing import Any -from pydantic import BaseModel, ConfigDict, Field, SerializationInfo, model_serializer +from pydantic import BaseModel, ConfigDict, Field, model_serializer __all__ = [ "AggregateMetrics", @@ -147,24 +161,26 @@ class AggregateMetrics(BaseModel): total_tasks: int total_trials: int - # Weighted (micro) averages — set when ``weighted=True``. Union of - # ``int | float | None`` so an empty-aggregate ``int 0`` from the - # producer preserves the JSON wire type verbatim (see - # :class:`PerTaskMetrics` for the same convention). - success_rate_micro: int | float | None = None - avg_score_micro: int | float | None = None + # Weighted (micro) averages — set when ``weighted=True``. Rate-shaped + # ([0, 1] scalar from ``sum(...) / n``): always ``float`` at the + # producer. Stays narrow — no ``int`` widening — matching the + # ``PerTaskMetrics.success_rate`` / ``avg_score`` siblings. + success_rate_micro: float | None = None + avg_score_micro: float | None = None - # Unweighted (macro) averages — set when ``weighted=False``. - success_rate_macro: int | float | None = None - avg_score_macro: int | float | None = None + # Unweighted (macro) averages — set when ``weighted=False``. Same + # rate-shaped rationale. + success_rate_macro: float | None = None + avg_score_macro: float | None = None # pass@k macro averages across tasks, plus their pass_hat aliases. - pass_at_1_macro: int | float | None = Field(default=None, alias="pass@1_macro") - pass_at_5_macro: int | float | None = Field(default=None, alias="pass@5_macro") - pass_at_10_macro: int | float | None = Field(default=None, alias="pass@10_macro") - pass_hat_at_1_macro: int | float | None = Field(default=None, alias="pass_hat@1_macro") - pass_hat_at_5_macro: int | float | None = Field(default=None, alias="pass_hat@5_macro") - pass_hat_at_10_macro: int | float | None = Field(default=None, alias="pass_hat@10_macro") + # Also always ``float`` (macro-average of ``pass@k`` rates). + pass_at_1_macro: float | None = Field(default=None, alias="pass@1_macro") + pass_at_5_macro: float | None = Field(default=None, alias="pass@5_macro") + pass_at_10_macro: float | None = Field(default=None, alias="pass@10_macro") + pass_hat_at_1_macro: float | None = Field(default=None, alias="pass_hat@1_macro") + pass_hat_at_5_macro: float | None = Field(default=None, alias="pass_hat@5_macro") + pass_hat_at_10_macro: float | None = Field(default=None, alias="pass_hat@10_macro") # Simple cross-task averages. avg_latency_s: int | float = 0 @@ -221,11 +237,16 @@ class RunAggregate(AggregateMetrics): schema_version: int = 1 @model_serializer(mode="wrap") - def _always_include_schema_version( - self, handler: Any, info: SerializationInfo - ) -> dict[str, Any]: + def _always_include_schema_version(self, handler: Any) -> dict[str, Any]: """Wrap the default serializer so ``schema_version`` is always - present in the dumped dict, even under ``exclude_unset=True``.""" + present in the dumped dict, even under ``exclude_unset=True``. + + The wrapped handler produces whatever dict the caller's dump + options normally would; this method injects ``schema_version`` + only when the handler's output omitted it (which happens under + ``exclude_unset=True`` when the field was left at its default). + No forced overwrite of a caller-set value. + """ data = handler(self) if isinstance(data, dict) and "schema_version" not in data: data["schema_version"] = self.schema_version From 9cd7ea437943c41f3b555d9af4328dabc2e44792 Mon Sep 17 00:00:00 2001 From: cirogam22 Date: Sun, 5 Jul 2026 15:44:16 +0200 Subject: [PATCH 3/3] review: hoist function-local json import to module level MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up on the review nit — the two `import json` statements inside `test_int_valued_numeric_fields_preserve_int_type` and `test_current_producer_output_matches_model_dump_byte_for_byte` are now a single top-of-module `import json` (the rest of the file already imports at module level). Purely cosmetic; no behaviour change. --- tests/canonical/test_run_aggregate_models_snapshot.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/canonical/test_run_aggregate_models_snapshot.py b/tests/canonical/test_run_aggregate_models_snapshot.py index 29d84ad6..20cbd635 100644 --- a/tests/canonical/test_run_aggregate_models_snapshot.py +++ b/tests/canonical/test_run_aggregate_models_snapshot.py @@ -23,6 +23,7 @@ from __future__ import annotations +import json from datetime import UTC, datetime from typing import Any @@ -404,8 +405,6 @@ def test_int_valued_numeric_fields_preserve_int_type() -> None: future producer refactor introducing ``int`` output — the model must handle it correctly without a coordinated widening. Closes #153. """ - import json - # Hand-construct a payload with int-valued token counts + latency # percentiles. The producer path today never emits these as int, # so this is the only way to exercise the int branch of the union. @@ -544,7 +543,6 @@ def test_current_producer_output_matches_model_dump_byte_for_byte() -> None: catches any accidental type coercion in the round trip for the fields we DID narrow (``success_rate_*``, ``avg_score_*``, ``pass_at_*_macro``).""" - import json def _canonical(payload: dict[str, Any]) -> str: return json.dumps(payload, sort_keys=True, default=str)