Track LLM token usage in reports#244
Conversation
Signed-off-by: Jacob <[email protected]>
3bb9d24 to
622ab7d
Compare
rng1995
left a comment
There was a problem hiding this comment.
[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.
| ): | ||
| 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)] |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.
rng1995
left a comment
There was a problem hiding this comment.
[Automated SkillSpector Review]
Re-review of head 8b45e81. The PR adds token counters to LLMCallRecord, captures provider usage in LLMAnalyzerBase (sync/async, structured via include_raw=True and raw paths), instruments the TP4 chat_completion path via a new chat_completion_with_usage(), and aggregates totals into report JSON under metadata.llm_usage.
Prior-issue resolution checklist
- TP4
chat_completion()path bypassed usage extraction; test baked in zero tokens → Resolved. Commit8b45e81addschat_completion_with_usage()insrc/skillspector/llm_utils.pyreturning(content, usage);_check_tp4now records**usageon both the ok record and the post-attempt failure record (so usage from a successful call whose JSON fails to parse is still counted). The regressiontest_successful_call_records_provider_usageasserts nonzero provider usage (12/3/15) lands in thellm_call_logrecord and flows through_build_metadataintometadata.llm_usage.chat_completion()is kept as a thin back-compat wrapper, and the integration test's patch target was updated.
New-commit scan (8b45e81)
No new blockers. The token-usage helpers were sensibly moved from llm_analyzer_base into llm_utils (LLMTokenUsage, empty_token_usage, extract_token_usage) so both transports share one normalizer. Parse-failure behavior is preserved (_unwrap_structured_response re-raises parsing_error, so the ValidationError degradation branches still trigger), and the CLI adapter's fail-closed ValueError on garbage JSON is unchanged for the default include_raw=False path.
Schema contract
metadata.llm_usageis additive in the JSON report; terminal/markdown/SARIF outputs are untouched.LLMCallRecordgains three int fields, but every producer goes throughllm_call_record()(defaults 0) and test assertions were migrated to the builder, so the log shape stays consistent. Non-breaking.
Non-blocking nits
extract_token_usage(src/skillspector/llm_utils.py:73-75):int(...)on malformed provider metadata can raise inside the live call path — see inline comment.- The
_NoUsageshim +locals().get("analyzer", _NoUsage())pattern is duplicated in four modules (semantic_developer_intent.py,semantic_quality_policy.py,semantic_security_discovery.py,meta_analyzer.py). Initializinganalyzer: LLMAnalyzerBase | None = Nonebefore thetryand usinganalyzer.llm_usage if analyzer else empty_token_usage()would be clearer and avoid the shared mutable class-attribute dict. - The PR description claims a test for "structured-output
include_raw=Trueparsing", but_StructuredAgentCLIModel.invokewithinclude_raw=True(the dict-returning branch, including a parse failure captured asparsing_error) has no direct test — worth adding alongsidetest_structured_output_parses_and_validatesintests/unit/test_llm_utils.py.
| usage = getattr(raw, "usage_metadata", None) or {} | ||
| if not isinstance(usage, dict): | ||
| return empty_token_usage() | ||
| input_tokens = int(usage.get("input_tokens") or usage.get("prompt_tokens") or 0) |
There was a problem hiding this comment.
Non-blocking: these int(...) conversions run inline in the live LLM call path (now including TP4 via chat_completion_with_usage). If a provider ever returns a non-numeric usage value (e.g. "input_tokens": "n/a"), the raised ValueError/TypeError would convert a successful analyzer call into a failed one — degrading the scan and dropping findings over telemetry. Consider wrapping the conversions in a try/except (TypeError, ValueError): return empty_token_usage() so usage extraction can never fail the detection path.
Motivation
rawpayload).Description
LLMCallRecordandllm_call_record()defaults:input_tokens,output_tokens, andtotal_tokens.with_structured_output(..., include_raw=True)and unwrap{"raw","parsed","parsing_error"}responses, extracting usage fromraw.usage_metadata(supportsinput_tokens/output_tokensandprompt_tokens/completion_tokensvariants) and preserving parsing-error behavior.run_batches) and async (arun_batches) paths (including non-structured raw responses), and propagateanalyzer.llm_usageinto node-levelllm_call_logentries (success and failure).llm_call_logtoken counters into report JSON metadata undermetadata.llm_usagewhile keeping existing metadata intact.include_raw=Trueso local/CLI providers exposerawusage without changing parser expectations.include_raw=Trueparsing, and metadata aggregation; closes feat: expose LLM token usage in JSON report output #242.Testing
ruff format .andruff check .with no issues.pytest -q: all tests passed (1261 passed, 12 skipped, 34 deselected, 6 xfailed).mypy srcwas 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).