diff --git a/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/fabric/runtime.py b/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/fabric/runtime.py index 22b650efb5..381f42849d 100644 --- a/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/fabric/runtime.py +++ b/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/fabric/runtime.py @@ -55,13 +55,14 @@ from nemo_evaluator_sdk.agent_eval.tasks import AgentEvalRunConfig, AgentEvalTask from nemo_evaluator_sdk.agent_eval.trials import AgentEvalTrial, AgentEvalTrialStatus, AgentOutput, RunnerInfo from nemo_evaluator_sdk.agent_eval.workspace_seeds import SEED_FILES_INPUT_KEY, seed_workspace +from nemo_evaluator_sdk.values.atif import FinalMetrics from nemo_evaluator_sdk.values.evidence import ( EVIDENCE_FORMAT_ATIF, EVIDENCE_TRACE, CandidateEvidence, EvidenceDescriptor, ) -from pydantic import JsonValue +from pydantic import JsonValue, ValidationError if TYPE_CHECKING: # Annotations use nemo_fabric's real types (single source of truth). nemo_fabric is an optional @@ -358,9 +359,13 @@ async def _run_task( raise hook_extras = None except TimeoutError as exc: - return self._failed_trial(task, evidence_dir, exc, extra_metadata=self._skill_metadata(skill_provenances)) + return self._failed_trial( + task, evidence_dir, exc, extra_metadata=self._failed_metadata(skill_provenances, evidence_dir) + ) except Exception as exc: # noqa: BLE001 - a task failure must not abort the whole run - return self._failed_trial(task, evidence_dir, exc, extra_metadata=self._skill_metadata(skill_provenances)) + return self._failed_trial( + task, evidence_dir, exc, extra_metadata=self._failed_metadata(skill_provenances, evidence_dir) + ) finally: if self._task_hook is not None: try: @@ -396,6 +401,18 @@ def _skill_metadata(provenances: list[SkillProvenance]) -> dict[str, Any]: """ return {"skill": provenances[0] if len(provenances) == 1 else None, "skills": provenances} + @staticmethod + def _failed_metadata(provenances: list[SkillProvenance], evidence_dir: Path) -> dict[str, Any]: + """Trial metadata for a timed-out/errored task: skill provenance plus whatever tokens Relay flushed. + + Timeouts never reach ``_to_trial``, and there is no ``RunResult`` here, so the trajectory is read + straight from the relay dir — these are the long, expensive rows the token count matters most for. + """ + return { + **FabricAgentRuntime._skill_metadata(provenances), + **_atif_token_metadata(_relay_atif_path(evidence_dir)), + } + def _to_trial( self, task: AgentEvalTask, @@ -421,6 +438,9 @@ def _to_trial( # Skill provenance (name + content hash + injection mode) for the A/B diff. **self._skill_metadata(skill_provenances or []), **extras, + # Token usage from the Relay ATIF trajectory; Fabric's RunResult carries no usage of its + # own. Merged last so a hook extra can't shadow it. + **_atif_token_metadata(_atif_artifact_path(result)), } if result.status != "succeeded": @@ -715,6 +735,83 @@ def _result_error(result: RunResult) -> Mapping[str, Any]: return {"stage": error.stage, "code": error.code, "message": error.message} +def _atif_artifact_path(result: RunResult) -> Path | None: + """Path of the ATIF trajectory Fabric promoted as an artifact, if any.""" + for artifact in result.artifacts.artifacts: + if artifact.kind == _ATIF_ARTIFACT_KIND: + return Path(artifact.path) + return None + + +def _relay_atif_path(evidence_dir: Path) -> Path | None: + """Path of the Relay-written ATIF trajectory, used when no ``RunResult`` exists (timeout/error). + + Relay's filename template is per-session, so more than one file can land when subagents emit + their own sessions. Picking one under-reports and summing double-counts a root that already + aggregates, so anything other than a single match reports nothing rather than a wrong number. + """ + matches = sorted((evidence_dir / _RELAY_SUBDIR).glob(_ATIF_FILENAME_TEMPLATE.format(session_id="*"))) + if len(matches) == 1: + return matches[0] + if matches: + logger.warning("Fabric token capture: %d ATIF trajectories under %s; skipping", len(matches), evidence_dir) + return None + + +def _atif_token_metadata(path: Path | None) -> dict[str, int]: + """Project an ATIF trajectory's token totals onto the trial-metadata ``TOKEN_KEYS``. + + Each token field is resolved on its own: the trajectory-level ``final_metrics`` aggregate when it + reports that field, else the sum of the matching per-step ``metrics``. Every ``final_metrics`` + field is optional, so a block carrying only ``total_steps`` or a cost — or one that fails to + validate — must not suppress counts the steps do carry. That partial shape is likeliest on the + timeout path, where the trajectory was flushed mid-run and the counts matter most. + + ``total_tokens`` and ``cache_creation_tokens`` have no ATIF source and stay unset — Intake + recomputes the total. A missing or unreadable trajectory yields ``{}``: an absent token count + must not fail the trial. + """ + if path is None: + return {} + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + logger.warning("Fabric token capture: unreadable ATIF trajectory %s (%s)", path, exc) + return {} + if not isinstance(payload, Mapping): + return {} + + totals = FinalMetrics() + final_metrics = payload.get("final_metrics") + if isinstance(final_metrics, Mapping): + try: + totals = FinalMetrics.model_validate(final_metrics) + except ValidationError as exc: + logger.warning("Fabric token capture: invalid final_metrics in %s (%s)", path, exc) + + captured = { + "prompt_tokens": (totals.total_prompt_tokens, "prompt_tokens"), + "completion_tokens": (totals.total_completion_tokens, "completion_tokens"), + "cache_read_tokens": (totals.total_cached_tokens, "cached_tokens"), + } + resolved = { + key: total if total is not None else _sum_step_metric(payload, step_key) + for key, (total, step_key) in captured.items() + } + return {key: value for key, value in resolved.items() if value is not None} + + +def _sum_step_metric(payload: Mapping[str, Any], key: str) -> int | None: + """Sum one per-step ATIF metric across the trajectory, or ``None`` when no step reported it.""" + total: int | None = None + for step in payload.get("steps") or []: + metrics = step.get("metrics") if isinstance(step, Mapping) else None + value = metrics.get(key) if isinstance(metrics, Mapping) else None + if isinstance(value, int) and not isinstance(value, bool): + total = value if total is None else total + value + return total + + def _safe_path_name(value: str) -> str: return "".join(char if char.isalnum() or char in "._-" else "-" for char in value).strip(".-")[:120] diff --git a/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_inference.py b/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_inference.py index a9b95197f2..dcfa5b743b 100644 --- a/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_inference.py +++ b/packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_inference.py @@ -369,6 +369,12 @@ async def _invoke_post() -> dict[str, Any]: field_name=invocation.response_path_field, ) response = _openai_response(str(response_value)) + # The synthesized response keeps only the extracted text, so carry the agent's own token + # usage across: it is the sole token source for a row evaluation, whose trials are otherwise + # built with no measurements at all. Model targets already return their full completion. + usage = result_data.get("usage") if isinstance(result_data, Mapping) else None + if isinstance(usage, Mapping): + response["usage"] = dict(usage) if invocation.trajectory_path: trajectory = _extract_jsonpath( result_data, diff --git a/packages/nemo_evaluator_sdk/tests/agent_eval/test_fabric_runtime.py b/packages/nemo_evaluator_sdk/tests/agent_eval/test_fabric_runtime.py index d13b728628..55e190c6d0 100644 --- a/packages/nemo_evaluator_sdk/tests/agent_eval/test_fabric_runtime.py +++ b/packages/nemo_evaluator_sdk/tests/agent_eval/test_fabric_runtime.py @@ -1100,3 +1100,107 @@ def handler(agent: Any, kwargs: dict[str, Any]) -> _FakeResult: assert not (workspace / ".agents").exists() # Provenance is still stamped on the failed trial for the A/B diff. assert [prov["name"] for prov in trials[0].metadata["skills"]] == list(names) + + +# --- ATIF token capture ----------------------------------------------------- + + +def _atif(tmp_path: Path, payload: Mapping[str, Any]) -> Path: + path = tmp_path / "trajectory-abc.atif.json" + path.write_text(json.dumps(payload), encoding="utf-8") + return path + + +def _step(**metrics: int) -> dict[str, Any]: + return {"source": "agent", "message": "", "metrics": metrics} + + +def test_final_metrics_totals_are_projected_onto_the_token_keys(tmp_path: Path) -> None: + path = _atif( + tmp_path, + { + "schema_version": "ATIF-v1.7", + "steps": [_step(prompt_tokens=1)], + "final_metrics": { + "total_prompt_tokens": 358, + "total_completion_tokens": 19324, + "total_cached_tokens": 3984621, + }, + }, + ) + assert fabric_runtime._atif_token_metadata(path) == { + "prompt_tokens": 358, + "completion_tokens": 19324, + "cache_read_tokens": 3984621, + } + + +def test_steps_are_summed_when_the_trajectory_has_no_aggregate_block(tmp_path: Path) -> None: + path = _atif( + tmp_path, + { + "schema_version": "ATIF-v1.7", + "steps": [_step(prompt_tokens=100, completion_tokens=10), _step(prompt_tokens=58, cached_tokens=7)], + }, + ) + assert fabric_runtime._atif_token_metadata(path) == { + "prompt_tokens": 158, + "completion_tokens": 10, + "cache_read_tokens": 7, + } + + +def test_a_partial_aggregate_does_not_suppress_the_fields_only_the_steps_report(tmp_path: Path) -> None: + # Every final_metrics field is optional, so a block reporting only steps/cost validates cleanly. + # Resolving per field is what keeps a mid-run flush (the timeout path) from publishing nothing. + path = _atif( + tmp_path, + { + "schema_version": "ATIF-v1.7", + "steps": [_step(prompt_tokens=100, completion_tokens=10), _step(prompt_tokens=58)], + "final_metrics": {"total_steps": 2, "total_cost_usd": 0.12}, + }, + ) + assert fabric_runtime._atif_token_metadata(path) == {"prompt_tokens": 158, "completion_tokens": 10} + + +def test_a_reported_total_wins_over_the_step_sum_for_that_field_alone(tmp_path: Path) -> None: + # The aggregate is authoritative where it speaks; the steps fill only the fields it omits. + path = _atif( + tmp_path, + { + "schema_version": "ATIF-v1.7", + "steps": [_step(prompt_tokens=1, completion_tokens=10)], + "final_metrics": {"total_prompt_tokens": 358}, + }, + ) + assert fabric_runtime._atif_token_metadata(path) == {"prompt_tokens": 358, "completion_tokens": 10} + + +def test_an_unvalidatable_aggregate_still_falls_back_to_the_steps(tmp_path: Path) -> None: + path = _atif( + tmp_path, + { + "schema_version": "ATIF-v1.7", + "steps": [_step(prompt_tokens=158)], + "final_metrics": {"total_prompt_tokens": "not-an-int"}, + }, + ) + assert fabric_runtime._atif_token_metadata(path) == {"prompt_tokens": 158} + + +def test_a_trajectory_with_no_counts_anywhere_records_nothing(tmp_path: Path) -> None: + path = _atif(tmp_path, {"schema_version": "ATIF-v1.7", "steps": [_step()], "final_metrics": {}}) + assert fabric_runtime._atif_token_metadata(path) == {} + + +@pytest.mark.parametrize("payload", ["[]", "{ not json", '"a string"']) +def test_an_unreadable_trajectory_records_nothing_rather_than_failing(tmp_path: Path, payload: str) -> None: + path = tmp_path / "trajectory-abc.atif.json" + path.write_text(payload, encoding="utf-8") + assert fabric_runtime._atif_token_metadata(path) == {} + + +def test_a_missing_trajectory_records_nothing(tmp_path: Path) -> None: + assert fabric_runtime._atif_token_metadata(None) == {} + assert fabric_runtime._atif_token_metadata(tmp_path / "absent.atif.json") == {} diff --git a/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py b/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py index 424b07d6e8..a3906fd55b 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py @@ -18,7 +18,10 @@ import hashlib import json +import logging +from collections.abc import Mapping from datetime import datetime +from typing import Any from nemo_evaluator_sdk.agent_eval.results import AgentEvalResult, AgentEvalSummary, RunMetadata from nemo_evaluator_sdk.agent_eval.scores import ( @@ -32,6 +35,8 @@ from nemo_evaluator_sdk.values.multi_metric_results import BenchmarkEvaluationResult from nemo_evaluator_sdk.values.results import EvaluationResult, RowScore +logger = logging.getLogger(__name__) + #: Key ``sample`` carries when generation itself failed, rather than the metric. _INFERENCE_ERROR = "inference_error" @@ -80,6 +85,61 @@ def _output(row: RowScore) -> AgentOutput | None: return AgentOutput(output_text=output_text, response=response) +def _first_int(usage: Mapping[str, Any], *keys: str) -> int | None: + """First key in ``keys`` holding a token count, or ``None``. + + A count is a non-negative int that is not a bool. Negatives are rejected because they are used + as an unknown-value sentinel rather than a measurement, and nothing downstream would catch one: + Intake's ``total_prompt_tokens`` is an unconstrained ``int | None``, so a negative would be + summed into the evaluation rollup and shown as a real total. + """ + for key in keys: + value = usage.get(key) + if isinstance(value, int) and not isinstance(value, bool) and value >= 0: + return value + return None + + +def _token_metadata(row: RowScore) -> dict[str, int]: + """Project the generation response's ``usage`` block onto the trial-metadata token keys. + + A row's generation response is the only place its token usage survives: ``row.requests`` mixes + the generation call with each metric's judge calls, so summing that would credit judge tokens to + the agent. ``total_tokens`` is left unset — Intake recomputes it from the parts. + + Both usage schemas a target can return are read, OpenAI's first: a ``GenericAgent`` points at an + arbitrary URL, so the response is whatever that endpoint emits, and an Anthropic-shaped block + would otherwise be dropped whole. Note the two disagree on whether cache reads are already + counted in the prompt total (OpenAI includes them, Anthropic does not); the values are recorded + as reported rather than reconciled, since nothing downstream adds them together. + + No key list covers every provider, and a renamed key would fail the same silent way this + function exists to fix, so a usage block that yields nothing is logged with the keys it actually + carried. An unrecognized schema is then a log line naming what to add, not an unexplained blank. + """ + response = row.sample.get("response") + usage = response.get("usage") if isinstance(response, Mapping) else None + if not isinstance(usage, Mapping): + return {} + details = usage.get("prompt_tokens_details") + cache_read = _first_int(details, "cached_tokens") if isinstance(details, Mapping) else None + if cache_read is None: + cache_read = _first_int(usage, "cache_read_input_tokens") + captured = { + "prompt_tokens": _first_int(usage, "prompt_tokens", "input_tokens"), + "completion_tokens": _first_int(usage, "completion_tokens", "output_tokens"), + "cache_read_tokens": cache_read, + "cache_creation_tokens": _first_int(usage, "cache_creation_input_tokens"), + } + recorded = {key: value for key, value in captured.items() if value is not None} + if not recorded: + logger.warning( + "No token counts recognized in the generation response's usage block; keys present: %s", + sorted(str(key) for key in usage), + ) + return recorded + + def _scores(row: RowScore, *, run_id: str, task_id: str, trial_id: str) -> list[AgentEvalTaskScore]: """One score per metric key on the row; ``metrics`` values are already ``MetricOutput``.""" errors = row.metric_errors or {} @@ -152,6 +212,7 @@ def row_result_to_agent_eval_result( task_id=task_id, status=AgentEvalTrialStatus.COMPLETED if output is not None else AgentEvalTrialStatus.FAILED, output=output, + metadata=_token_metadata(row), ) ) scores.extend(_scores(row, run_id=run_id, task_id=task_id, trial_id=trial_id)) diff --git a/plugins/nemo-evaluator/tests/intake/test_row_adapter.py b/plugins/nemo-evaluator/tests/intake/test_row_adapter.py index c190c81ef5..8718bb42c8 100644 --- a/plugins/nemo-evaluator/tests/intake/test_row_adapter.py +++ b/plugins/nemo-evaluator/tests/intake/test_row_adapter.py @@ -5,6 +5,7 @@ from __future__ import annotations +import logging from datetime import UTC, datetime from typing import Any @@ -218,3 +219,94 @@ def test_a_metric_error_does_not_fail_the_trial_itself() -> None: # The agent answered; only scoring failed. The trajectory is still worth publishing. result = _adapt([_row(metric_errors={"exact_match": "judge timed out"})]) assert result.trials[0].status == AgentEvalTrialStatus.COMPLETED + + +# --- token usage ------------------------------------------------------------ + + +def _usage_metadata(usage: object) -> dict[str, Any]: + return _adapt([_row(sample={"output_text": "4", "response": {"usage": usage}})]).trials[0].metadata + + +def test_openai_usage_becomes_trial_token_measurements() -> None: + # Publishing reads these keys off trial metadata, so the names are the contract with + # TrialMeasurements.from_metadata; total_tokens is deliberately absent (Intake recomputes it). + metadata = _usage_metadata( + { + "prompt_tokens": 22635, + "completion_tokens": 2949, + "total_tokens": 25584, + "prompt_tokens_details": {"cached_tokens": 1200}, + } + ) + assert metadata == {"prompt_tokens": 22635, "completion_tokens": 2949, "cache_read_tokens": 1200} + + +def test_anthropic_usage_is_read_under_its_own_key_names() -> None: + # A GenericAgent can target any endpoint, so an Anthropic-shaped block must not be dropped whole. + metadata = _usage_metadata( + { + "input_tokens": 358, + "output_tokens": 19324, + "cache_read_input_tokens": 3984621, + "cache_creation_input_tokens": 512, + } + ) + assert metadata == { + "prompt_tokens": 358, + "completion_tokens": 19324, + "cache_read_tokens": 3984621, + "cache_creation_tokens": 512, + } + + +def test_openai_keys_win_when_a_response_carries_both_schemas() -> None: + metadata = _usage_metadata({"prompt_tokens": 10, "input_tokens": 999, "completion_tokens": 20}) + assert metadata["prompt_tokens"] == 10 + + +@pytest.mark.parametrize("usage", [None, {}, "1000", {"prompt_tokens": None}, {"prompt_tokens": "22635"}]) +def test_unusable_usage_records_no_tokens(usage: object) -> None: + # A missing count must stay missing rather than land as a wrong number. + assert _usage_metadata(usage) == {} + + +def test_zero_counts_are_recorded_rather_than_dropped() -> None: + # NAT reports prompt_tokens=0; 0 is a reported measurement and must survive a truthiness filter. + assert _usage_metadata({"prompt_tokens": 0, "completion_tokens": 38}) == { + "prompt_tokens": 0, + "completion_tokens": 38, + } + + +def test_booleans_are_not_counted_as_token_counts() -> None: + assert _usage_metadata({"prompt_tokens": True, "completion_tokens": 38}) == {"completion_tokens": 38} + + +def test_negative_counts_are_rejected_rather_than_summed_into_a_total() -> None: + # -1 is an unknown-value sentinel, not a measurement, and Intake's token fields carry no ge=0 + # constraint — so publishing one would deflate the evaluation rollup with no error anywhere. + assert _usage_metadata({"prompt_tokens": -1, "completion_tokens": 38}) == {"completion_tokens": 38} + + +def test_a_negative_falls_through_to_the_next_known_key() -> None: + assert _usage_metadata({"prompt_tokens": -1, "input_tokens": 500})["prompt_tokens"] == 500 + + +def test_a_row_without_a_response_records_no_tokens() -> None: + assert _adapt([_row(sample={"output_text": "4"})]).trials[0].metadata == {} + + +def test_an_unrecognized_usage_schema_is_logged_with_its_keys(caplog: pytest.LogCaptureFixture) -> None: + # No key list covers every provider, so the one thing that must not happen is a silent blank: + # the log has to name the real keys, which is what tells us what to add. + with caplog.at_level(logging.WARNING, logger="nemo_evaluator.intake.row_adapter"): + assert _usage_metadata({"promptTokenCount": 12, "candidatesTokenCount": 34}) == {} + assert "promptTokenCount" in caplog.text + assert "candidatesTokenCount" in caplog.text + + +def test_a_recognized_usage_block_logs_nothing(caplog: pytest.LogCaptureFixture) -> None: + with caplog.at_level(logging.WARNING, logger="nemo_evaluator.intake.row_adapter"): + _usage_metadata({"prompt_tokens": 1, "completion_tokens": 2}) + assert caplog.text == "" diff --git a/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/runtimes/fabric/runtime.py b/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/runtimes/fabric/runtime.py index 712128bd4e..8e0bf51275 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/runtimes/fabric/runtime.py +++ b/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_eval/runtimes/fabric/runtime.py @@ -55,13 +55,14 @@ from nemo_platform.beta.evaluator.agent_eval.tasks import AgentEvalRunConfig, AgentEvalTask from nemo_platform.beta.evaluator.agent_eval.trials import AgentEvalTrial, AgentEvalTrialStatus, AgentOutput, RunnerInfo from nemo_platform.beta.evaluator.agent_eval.workspace_seeds import SEED_FILES_INPUT_KEY, seed_workspace +from nemo_platform.beta.evaluator.values.atif import FinalMetrics from nemo_platform.beta.evaluator.values.evidence import ( EVIDENCE_FORMAT_ATIF, EVIDENCE_TRACE, CandidateEvidence, EvidenceDescriptor, ) -from pydantic import JsonValue +from pydantic import JsonValue, ValidationError if TYPE_CHECKING: # Annotations use nemo_fabric's real types (single source of truth). nemo_fabric is an optional @@ -358,9 +359,13 @@ async def _run_task( raise hook_extras = None except TimeoutError as exc: - return self._failed_trial(task, evidence_dir, exc, extra_metadata=self._skill_metadata(skill_provenances)) + return self._failed_trial( + task, evidence_dir, exc, extra_metadata=self._failed_metadata(skill_provenances, evidence_dir) + ) except Exception as exc: # noqa: BLE001 - a task failure must not abort the whole run - return self._failed_trial(task, evidence_dir, exc, extra_metadata=self._skill_metadata(skill_provenances)) + return self._failed_trial( + task, evidence_dir, exc, extra_metadata=self._failed_metadata(skill_provenances, evidence_dir) + ) finally: if self._task_hook is not None: try: @@ -396,6 +401,18 @@ def _skill_metadata(provenances: list[SkillProvenance]) -> dict[str, Any]: """ return {"skill": provenances[0] if len(provenances) == 1 else None, "skills": provenances} + @staticmethod + def _failed_metadata(provenances: list[SkillProvenance], evidence_dir: Path) -> dict[str, Any]: + """Trial metadata for a timed-out/errored task: skill provenance plus whatever tokens Relay flushed. + + Timeouts never reach ``_to_trial``, and there is no ``RunResult`` here, so the trajectory is read + straight from the relay dir — these are the long, expensive rows the token count matters most for. + """ + return { + **FabricAgentRuntime._skill_metadata(provenances), + **_atif_token_metadata(_relay_atif_path(evidence_dir)), + } + def _to_trial( self, task: AgentEvalTask, @@ -421,6 +438,9 @@ def _to_trial( # Skill provenance (name + content hash + injection mode) for the A/B diff. **self._skill_metadata(skill_provenances or []), **extras, + # Token usage from the Relay ATIF trajectory; Fabric's RunResult carries no usage of its + # own. Merged last so a hook extra can't shadow it. + **_atif_token_metadata(_atif_artifact_path(result)), } if result.status != "succeeded": @@ -715,6 +735,83 @@ def _result_error(result: RunResult) -> Mapping[str, Any]: return {"stage": error.stage, "code": error.code, "message": error.message} +def _atif_artifact_path(result: RunResult) -> Path | None: + """Path of the ATIF trajectory Fabric promoted as an artifact, if any.""" + for artifact in result.artifacts.artifacts: + if artifact.kind == _ATIF_ARTIFACT_KIND: + return Path(artifact.path) + return None + + +def _relay_atif_path(evidence_dir: Path) -> Path | None: + """Path of the Relay-written ATIF trajectory, used when no ``RunResult`` exists (timeout/error). + + Relay's filename template is per-session, so more than one file can land when subagents emit + their own sessions. Picking one under-reports and summing double-counts a root that already + aggregates, so anything other than a single match reports nothing rather than a wrong number. + """ + matches = sorted((evidence_dir / _RELAY_SUBDIR).glob(_ATIF_FILENAME_TEMPLATE.format(session_id="*"))) + if len(matches) == 1: + return matches[0] + if matches: + logger.warning("Fabric token capture: %d ATIF trajectories under %s; skipping", len(matches), evidence_dir) + return None + + +def _atif_token_metadata(path: Path | None) -> dict[str, int]: + """Project an ATIF trajectory's token totals onto the trial-metadata ``TOKEN_KEYS``. + + Each token field is resolved on its own: the trajectory-level ``final_metrics`` aggregate when it + reports that field, else the sum of the matching per-step ``metrics``. Every ``final_metrics`` + field is optional, so a block carrying only ``total_steps`` or a cost — or one that fails to + validate — must not suppress counts the steps do carry. That partial shape is likeliest on the + timeout path, where the trajectory was flushed mid-run and the counts matter most. + + ``total_tokens`` and ``cache_creation_tokens`` have no ATIF source and stay unset — Intake + recomputes the total. A missing or unreadable trajectory yields ``{}``: an absent token count + must not fail the trial. + """ + if path is None: + return {} + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + logger.warning("Fabric token capture: unreadable ATIF trajectory %s (%s)", path, exc) + return {} + if not isinstance(payload, Mapping): + return {} + + totals = FinalMetrics() + final_metrics = payload.get("final_metrics") + if isinstance(final_metrics, Mapping): + try: + totals = FinalMetrics.model_validate(final_metrics) + except ValidationError as exc: + logger.warning("Fabric token capture: invalid final_metrics in %s (%s)", path, exc) + + captured = { + "prompt_tokens": (totals.total_prompt_tokens, "prompt_tokens"), + "completion_tokens": (totals.total_completion_tokens, "completion_tokens"), + "cache_read_tokens": (totals.total_cached_tokens, "cached_tokens"), + } + resolved = { + key: total if total is not None else _sum_step_metric(payload, step_key) + for key, (total, step_key) in captured.items() + } + return {key: value for key, value in resolved.items() if value is not None} + + +def _sum_step_metric(payload: Mapping[str, Any], key: str) -> int | None: + """Sum one per-step ATIF metric across the trajectory, or ``None`` when no step reported it.""" + total: int | None = None + for step in payload.get("steps") or []: + metrics = step.get("metrics") if isinstance(step, Mapping) else None + value = metrics.get(key) if isinstance(metrics, Mapping) else None + if isinstance(value, int) and not isinstance(value, bool): + total = value if total is None else total + value + return total + + def _safe_path_name(value: str) -> str: return "".join(char if char.isalnum() or char in "._-" else "-" for char in value).strip(".-")[:120] diff --git a/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_inference.py b/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_inference.py index 181b29bbdf..a621d40c81 100644 --- a/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_inference.py +++ b/sdk/python/nemo-platform/src/nemo_platform/beta/evaluator/agent_inference.py @@ -369,6 +369,12 @@ async def _invoke_post() -> dict[str, Any]: field_name=invocation.response_path_field, ) response = _openai_response(str(response_value)) + # The synthesized response keeps only the extracted text, so carry the agent's own token + # usage across: it is the sole token source for a row evaluation, whose trials are otherwise + # built with no measurements at all. Model targets already return their full completion. + usage = result_data.get("usage") if isinstance(result_data, Mapping) else None + if isinstance(usage, Mapping): + response["usage"] = dict(usage) if invocation.trajectory_path: trajectory = _extract_jsonpath( result_data, diff --git a/web/packages/studio/src/components/dataViews/EvaluationSessionsDataView/index.tsx b/web/packages/studio/src/components/dataViews/EvaluationSessionsDataView/index.tsx index 87eba33c3a..60c382adc4 100644 --- a/web/packages/studio/src/components/dataViews/EvaluationSessionsDataView/index.tsx +++ b/web/packages/studio/src/components/dataViews/EvaluationSessionsDataView/index.tsx @@ -29,6 +29,7 @@ import { IntakePayloadPreviewCell } from '@studio/components/IntakeLists/IntakeP import { useWorkspaceFromPath } from '@studio/hooks/useWorkspaceFromPath'; import { getEvaluationSessionTraceDetailRoute } from '@studio/routes/utils'; import { tooltipClassName } from '@studio/styles/common'; +import { formatInteger } from '@studio/util/intakeTelemetry'; import { keepPreviousData } from '@tanstack/react-query'; import { isAxiosError } from 'axios'; import { Columns3 } from 'lucide-react'; @@ -265,7 +266,7 @@ export const EvaluationSessionsDataView: FC = ( cell: ({ row }) => { const { input_tokens, output_tokens } = row.original; if (input_tokens == null && output_tokens == null) return -; - return {String((input_tokens ?? 0) + (output_tokens ?? 0))}; + return {formatInteger((input_tokens ?? 0) + (output_tokens ?? 0))}; }, } ), diff --git a/web/packages/studio/src/components/dataViews/ExperimentDataView/index.tsx b/web/packages/studio/src/components/dataViews/ExperimentDataView/index.tsx index 6392616b7c..5ba4248dcb 100644 --- a/web/packages/studio/src/components/dataViews/ExperimentDataView/index.tsx +++ b/web/packages/studio/src/components/dataViews/ExperimentDataView/index.tsx @@ -54,6 +54,7 @@ const STATIC_SORT_FIELD_MAP: Readonly> = { latency_ms: 'latency_ms.mean', total_latency_ms: 'latency_ms.sum', tokens: 'tokens.mean', + total_tokens: 'tokens.sum', test_case_count: 'test_case_count', }; @@ -67,6 +68,7 @@ const sortFieldToColumnId = (field: string): string | undefined => { if (field.startsWith('cost_usd.')) return 'cost_usd'; if (field === 'latency_ms.sum') return 'total_latency_ms'; if (field.startsWith('latency_ms.')) return 'latency_ms'; + if (field === 'tokens.sum') return 'total_tokens'; if (field.startsWith('tokens.')) return 'tokens'; const evaluatorMatch = field.match(/^evaluators\.(.+)\.[^.]+$/); if (evaluatorMatch) return `evaluator-${evaluatorMatch[1]}`; @@ -99,6 +101,7 @@ const getEvaluationFilterField = (id: string): string | undefined => { if (id === 'latency_ms') return 'latency_ms.mean'; if (id === 'total_latency_ms') return 'latency_ms.sum'; if (id === 'tokens') return 'tokens.mean'; + if (id === 'total_tokens') return 'tokens.sum'; // The list columns hold the plural name facets; the API filter params are singular contains-matches. if (id === 'agent_names') return 'agent_name'; if (id === 'agent_versions') return 'agent_version'; @@ -475,6 +478,18 @@ export const ExperimentDataView: FC = ({ group, paretoV ); }, }), + accessor((original) => original.tokens?.sum, { + id: 'total_tokens', + // Sum of per-test-case means: the rollup is deliberately k-invariant, so this is one pass + // through the dataset, not total tokens burned across repeated attempts. + header: 'Total tokens (per test case)', + enableSorting: true, + meta: { title: false, filter: numberRangeFilter('Total tokens (per test case)') }, + cell: ({ row }) => { + const { tokens } = row.original; + return {tokens?.sum != null ? Math.round(tokens.sum).toLocaleString() : '-'}; + }, + }), accessor((original) => original.test_case_count, { id: 'test_case_count', header: 'Total test cases', diff --git a/web/packages/studio/src/routes/EvaluationDetailRoute/EvaluationDetailMetrics.tsx b/web/packages/studio/src/routes/EvaluationDetailRoute/EvaluationDetailMetrics.tsx index 11f718b7b1..82e901a21c 100644 --- a/web/packages/studio/src/routes/EvaluationDetailRoute/EvaluationDetailMetrics.tsx +++ b/web/packages/studio/src/routes/EvaluationDetailRoute/EvaluationDetailMetrics.tsx @@ -25,6 +25,15 @@ export const EvaluationDetailMetrics: FC = ({ eval // formatDurationMs returns '—' for null/undefined, which is also KVPair's default empty value. const avgLatency = formatDurationMs(experiment?.latency_ms?.mean); + const tokenSum = experiment?.tokens?.sum; + const totalTokens = + tokenSum != null + ? Math.round(tokenSum).toLocaleString(undefined, { + notation: 'compact', + maximumFractionDigits: 0, + }) + : undefined; + const modelNames = experiment?.model_names ?? []; const modelNamesJoined = modelNames.length > 0 ? modelNames.join(', ') : undefined; const modelNamesValue: ReactNode = modelNamesJoined ? ( @@ -96,6 +105,8 @@ export const EvaluationDetailMetrics: FC = ({ eval + +