From fd7a1dc56e419017d40b70961021391ff9d3dc37 Mon Sep 17 00:00:00 2001 From: Octavian Drulea Date: Fri, 14 Aug 2026 22:24:11 -0700 Subject: [PATCH 1/7] feat(evaluator): surface token counts for agent evaluations Row evaluations invoke the agent over chat/completions, which reports its token usage, but the count was dropped twice before publish: the agent inference path synthesized a bare {"choices": [...]} response holding only the extracted text, and the row adapter built trials with no metadata at all. Publish therefore sent no ATIF final_metrics and Intake stored null, so every Tokens cell in Studio rendered empty. Carry the response's usage block onto the synthesized response, and project it onto the trial's token measurements in the row adapter. Cached tokens come from prompt_tokens_details; total_tokens is left unset because Intake recomputes it from the parts. The row's requests log is deliberately not used as the source: it concatenates the generation call with each metric's judge calls, so summing it would credit judge tokens to the agent. The Fabric agent-eval runtime had the same gap on its own path, where the numbers live in the Relay ATIF trajectory it writes and never reads back. It now reads them, preferring the artifact Fabric promoted as the trajectory and falling back to the relay directory on the timeout path, where no RunResult exists. A missing or unparseable trajectory yields no tokens rather than failing the trial. Studio gains a total-tokens column on the evaluation list, labelled per test case because the rollup sums per-test-case means and is deliberately k-invariant, and the per-row Tokens cell is now localized. Signed-off-by: Octavian Drulea --- .../agent_eval/runtimes/fabric/runtime.py | 99 ++++++++++++++++++- .../src/nemo_evaluator_sdk/agent_inference.py | 6 ++ .../src/nemo_evaluator/intake/row_adapter.py | 22 +++++ .../EvaluationSessionsDataView/index.tsx | 3 +- .../dataViews/ExperimentDataView/index.tsx | 15 +++ 5 files changed, 141 insertions(+), 4 deletions(-) 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..fab2d142f8 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,79 @@ 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``. + + Prefers the trajectory-level ``final_metrics``, falling back to summing per-step ``metrics`` (a + trajectory flushed mid-run may carry steps but no aggregate block). ``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 {} + + 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) + return {} + captured = { + "prompt_tokens": totals.total_prompt_tokens, + "completion_tokens": totals.total_completion_tokens, + "cache_read_tokens": totals.total_cached_tokens, + } + else: + captured = { + "prompt_tokens": _sum_step_metric(payload, "prompt_tokens"), + "completion_tokens": _sum_step_metric(payload, "completion_tokens"), + "cache_read_tokens": _sum_step_metric(payload, "cached_tokens"), + } + return {key: value for key, value in captured.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/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py b/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py index 424b07d6e8..b95737b8fc 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py @@ -18,6 +18,7 @@ import hashlib import json +from collections.abc import Mapping from datetime import datetime from nemo_evaluator_sdk.agent_eval.results import AgentEvalResult, AgentEvalSummary, RunMetadata @@ -80,6 +81,26 @@ def _output(row: RowScore) -> AgentOutput | None: return AgentOutput(output_text=output_text, response=response) +def _token_metadata(row: RowScore) -> dict[str, int]: + """Project the generation response's OpenAI ``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. + """ + 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") + captured = { + "prompt_tokens": usage.get("prompt_tokens"), + "completion_tokens": usage.get("completion_tokens"), + "cache_read_tokens": details.get("cached_tokens") if isinstance(details, Mapping) else None, + } + return {key: value for key, value in captured.items() if isinstance(value, int) and not isinstance(value, bool)} + + 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 +173,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/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', From 34f15926cf0c2641b794b03ea49fc2870230f86d Mon Sep 17 00:00:00 2001 From: Octavian Drulea Date: Fri, 14 Aug 2026 23:38:54 -0700 Subject: [PATCH 2/7] fix(studio): revert NAT readd Signed-off-by: Octavian Drulea --- .../EvaluationDetailRoute/EvaluationDetailMetrics.tsx | 11 +++++++++++ 1 file changed, 11 insertions(+) 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 + + From f2023cfe14d88d82f49eb9073bd5e03df73add67 Mon Sep 17 00:00:00 2001 From: Octavian Drulea Date: Sun, 16 Aug 2026 06:14:32 -0700 Subject: [PATCH 3/7] chore(sdk): re-vendor the evaluator SDK after the token-capture change `make vendor` copies nemo_evaluator_sdk into the SDK tree, so edits to agent_inference.py and the Fabric agent-eval runtime leave their vendored counterparts stale until it is re-run. lint-sdk-vendored catches exactly that, and lint-cli then fails as a cascade: the first script stages sdk/python/ before the second diffs it against the index, so one stale vendor shows up as two failing lints. Regenerated; no source change. Signed-off-by: Octavian Drulea --- .../agent_eval/runtimes/fabric/runtime.py | 99 ++++++++++++++++++- .../beta/evaluator/agent_inference.py | 6 ++ 2 files changed, 102 insertions(+), 3 deletions(-) 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..23d8610c33 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,79 @@ 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``. + + Prefers the trajectory-level ``final_metrics``, falling back to summing per-step ``metrics`` (a + trajectory flushed mid-run may carry steps but no aggregate block). ``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 {} + + 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) + return {} + captured = { + "prompt_tokens": totals.total_prompt_tokens, + "completion_tokens": totals.total_completion_tokens, + "cache_read_tokens": totals.total_cached_tokens, + } + else: + captured = { + "prompt_tokens": _sum_step_metric(payload, "prompt_tokens"), + "completion_tokens": _sum_step_metric(payload, "completion_tokens"), + "cache_read_tokens": _sum_step_metric(payload, "cached_tokens"), + } + return {key: value for key, value in captured.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, From 3413f8090f512f6111a70834c0183219bc1d3083 Mon Sep 17 00:00:00 2001 From: Octavian Drulea Date: Sun, 16 Aug 2026 06:18:45 -0700 Subject: [PATCH 4/7] fix(evaluator): read Anthropic-shaped token usage from row evaluations A GenericAgent target is an arbitrary URL, so its response is whatever that endpoint emits. Only OpenAI's usage key names were read, which meant an Anthropic-shaped block was dropped whole and the row published no token counts at all -- the same silent-null failure the OpenAI path was just fixed for. Read either schema, preferring OpenAI's names so an already-normalized response is unaffected, and pick up cache creation while there since Anthropic reports it. The two schemas disagree on whether cache reads are already counted in the prompt total; the values are recorded as reported rather than reconciled, because nothing downstream adds them together. Signed-off-by: Octavian Drulea --- .../src/nemo_evaluator/intake/row_adapter.py | 30 +++++++-- .../tests/intake/test_row_adapter.py | 66 +++++++++++++++++++ 2 files changed, 91 insertions(+), 5 deletions(-) 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 b95737b8fc..d3666683e1 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py @@ -20,6 +20,7 @@ import json 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 ( @@ -81,24 +82,43 @@ 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 real int, or ``None``. Bools are not counts.""" + for key in keys: + value = usage.get(key) + if isinstance(value, int) and not isinstance(value, bool): + return value + return None + + def _token_metadata(row: RowScore) -> dict[str, int]: - """Project the generation response's OpenAI ``usage`` block onto the trial-metadata token keys. + """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. """ 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": usage.get("prompt_tokens"), - "completion_tokens": usage.get("completion_tokens"), - "cache_read_tokens": details.get("cached_tokens") if isinstance(details, Mapping) else None, + "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"), } - return {key: value for key, value in captured.items() if isinstance(value, int) and not isinstance(value, bool)} + return {key: value for key, value in captured.items() if value is not None} def _scores(row: RowScore, *, run_id: str, task_id: str, trial_id: str) -> list[AgentEvalTaskScore]: diff --git a/plugins/nemo-evaluator/tests/intake/test_row_adapter.py b/plugins/nemo-evaluator/tests/intake/test_row_adapter.py index c190c81ef5..c0560b8b88 100644 --- a/plugins/nemo-evaluator/tests/intake/test_row_adapter.py +++ b/plugins/nemo-evaluator/tests/intake/test_row_adapter.py @@ -218,3 +218,69 @@ 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_a_row_without_a_response_records_no_tokens() -> None: + assert _adapt([_row(sample={"output_text": "4"})]).trials[0].metadata == {} From 33434bd47aa623e867268ec2eea02517294467fb Mon Sep 17 00:00:00 2001 From: Octavian Drulea Date: Sun, 16 Aug 2026 06:45:47 -0700 Subject: [PATCH 5/7] feat(evaluator): emit log when encountering unknown prompt token schema in response Signed-off-by: Octavian Drulea --- .../src/nemo_evaluator/intake/row_adapter.py | 15 ++++++++++++++- .../tests/intake/test_row_adapter.py | 16 ++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) 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 d3666683e1..40ab21af74 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py @@ -18,6 +18,7 @@ import hashlib import json +import logging from collections.abc import Mapping from datetime import datetime from typing import Any @@ -34,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" @@ -103,6 +106,10 @@ def _token_metadata(row: RowScore) -> dict[str, int]: 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 @@ -118,7 +125,13 @@ def _token_metadata(row: RowScore) -> dict[str, int]: "cache_read_tokens": cache_read, "cache_creation_tokens": _first_int(usage, "cache_creation_input_tokens"), } - return {key: value for key, value in captured.items() if value is not None} + 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]: diff --git a/plugins/nemo-evaluator/tests/intake/test_row_adapter.py b/plugins/nemo-evaluator/tests/intake/test_row_adapter.py index c0560b8b88..101bb797f9 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 @@ -284,3 +285,18 @@ def test_booleans_are_not_counted_as_token_counts() -> None: 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 == "" From be78da0cbd886458eed77a7b0c5cf700ec9d8393 Mon Sep 17 00:00:00 2001 From: Octavian Drulea Date: Sun, 16 Aug 2026 07:21:53 -0700 Subject: [PATCH 6/7] fix(evaluator): address PR review comments Signed-off-by: Octavian Drulea --- .../src/nemo_evaluator/intake/row_adapter.py | 10 ++++++++-- .../nemo-evaluator/tests/intake/test_row_adapter.py | 10 ++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) 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 40ab21af74..a3906fd55b 100644 --- a/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py +++ b/plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py @@ -86,10 +86,16 @@ def _output(row: RowScore) -> AgentOutput | None: def _first_int(usage: Mapping[str, Any], *keys: str) -> int | None: - """First key in ``keys`` holding a real int, or ``None``. Bools are not counts.""" + """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): + if isinstance(value, int) and not isinstance(value, bool) and value >= 0: return value return None diff --git a/plugins/nemo-evaluator/tests/intake/test_row_adapter.py b/plugins/nemo-evaluator/tests/intake/test_row_adapter.py index 101bb797f9..8718bb42c8 100644 --- a/plugins/nemo-evaluator/tests/intake/test_row_adapter.py +++ b/plugins/nemo-evaluator/tests/intake/test_row_adapter.py @@ -283,6 +283,16 @@ 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 == {} From 02cdcc71dbeca7fed435bcc3a85b4ff6bcb33646 Mon Sep 17 00:00:00 2001 From: Octavian Drulea Date: Mon, 17 Aug 2026 06:32:30 -0700 Subject: [PATCH 7/7] fix(evaluator): resolve each ATIF token field independently of the aggregate Signed-off-by: Octavian Drulea --- .../agent_eval/runtimes/fabric/runtime.py | 38 ++++--- .../tests/agent_eval/test_fabric_runtime.py | 104 ++++++++++++++++++ .../agent_eval/runtimes/fabric/runtime.py | 38 ++++--- 3 files changed, 146 insertions(+), 34 deletions(-) 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 fab2d142f8..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 @@ -761,10 +761,15 @@ def _relay_atif_path(evidence_dir: Path) -> Path | None: def _atif_token_metadata(path: Path | None) -> dict[str, int]: """Project an ATIF trajectory's token totals onto the trial-metadata ``TOKEN_KEYS``. - Prefers the trajectory-level ``final_metrics``, falling back to summing per-step ``metrics`` (a - trajectory flushed mid-run may carry steps but no aggregate block). ``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. + 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 {} @@ -776,25 +781,24 @@ def _atif_token_metadata(path: Path | None) -> dict[str, int]: 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) - return {} - captured = { - "prompt_tokens": totals.total_prompt_tokens, - "completion_tokens": totals.total_completion_tokens, - "cache_read_tokens": totals.total_cached_tokens, - } - else: - captured = { - "prompt_tokens": _sum_step_metric(payload, "prompt_tokens"), - "completion_tokens": _sum_step_metric(payload, "completion_tokens"), - "cache_read_tokens": _sum_step_metric(payload, "cached_tokens"), - } - return {key: value for key, value in captured.items() if value is not None} + + 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: 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/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 23d8610c33..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 @@ -761,10 +761,15 @@ def _relay_atif_path(evidence_dir: Path) -> Path | None: def _atif_token_metadata(path: Path | None) -> dict[str, int]: """Project an ATIF trajectory's token totals onto the trial-metadata ``TOKEN_KEYS``. - Prefers the trajectory-level ``final_metrics``, falling back to summing per-step ``metrics`` (a - trajectory flushed mid-run may carry steps but no aggregate block). ``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. + 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 {} @@ -776,25 +781,24 @@ def _atif_token_metadata(path: Path | None) -> dict[str, int]: 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) - return {} - captured = { - "prompt_tokens": totals.total_prompt_tokens, - "completion_tokens": totals.total_completion_tokens, - "cache_read_tokens": totals.total_cached_tokens, - } - else: - captured = { - "prompt_tokens": _sum_step_metric(payload, "prompt_tokens"), - "completion_tokens": _sum_step_metric(payload, "completion_tokens"), - "cache_read_tokens": _sum_step_metric(payload, "cached_tokens"), - } - return {key: value for key, value in captured.items() if value is not None} + + 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: