Skip to content

Track LLM token usage in reports#244

Open
HexSleeves wants to merge 2 commits into
NVIDIA:mainfrom
HexSleeves:codex/add-llm-token-usage-reporting
Open

Track LLM token usage in reports#244
HexSleeves wants to merge 2 commits into
NVIDIA:mainfrom
HexSleeves:codex/add-llm-token-usage-reporting

Conversation

@HexSleeves

Copy link
Copy Markdown

Motivation

  • Capture per-call token telemetry so report metadata can surface LLM token consumption and detect cost/usage trends.
  • Preserve existing structured-output parsing while retaining raw provider metadata needed for usage (LangChain raw payload).

Description

  • Add token counters to LLMCallRecord and llm_call_record() defaults: input_tokens, output_tokens, and total_tokens.
  • Use with_structured_output(..., include_raw=True) and unwrap {"raw","parsed","parsing_error"} responses, extracting usage from raw.usage_metadata (supports input_tokens/output_tokens and prompt_tokens/completion_tokens variants) and preserving parsing-error behavior.
  • Record token usage in both sync (run_batches) and async (arun_batches) paths (including non-structured raw responses), and propagate analyzer.llm_usage into node-level llm_call_log entries (success and failure).
  • Aggregate llm_call_log token counters into report JSON metadata under metadata.llm_usage while keeping existing metadata intact.
  • Add support in the CLI structured-output adapter for include_raw=True so local/CLI providers expose raw usage without changing parser expectations.
  • Add unit tests covering the new fields, sync/async usage capture, missing usage metadata handling, structured-output include_raw=True parsing, and metadata aggregation; closes feat: expose LLM token usage in JSON report output #242.

Testing

  • Ran ruff format . and ruff check . with no issues.
  • Ran the full test suite with pytest -q: all tests passed (1261 passed, 12 skipped, 34 deselected, 6 xfailed).
  • mypy src was executed and reported pre-existing typing issues in unrelated modules; no new type regressions introduced by these changes (mypy warnings are repository pre-existing and not caused by this PR).

Signed-off-by: Jacob <lecoqjacob@gmail.com>
@HexSleeves HexSleeves force-pushed the codex/add-llm-token-usage-reporting branch from 3bb9d24 to 622ab7d Compare July 6, 2026 02:16

@rng1995 rng1995 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Automated SkillSpector Review]

Requesting changes. Usage capture works for LLMAnalyzerBase, but the report still undercounts the real TP4 graph LLM call. Instrument the direct chat_completion() path and add a regression asserting provider usage is included in metadata.llm_usage.

Comment thread tests/test_mcp_tool_poisoning.py Outdated
):
result = node(state)
assert result["llm_call_log"] == [{"node": "mcp_tool_poisoning", "ok": True, "error": None}]
assert result["llm_call_log"] == [llm_call_record("mcp_tool_poisoning", ok=True)]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blocking: this expectation bakes in zero tokens for TP4 even though production mcp_tool_poisoning.py calls chat_completion() directly and records a successful LLM call. That path bypasses the new raw-response usage extraction, so report totals are incomplete. Instrument chat_completion/TP4 and assert the provider's nonzero usage metadata here.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 8b45e81. chat_completion_with_usage() now preserves normalized provider token metadata, TP4 forwards it into llm_call_log (including parse failures), and the regression verifies the nonzero TP4 usage reaches metadata.llm_usage. Verification: Ruff clean; 1262 tests passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: expose LLM token usage in JSON report output

2 participants