Module: Models::AssistantConversationTokenTrackable

Extended by:
ActiveSupport::Concern
Included in:
AssistantConversation
Defined in:
app/concerns/models/assistant_conversation_token_trackable.rb

Overview

Aggregates LLM token usage, cost, and response-time metrics on an
AssistantConversation.

Mixed into AssistantConversation to keep its token/cost bookkeeping out of
the model itself: per-query counters via #track_query!, recomputation from
RubyLLM's per-attempt usage ledger via #computed_token_totals, and the
authoritative cost rollup via #logged_cost_usd.

See Also:

Instance Method Summary collapse

Instance Method Details

#computed_token_totals ⇒ Hash{Symbol => Integer}

Compute token totals from RubyLLM's usage ledger (source of truth).

Was a SUM over assistant_messages' own token columns until RubyLLM 2.0
moved usage into ruby_llm_usage_entries — one row per API attempt rather
than one per message, so a turn that retried now contributes every attempt
it was billed for instead of only the one that produced a message.

Returns { input: N, output: N, thinking: N, cached: N, cache_creation: N, total: N }

Returns:

  • (Hash{Symbol => Integer}) —

    token totals keyed by usage bucket;
    +:total+ is +:input+ + +:output+



82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
 
# File 'app/concerns/models/assistant_conversation_token_trackable.rb', line 82

def computed_token_totals
  sums = ruby_llm_usages.unscope(:order)
                        .pick(
                          Arel.sql('COALESCE(SUM(input_tokens), 0)'),
                          Arel.sql('COALESCE(SUM(output_tokens), 0)'),
                          Arel.sql('COALESCE(SUM(thinking_tokens), 0)'),
                          Arel.sql('COALESCE(SUM(cache_read_tokens), 0)'),
                          Arel.sql('COALESCE(SUM(cache_write_tokens), 0)')
                        ) || [0, 0, 0, 0, 0]
  {
    input: sums[0].to_i,
    output: sums[1].to_i,
    thinking: sums[2].to_i,
    cached: sums[3].to_i,
    cache_creation: sums[4].to_i,
    total: sums[0].to_i + sums[1].to_i
  }
end

#computed_total_cost ⇒ Float

Note:

Prefer #logged_cost_usd for display — this re-prices through
Assistant::CostCalculator and can drift from the authoritative ledger.

Compute total cost (USD) for this conversation from per-message token data.
Uses each message's associated LlmModel to look up the correct pricing.
Falls back to the conversation-level model when a message has no model association.

Returns:

  • (Float) —

    total cost in USD



108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
 
# File 'app/concerns/models/assistant_conversation_token_trackable.rb', line 108

def computed_total_cost
  # Build a model_id → model_key lookup from ChatService::MODELS
  model_id_to_key = Assistant::ChatService::MODELS.transform_values { |v| v[:id] }.invert

  assistant_messages
    .where(role: 'assistant')
    .includes(:ruby_llm_usages)
    .sum do |message|
      model_key = model_id_to_key[message.model_id] || llm_model_name
      # RubyLLM 2.0 keeps token counts in its own per-attempt ledger rather
      # than on the message; #tokens aggregates a message's entries.
      tokens = message.tokens

      Assistant::CostCalculator.cost_for(
        model_key,
        input_tokens: tokens.input.to_i,
        output_tokens: tokens.output.to_i,
        cached_tokens: tokens.cache_read.to_i,
        cache_creation_tokens: tokens.cache_write.to_i
      )
    end
end

#logged_cost_usd ⇒ Float

Conversation cost (USD) summed from ai_usage_logs — the same authoritative
per-turn ledger the admin dashboard reads, so the figure shown in-chat
reconciles with the admin chat-by-user rollup. Prefer this over
computed_total_cost for display: the latter re-prices from assistant_messages
through a separate pricing table (Assistant::CostCalculator) and can drift.

Scoped to feature 'assistant_chat' to match that rollup — an
AssistantConversation is also the subject for context_compaction /
daily_focus_analysis logs, which would otherwise leak into the chat cost.

Returns:

  • (Float) —

    cost in USD (0.0 when nothing logged yet)



142
143
144
145
146
147
 
# File 'app/concerns/models/assistant_conversation_token_trackable.rb', line 142

def logged_cost_usd
  AiUsageLog
    .where(subject_type: 'AssistantConversation', subject_id: id, feature: 'assistant_chat')
    .sum(:cost_microdollars)
    .to_f / AiUsageLog::MICRODOLLARS_PER_DOLLAR
end

#sync_token_totals! ⇒ AssistantConversation

Sync aggregate metadata from assistant_messages (call after responses complete).
Fixes the issue where track_query! only captures last-chunk tokens.
Also computes and caches total cost for fast sidebar display.

Uses a database-level JSONB merge (metadata || patch) instead of a
Ruby-side Hash#merge so that keys written by concurrent callers — most
importantly compaction_summary set by ContextCompactor — are never
clobbered by a stale in-memory copy of metadata.

Returns:



159
160
161
162
163
164
165
166
167
168
169
170
 
# File 'app/concerns/models/assistant_conversation_token_trackable.rb', line 159

def sync_token_totals!
  totals = computed_token_totals
  cost = computed_total_cost
  patch = {
    'total_input_tokens' => totals[:input],
    'total_output_tokens' => totals[:output],
    'total_cost_cents' => cost
  }.to_json

  self.class.where(id: id).update_all(["metadata = metadata || ?::jsonb", patch])
  reload
end

#total_tokens ⇒ Integer

Total tokens used — returns cached metadata totals (synced after each response),
falls back to computing from assistant_messages only when metadata is empty.
This avoids N+1 queries when displaying token counts in conversation lists.

Returns:

  • (Integer) —

    total input + output tokens



64
65
66
67
68
69
 
# File 'app/concerns/models/assistant_conversation_token_trackable.rb', line 64

def total_tokens
  cached = (total_input_tokens || 0) + (total_output_tokens || 0)
  return cached if cached.positive?

  computed_token_totals[:total]
end

#track_error! ⇒ Boolean

Track an error

Returns:

  • (Boolean) —

    result of +save!+

Raises:

  • (ActiveRecord::RecordInvalid) —

    if the conversation fails validation



54
55
56
57
 
# File 'app/concerns/models/assistant_conversation_token_trackable.rb', line 54

def track_error!
  self.error_count = (error_count || 0) + 1
  save!
end

#track_query!(model:, input_tokens: 0, output_tokens: 0, response_time: nil, tool_stats: {}) ⇒ Boolean

Track a completed query with its metrics

Parameters:

  • model (String) —

    LLM model name that served the query (cached on
    +llm_model_name+)

  • input_tokens (Integer) (defaults to: 0) —

    prompt tokens billed for this query

  • output_tokens (Integer) (defaults to: 0) —

    completion tokens billed for this query

  • response_time (Float, nil) (defaults to: nil) —

    wall-clock seconds for the query; folded
    into the running +average_response_time+ when present

  • tool_stats (Hash) (defaults to: {}) —

    per-query tool telemetry (+:total_tool_calls+,
    +:sql_errors+, +:patch_errors+); empty Hash means no tool activity

Returns:

  • (Boolean) —

    result of +save!+

Raises:

  • (ActiveRecord::RecordInvalid) —

    if the conversation fails validation



28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
 
# File 'app/concerns/models/assistant_conversation_token_trackable.rb', line 28

def track_query!(model:, input_tokens: 0, output_tokens: 0, response_time: nil, tool_stats: {})
  self.llm_model_name = model
  self.total_input_tokens = (total_input_tokens || 0) + input_tokens
  self.total_output_tokens = (total_output_tokens || 0) + output_tokens
  self.total_queries = (total_queries || 0) + 1
  self.last_query_at = Time.current

  if response_time
    current_avg = average_response_time || 0
    current_count = (total_queries || 1) - 1
    self.average_response_time = ((current_avg * current_count) + response_time) / total_queries
  end

  if tool_stats.present?
    self.total_tool_calls = (total_tool_calls || 0) + (tool_stats[:total_tool_calls] || 0)
    self.total_tool_errors = (total_tool_errors || 0) +
                             (tool_stats[:sql_errors] || 0) + (tool_stats[:patch_errors] || 0)
  end

  save!
end