feat(sdk): Run Receipt — per-run token/cost/timing visibility (fixes silent token-tracking crash)#467
feat(sdk): Run Receipt — per-run token/cost/timing visibility (fixes silent token-tracking crash)#467segevshlomovIBM wants to merge 8 commits into
Conversation
- TokenUsageTracker.on_llm_end dereferenced response.llm_output without a None-guard, raising inside the callback for LiteLLM/watsonx/streaming responses and silently losing token counts - keep the legacy llm_output.token_usage path, fall back to the message's usage_metadata (modern LangChain field) when absent - add unit tests covering both paths, the fallback, and the no-usage no-op Refs #466
- (pattern, input, output) USD-per-1M-token entries for the models shipped in configurations/models/*.toml; longest-pattern-first substring matching on a normalized model name (provider prefixes stripped) - estimate_cost() degrades to None for unknown/self-hosted models and never raises, so callers can rely on it unconditionally - groundwork for per-run cost reporting on InvokeResult Refs #466
- RunMetricsCollector: per-invoke LangChain callback accumulating token usage per model (usage_metadata with legacy token_usage fallback) and LLM wall time; one instance per run so concurrent runs never share state - RunReceipt pydantic model with a dependency-free box-style __str__ render - build_run_receipt() aggregates collector + tracked tool durations; cost reported only when every used model has a known price - new advanced_features.run_receipt flag, default false (zero change when off); every code path is fail-safe and can never affect the run Refs #466
- new InvokeResult.receipt field, populated only when advanced_features.run_receipt is on; None otherwise - invoke() creates a per-run collector, passes it via _apply_callbacks extra_callbacks on both the normal and HITL-resume paths, and forces track_tool_calls internally for timings while result.tool_calls stays gated on the caller's flag (behavior unchanged) - multi-turn threads: snapshot prior tool_calls count so the receipt covers only this invocation; resume path is thread-cumulative (documented) - mock model now emits estimated usage_metadata so receipts work offline - integration test drives CugaAgent.invoke() with a real BaseChatModel fake: flag on -> full receipt; flag off -> receipt is None Refs #466
- new Run Receipt section under the SDK docs: flag snippet, example box output, structured access, pricing-table pointer - note tool-timing semantics (same coverage as track_tool_calls) Refs #466
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (11)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughAdds an opt-in “Run Receipt” to ChangesRun Receipt implementation
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant CugaAgent
participant RunMetricsCollector
participant AgentGraph
participant build_run_receipt
Caller->>CugaAgent: invoke(message)
CugaAgent->>RunMetricsCollector: start collection
CugaAgent->>AgentGraph: run with receipt callback
AgentGraph-->>CugaAgent: result and tool timings
CugaAgent->>build_run_receipt: aggregate metrics and estimate cost
build_run_receipt-->>CugaAgent: RunReceipt
CugaAgent-->>Caller: InvokeResult with receipt
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
coderabbitai
Bot
added
complexity: medium
haroldship
left a comment
haroldship
left a comment
There was a problem hiding this comment.
Thanks — this is a well-built feature: the fail-safe discipline is consistent, the per-run collector correctly avoids the global ActivityTracker's concurrency problem, and test coverage of the aggregation logic is solid. The on_llm_end crash fix alone is worth landing.
Requesting changes for two behavioral issues, plus a few accuracy/design questions in the inline comments:
Blocking
- Enabling run_receipt silently turns on full tool-payload capture (sdk.py). Forcing track_tool_calls=True means every tool call's arguments and results are appended to state.tool_calls and persisted by the checkpointer, accumulating unboundedly across turns on a thread — data retention and state growth the caller never opted into. A metrics flag shouldn't capture payloads; a lightweight mode (name + duration only) or at minimum prominent documentation is needed.
- Resume-path receipt mixes scopes (sdk.py). Tokens/LLM time cover only the resume segment while tool timings are thread-cumulative, so tool_time_s can exceed wall_time_s and counts include earlier turns. The consumer has no way to tell. The normal path's snapshot-slice approach can be applied here too via graph.get_state() before resuming.
Non-blocking, see inline comments
- Pricing: consider litellm.model_cost (already a core dependency, maintained upstream) as the primary source with the static table as fallback — the table already disagrees with litellm on both gpt-oss models.
- The docstring says self-hosted RITS models "resolve to None", but rits/openai/gpt-oss-120b normalizes to a priced entry — which behavior is intended?
- Cached/reasoning token details (input_token_details / output_token_details) are ignored, so cost overestimates on cache-heavy runs.
- response.generations[0][0] in TokenUsageTracker.on_llm_end can still IndexError on an empty generations list — same silent-loss class this PR fixes; a guard would make the fix complete.
| token_usage = (response.llm_output or {}).get("token_usage") or {} | ||
| total_tokens = token_usage.get("total_tokens") | ||
| if total_tokens is None: | ||
| message = getattr(response.generations[0][0], "message", None) |
There was a problem hiding this comment.
generations[0][0] can crash if malformed response
| cost = f"${self.cost_usd:.4f}" if self.cost_usd is not None else "n/a (unknown model)" | ||
| lines = [ | ||
| f"model: {', '.join(self.models) if self.models else 'unknown'}", | ||
| f"tokens: {self.input_tokens:,} in / {self.output_tokens:,} out ({self.total_tokens:,})", |
There was a problem hiding this comment.
could there be cached or reasoning tokens?
|
|
||
| # (substring pattern, USD per 1M input tokens, USD per 1M output tokens) | ||
| MODEL_PRICES: list[Tuple[str, float, float]] = [ | ||
| ("gpt-4.1-mini", 0.40, 1.60), |
There was a problem hiding this comment.
Since litellm is already a core dependency, could litellm.get_model_info() / litellm.model_cost be the primary price source, with this table only as a fallback for models litellm doesn't know? It's maintained upstream and includes cache-read pricing. Two caveats from trying it:
- Bare open-weight names aren't keys —
gpt-oss-120bneeds a provider prefix (groq/openai/gpt-oss-120b,azure_ai/gpt-oss-120b), so the lookup should try the raw (unstripped) model name before the normalized one. - This table already disagrees with litellm on both gpt-oss models (120b: 0.15/0.75 here vs 0.15/0.60 on groq/azure/bedrock in litellm; 20b: 0.10/0.50 vs 0.07/0.30). These look like Groq's launch prices, since reduced — which is the drift this hard-coded approach invites. (minimax-m3 and all the gpt-4.x/gpt-4o entries match litellm exactly — only the two gpt-oss models drifted.)
Note: this is the bundled/GitHub price map in the litellm package, not our LiteLLM proxy — no network dependency beyond an optional GitHub fetch that already falls back to a local copy (LITELLM_LOCAL_MODEL_COST_MAP=True forces offline). If we ever pull live prices from our proxy's /model/info instead, that one is behind VPN and would need fail-safe + caching.
| @@ -0,0 +1,56 @@ | |||
| """Static price estimates for the models shipped in configurations/models/*.toml. | |||
|
|
|||
| Prices are public list prices in USD per 1M tokens and are estimates by design — | |||
There was a problem hiding this comment.
rits/openai/gpt-oss-120b normalizes to gpt-oss-120b and gets the public list price, but the docstring says self-hosted RITS deployments "resolve to None." If pricing RITS usage at notional market rate is intended, the docstring should say so; if not, RITS prefixes should short-circuit to None.
| thread_id=thread_id, | ||
| error=error_msg, | ||
| variables=_hitl_variables, | ||
| receipt=self._build_run_receipt( |
There was a problem hiding this comment.
On the resume path the collector is fresh (tokens/llm_time cover only this resume segment) but tool timings cover the whole thread history, so tool_time_s can exceed wall_time_s and tool_call_count includes calls from before the interrupt and from earlier turns. The code comment acknowledges it, but the receipt gives its consumer no signal that its fields have different scopes. This is fixable the same way as the normal path: read len(state.values.get("tool_calls", [])) via self.graph.get_state(run_config) (already called later at line 2331) before ainvoke and slice.
| receipt_started_at = time.monotonic() | ||
| # Tool durations for the receipt require tracking; the returned | ||
| # result.tool_calls stays gated on the caller's track_tool_calls. | ||
| run_config["configurable"]["track_tool_calls"] = True |
There was a problem hiding this comment.
Forcing track_tool_calls=True doesn't just collect timings. ToolCallTracker.record_call stores the full sanitized arguments and result of every tool call (tracking/tracker.py:98-107), and sandbox_node.py:241 appends those records to state.tool_calls, which is persisted by the checkpointer and accumulates across turns on a thread (that's why this PR needs the snapshot-slice at line 2455). So a metrics flag turns on full payload capture the caller never opted into: unbounded state growth on long-lived threads, and retention of potentially sensitive tool outputs. Consider a lightweight tracking mode (name + duration_ms only) when tracking is receipt-forced, or at minimum document this — the README says "zero overhead when disabled" but not what enabling costs.
- ToolCallTracker.start_tracking(timings_only=True) records tool name/app/ operation/duration but nulls arguments, results, and error payloads - sandbox node maps the "timings_only" track_tool_calls sentinel onto it - review feedback on #467: a metrics flag must not silently capture tool payloads into checkpointed thread state Refs #466
- when run_receipt forces tracking without the caller's track_tool_calls, pass the "timings_only" sentinel so no payloads reach thread state - resume path now snapshots the prior tool_calls count via graph.get_state() before resuming, so the receipt's tool metrics cover only the resume segment — matching the token/time scope of the per-invoke collector - README: document what enabling the receipt records, and the litellm-based pricing source - review feedback on #467 (both blocking items) Refs #466
…ards - litellm's bundled model_cost map is now the primary price source (raw model name tried before the normalized one), with the static table as fallback; gpt-oss entries corrected to litellm's current list prices - rits/... model names short-circuit to None — self-hosted serving is never billed at notional market rate (docstring now matches behavior) - cache-read input tokens are collected from usage_metadata, billed at the model's cache-read rate when known, and surfaced on the receipt - empty/malformed generations no longer abort usage collection in either TokenUsageTracker or RunMetricsCollector (llm_output fallback still lands) - review feedback on #467 (non-blocking items) Refs #466
|
All of the review points are fixed, verified end-to-end, and pushed. I added a timings_only tracking mode so that when the receipt forces tool tracking, only tool name/app/duration reach the checkpointed state, never arguments, results, or errors, and the resume path now snapshots the prior tool_calls count via graph.get_state() so the receipt covers only the resume segment (both blockers). Pricing switched to litellm's bundled model_cost map as the primary source (static table as fallback, gpt-oss prices corrected, rits/... short-circuits to None), with cache-read tokens now collected and billed at the cache-read rate. I also guarded the empty-generations crash in both usage collectors, updated docs/tests (95 passing, re-verified E2E on watsonx), and pushed the 3 commits to the PR |
|
@sami-marreed I've approved this PR |
2 participants
Summary
Adds a Run Receipt — per-run token usage, estimated cost, and LLM-vs-tools time split — behind a new
advanced_features.run_receiptflag (default off, zero behavior change when disabled), and fixes a silent token-tracking crash on the default providers.Closes #466
The problem
TokenUsageTracker.on_llm_enddereferencedresponse.llm_outputwithout a None-guard, so token usage silently crashed and was lost on LiteLLM/watsonx/streaming responses.duration_msalready being recorded.What this PR does
When
run_receipt = true, everyCugaAgent.invoke()returnsresult.receipt:(real output from a watsonx
gpt-oss-120brun)RunMetricsCollector) — one instance perinvoke(), so concurrent runs never share counters. ReadsAIMessage.usage_metadatawith a legacyllm_output["token_usage"]fallback, and captures every LLM call in the run (main loop, reflection, summarizer).backend/llm/pricing.py) seeded with the models shipped inconfigurations/models/*.toml; unknown/self-hosted models degrade toest. cost: n/a.duration_msfrom tool-call tracking (same semantics astrack_tool_calls); the returnedresult.tool_callsstays gated on the caller's flag — unchanged behavior.receipt=None, never affects the run.usage_metadatafallback inTokenUsageTracker— kept as a separate first commit so it can be split out if desired.Verification
run_tests.sh --skip-stability) green locally.CUGA_MOCK_LLM=true) and real (watsonxgpt-oss-120b— the exact provider whose token tracking previously crashed): flag on → receipt above; flag off →receipt is None, behavior identical tomain.Out of scope / follow-ups
/streampath.agent.stream()(e.g.agent.last_receipt).timings_onlytracker mode so receipt-forced tracking records only{name, duration_ms}instead of full tool arguments/results in thread state.Summary by CodeRabbit
run_receiptflag and the receipt fields/output.