Skip to content

feat: per-model cost tracking and message attribution - #945

Open
jmars wants to merge 1 commit into
mistralai:mainfrom
jmars:feat/per-model-cost-tracking
Open

feat: per-model cost tracking and message attribution#945
jmars wants to merge 1 commit into
mistralai:mainfrom
jmars:feat/per-model-cost-tracking

Conversation

@jmars

@jmars jmars commented Jul 27, 2026

Copy link
Copy Markdown

Problem

AgentStats.session_cost uses a single flat pricing pair for the entire session. When /model switches the active model mid-session, all accumulated tokens are retroactively re-priced — update_pricing() even documents this as "a known approximation." Additionally, session transcripts (messages.jsonl) have no way to identify which model generated each message.

Solution

Two changes:

1. Per-model cost tracking in AgentStats

  • New PerModelTokenUsage model tracks prompt_tokens, completion_tokens, and per-model pricing
  • AgentStats.per_model: dict[str, PerModelTokenUsage] accumulates usage keyed by model name
  • session_cost sums across all models at their respective rates (falls back to legacy single-model calc when per_model is empty)
  • add_tokens_for_model() attributes each turn's usage to the active model; _ensure_per_model() snapshots the pricing at first attribution so switching A→B→A preserves correct rates
  • update_pricing() now only sets the current-model default — past tokens are untouched

2. Model attribution on LLMMessage

  • New optional model and provider fields (default None, omitted via exclude_none in serialization)
  • Set on assistant messages in both streaming (_chat_streaming) and non-streaming (_complete) paths
  • Propagated through chunk accumulation (LLMMessage.__add__) and model_dump

Before / after (cost example)

Session: 1M prompt + 50K completion on deepseek ($0.14/$0.28/M)
Switch to claude ($3.00/$15.00/M), one turn: 1.5M prompt + 30K completion

Before:  $1.45  (all tokens at claude pricing — wrong)
After:   $0.15 + $4.95 = $5.10  (each model at its own rate — correct)

Backward compatibility

  • Old messages.jsonl (no model/provider): fields default to None, serialization omits them
  • Old meta.json (no per_model): defaults to empty dict, session_cost falls back to legacy calculation
  • No migration needed — per_model populates lazily on next turn

Token attribution note

Each turn reports the full context prompt_tokens the API billed for, not just the incremental delta. When a model switch occurs, the new model's turn includes the cost of re-encoding the entire history — so per-model totals correctly reflect what each model was actually charged. Prompt caching discounts and tokenization differences are provider-side and not reflected in usage responses.

Files changed

  • vibe/core/types.py: +80/-8 (PerModelTokenUsage, AgentStats.per_model + methods, LLMMessage.model/provider)
  • vibe/core/agent_loop/_loop.py: +24/-2 (message stamping, model_name param on _update_stats)

@jmars
jmars requested a review from a team as a code owner July 27, 2026 10:58
Add per-model token breakdown to AgentStats for accurate
session_cost when models change mid-session. Stamp each
assistant LLMMessage with model/provider attribution.

- New PerModelTokenUsage model tracks tokens + pricing per model
- AgentStats.per_model dict accumulates per-model usage
- session_cost sums across all models with correct per-model rates
- LLMMessage gains model/provider fields (None on non-assistant msgs)
- _update_stats accepts optional model_name for attribution
- Streaming and non-streaming paths both stamp messages

Backward compatible: old session logs without model/provider
or per_model load with defaults, falling back to legacy cost calc.
@jmars
jmars force-pushed the feat/per-model-cost-tracking branch from bdffdc9 to 931f9c1 Compare July 27, 2026 11:00
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.

1 participant