llm_cost_tracker 0.14.1 → 0.14.2

data/app/services/llm_cost_tracker/dashboard/top_models.rb

@@ -4,7 +4,6 @@ module LlmCostTracker
   module Dashboard
     class TopModels
       DEFAULT_LIMIT = 5
-      SORT_OPTIONS = %w[cost calls avg_cost latency tokens provider name].freeze
       DEFAULT_SORT = "cost"
       DEFAULT_DIRECTIONS = {
         "provider" => "asc",
@@ -43,7 +42,7 @@ module LlmCostTracker
       def initialize(scope:, limit:, sort: DEFAULT_SORT, direction: nil)
         @scope = scope
         @limit = limit
-        @sort = SORT_OPTIONS.include?(sort.to_s) ? sort.to_s : DEFAULT_SORT
+        @sort = DEFAULT_DIRECTIONS.key?(sort.to_s) ? sort.to_s : DEFAULT_SORT
         @direction = Sort::DIRECTIONS.include?(direction.to_s) ? direction.to_s : DEFAULT_DIRECTIONS[@sort]
       end
 

data/app/views/llm_cost_tracker/calls/show.html.erb

@@ -5,13 +5,10 @@
   value = @call[token_key]
   { label: component.fetch(:label), value: value, formatted_value: number_with_delimiter(value), css_class: component.fetch(:css_class) }
 end %>
-<% cost_segments = [] %>
-<% unless @call.total_cost.nil? %>
-  <% cost_segments = priced_components.map do |component|
-    value = line_item_costs_by_component[component.fetch(:price_key)]
-    { label: component.fetch(:label), value: value, formatted_value: optional_money(value), css_class: component.fetch(:css_class) }
-  end %>
-<% end %>
+<% cost_segments = priced_components.map do |component|
+  value = line_item_costs_by_component[component.fetch(:price_key)]
+  { label: component.fetch(:label), value: value, formatted_value: optional_money(value), css_class: component.fetch(:css_class) }
+end %>
 <% service_line_items = @call.line_items.reject { |li| li.unit == "token" }.sort_by(&:position) %>
 
 <p class="lct-breadcrumb-back"><%= link_to "← All calls", calls_path %></p>
@@ -84,22 +81,11 @@ end %>
         <dd><%= content_tag(:code, @call.provider_response_id, class: "lct-code-id") %></dd>
       </div>
     <% end %>
-    <% if @call.provider_project_id.present? %>
-      <div class="lct-meta-strip-item">
-        <dt>Project ID</dt>
-        <dd><%= content_tag(:code, LlmCostTracker::Dashboard::Masking.mask_value(:provider_project_id, @call.provider_project_id), class: "lct-code-id") %></dd>
-      </div>
-    <% end %>
-    <% if @call.provider_api_key_id.present? %>
-      <div class="lct-meta-strip-item">
-        <dt>API Key ID</dt>
-        <dd><%= content_tag(:code, LlmCostTracker::Dashboard::Masking.mask_value(:provider_api_key_id, @call.provider_api_key_id), class: "lct-code-id") %></dd>
-      </div>
-    <% end %>
-    <% if @call.provider_workspace_id.present? %>
+    <% { provider_project_id: "Project ID", provider_api_key_id: "API Key ID", provider_workspace_id: "Workspace ID" }.each do |attribute, label| %>
+      <% next if @call[attribute].blank? %>
       <div class="lct-meta-strip-item">
-        <dt>Workspace ID</dt>
-        <dd><%= content_tag(:code, LlmCostTracker::Dashboard::Masking.mask_value(:provider_workspace_id, @call.provider_workspace_id), class: "lct-code-id") %></dd>
+        <dt><%= label %></dt>
+        <dd><%= content_tag(:code, LlmCostTracker::Dashboard::Masking.mask_value(attribute, @call[attribute]), class: "lct-code-id") %></dd>
       </div>
     <% end %>
   </dl>

data/app/views/llm_cost_tracker/data_quality/index.html.erb

@@ -271,7 +271,7 @@
         <h2 class="lct-panel-title">Incomplete pricing by model</h2>
         <span class="lct-panel-meta"><%= link_to "View calls →", calls_path(current_query(cost_status: "incomplete", sort: nil)) %></span>
       </div>
-      <p class="lct-panel-intro">These models have line items without configured rates, so totals undercount (a full row is missing all rates; a partial row has some line items priced and others not). After the next price refresh or a <code>pricing_overrides</code> update, run <code>bin/rails llm_cost_tracker:backfill_unknown_pricing</code> to recompute these calls.</p>
+      <p class="lct-panel-intro">These models have line items without configured rates, so totals undercount (a full row is missing all rates; a partial row has some line items priced and others not). After the next price refresh or a <code>pricing_overrides</code> update, run <code>bin/rails llm_cost_tracker:backfill_unknown_pricing</code> to recompute these calls. A partial call is left as recorded when a rate it already used has changed since.</p>
 
       <div class="lct-table-wrap">
         <table class="lct-tbl">

data/app/views/llm_cost_tracker/shared/_bar.html.erb

@@ -1,5 +1,3 @@
-<% variant_class = local_assigns[:variant] == "budget" ? " lct-budget-fill" : "" %>
-
 <div class="lct-bar-track" aria-hidden="true">
-  <div data-lct-style="<%= inline_style("width: #{bar_width(value, max)}") %>" class="lct-bar-fill<%= variant_class %>"></div>
+  <div data-lct-style="<%= inline_style("width: #{bar_width(value, max)}") %>" class="lct-bar-fill"></div>
 </div>

data/app/views/llm_cost_tracker/shared/_filter_pill_date.html.erb

@@ -1,5 +1,3 @@
-<% extra_hidden = local_assigns.fetch(:extra_hidden, {}) %>
-<% extra_except = local_assigns.fetch(:extra_except, []) %>
 <details class="lct-filter-pop" name="lct-filter">
   <summary class="lct-filter-pill">
     <svg width="12" height="12" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><rect x="3" y="4" width="18" height="18" rx="2"/><path d="M16 2v4M8 2v4M3 10h18"/></svg>
@@ -8,8 +6,7 @@
     <svg class="lct-chev" width="10" height="6" viewBox="0 0 10 6" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round"><path d="M1 1l4 4 4-4"/></svg>
   </summary>
   <%= form_with url: path, method: :get, local: true, html: { class: "lct-filter-pop-body" } do %>
-    <% extra_hidden.each do |k, v| %><%= hidden_field_tag k, v %><% end %>
-    <%= hidden_query_fields(current_query.except(:from, :to, :page, :per, *extra_except)) %>
+    <%= hidden_query_fields(current_query.except(:from, :to, :page, :per)) %>
     <div class="lct-filter-pop-field"><label for="lct-filter-from">From</label><input type="date" name="from" id="lct-filter-from" value="<%= @from_date.iso8601 %>"></div>
     <div class="lct-filter-pop-field"><label for="lct-filter-to">To</label><input type="date" name="to" id="lct-filter-to" value="<%= @to_date.iso8601 %>"></div>
     <button type="submit" class="lct-button lct-button-primary">Apply</button>

data/app/views/llm_cost_tracker/shared/_filter_pill_model.html.erb

@@ -1,5 +1,3 @@
-<% extra_hidden = local_assigns.fetch(:extra_hidden, {}) %>
-<% extra_except = local_assigns.fetch(:extra_except, []) %>
 <% active = params[:model].presence %>
 <details class="lct-filter-pop" name="lct-filter">
   <summary class="lct-filter-pill <%= 'lct-active' if active %>">
@@ -9,8 +7,7 @@
     <svg class="lct-chev" width="10" height="6" viewBox="0 0 10 6" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round"><path d="M1 1l4 4 4-4"/></svg>
   </summary>
   <%= form_with url: path, method: :get, local: true, html: { class: "lct-filter-pop-body" } do %>
-    <% extra_hidden.each do |k, v| %><%= hidden_field_tag k, v %><% end %>
-    <%= hidden_query_fields(current_query.except(:model, :page, :per, *extra_except)) %>
+    <%= hidden_query_fields(current_query.except(:model, :page, :per)) %>
     <div class="lct-filter-pop-field">
       <label for="lct-filter-model">Model</label>
       <%= select_tag :model, options_for_select(filter_options_for(:model), active), include_blank: "All models", id: "lct-filter-model" %>

data/app/views/llm_cost_tracker/shared/_filter_pill_provider.html.erb

@@ -1,5 +1,3 @@
-<% extra_hidden = local_assigns.fetch(:extra_hidden, {}) %>
-<% extra_except = local_assigns.fetch(:extra_except, []) %>
 <% active = params[:provider].presence %>
 <details class="lct-filter-pop" name="lct-filter">
   <summary class="lct-filter-pill <%= 'lct-active' if active %>">
@@ -9,8 +7,7 @@
     <svg class="lct-chev" width="10" height="6" viewBox="0 0 10 6" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round"><path d="M1 1l4 4 4-4"/></svg>
   </summary>
   <%= form_with url: path, method: :get, local: true, html: { class: "lct-filter-pop-body" } do %>
-    <% extra_hidden.each do |k, v| %><%= hidden_field_tag k, v %><% end %>
-    <%= hidden_query_fields(current_query.except(:provider, :page, :per, *extra_except)) %>
+    <%= hidden_query_fields(current_query.except(:provider, :page, :per)) %>
     <div class="lct-filter-pop-field">
       <label for="lct-filter-provider">Provider</label>
       <%= select_tag :provider, options_for_select(filter_options_for(:provider), active), include_blank: "All providers", id: "lct-filter-provider" %>

data/app/views/llm_cost_tracker/shared/_filter_pill_stream.html.erb

@@ -1,5 +1,3 @@
-<% extra_hidden = local_assigns.fetch(:extra_hidden, {}) %>
-<% extra_except = local_assigns.fetch(:extra_except, []) %>
 <% active = params[:stream].presence %>
 <% display_value = case active when "yes" then "Streaming" when "no" then "Non-streaming" else "All" end %>
 <details class="lct-filter-pop" name="lct-filter">
@@ -10,8 +8,7 @@
     <svg class="lct-chev" width="10" height="6" viewBox="0 0 10 6" fill="none" stroke="currentColor" stroke-width="1.5" stroke-linecap="round" stroke-linejoin="round"><path d="M1 1l4 4 4-4"/></svg>
   </summary>
   <%= form_with url: path, method: :get, local: true, html: { class: "lct-filter-pop-body" } do %>
-    <% extra_hidden.each do |k, v| %><%= hidden_field_tag k, v %><% end %>
-    <%= hidden_query_fields(current_query.except(:stream, :page, :per, *extra_except)) %>
+    <%= hidden_query_fields(current_query.except(:stream, :page, :per)) %>
     <div class="lct-filter-pop-field">
       <label for="lct-filter-stream">Stream</label>
       <%= select_tag :stream, options_for_select(LlmCostTracker::Dashboard::Filter::STREAM_FILTER_OPTIONS, active), include_blank: "All calls", id: "lct-filter-stream" %>

data/app/views/llm_cost_tracker/tags/show.html.erb

@@ -5,9 +5,9 @@
   <h2 class="lct-page-title">Tag: <code class="lct-code-id"><%= params[:key] %></code> = <code class="lct-code-id"><%= @value %></code></h2>
 
   <div class="lct-filter-row">
-    <%= render "llm_cost_tracker/shared/filter_pill_date", path: form_url, extra_hidden: { tag_value: @value }, extra_except: [:tag_value] %>
-    <%= render "llm_cost_tracker/shared/filter_pill_provider", path: form_url, extra_hidden: { tag_value: @value }, extra_except: [:tag_value] %>
-    <%= render "llm_cost_tracker/shared/filter_pill_model", path: form_url, extra_hidden: { tag_value: @value }, extra_except: [:tag_value] %>
+    <%= render "llm_cost_tracker/shared/filter_pill_date", path: form_url %>
+    <%= render "llm_cost_tracker/shared/filter_pill_provider", path: form_url %>
+    <%= render "llm_cost_tracker/shared/filter_pill_model", path: form_url %>
 
     <% if params[:provider].present? || params[:model].present? %>
       <%= link_to "× Clear filters", tag_path(params[:key], current_query(provider: nil, model: nil, page: nil)), class: "lct-filter-clear" %>
@@ -35,7 +35,7 @@
       </div>
       <div class="lct-stat">
         <div class="lct-stat-head"><p class="lct-stat-label">Avg cost / call</p></div>
-        <p class="lct-stat-value"><%= money(@value_calls.positive? ? @value_total_cost / @value_calls : 0) %></p>
+        <p class="lct-stat-value"><%= money(@value_total_cost / @value_calls) %></p>
       </div>
     </div>
 
@@ -44,7 +44,7 @@
         <h2 class="lct-panel-title">Spend over time</h2>
         <span class="lct-panel-meta"><%= link_to "View calls →", calls_path(calls_query_for_tag(key: @key, value: @value)) %></span>
       </div>
-      <%= render "llm_cost_tracker/shared/spend_chart", series: @value_points, comparison_series: [] %>
+      <%= render "llm_cost_tracker/shared/spend_chart", series: @value_points %>
     </section>
   <% end %>
 

data/lib/llm_cost_tracker/budget/per_tag.rb

@@ -30,7 +30,7 @@ module LlmCostTracker
         end
 
         def blocking?
-          active? && configured.each_value.any? { |entry| behavior_for(entry) == :block_requests }
+          configured.each_value.any? { |entry| behavior_for(entry) == :block_requests } && active?
         end
 
         def rules_for(tags, blocking_only: false)
@@ -38,13 +38,13 @@ module LlmCostTracker
 
           normalized = (tags || {}).to_h.transform_keys(&:to_s)
           configured.filter_map do |key, entry|
-            raw = normalized[key]
-            next if raw.nil?
+            value = Ledger::Tags::Encoding.encode(normalized[key])
+            next if value.empty?
             next if blocking_only && behavior_for(entry) != :block_requests
 
             Rule.new(
               key: key,
-              value: Ledger::Tags::Encoding.encode(raw),
+              value: value,
               windows: entry[:windows],
               behavior: behavior_for(entry),
               on_exceeded: on_exceeded_for(entry)
@@ -53,21 +53,27 @@ module LlmCostTracker
         end
 
         def rules_for_events(events)
-          events.each_with_object({}) do |event, grouped|
+          events.select(&:total_cost).each_with_object({}) do |event, grouped|
             rules_for(event.tags).each { |rule| (grouped[rule] ||= []) << event }
           end
         end
 
         def spend(key, value, window, time:)
-          spend_by_value(key, [value], window, window_start(window, time)).fetch(value, 0).to_d
+          spend_by_value(key, [value], window, window_start(window, time)).fetch(value, [0]).first.to_d
         end
 
-        def spend_by_value(key, values, window, bucket)
+        def spend_by_value(key, values, window, bucket, upto = nil)
           started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          range = window_range(window, bucket)
+          upto_sum = LlmCostTracker::CallTag.sanitize_sql_array(
+            ["SUM(CASE WHEN #{TIME_COLUMN} <= ? THEN #{COST_COLUMN} ELSE 0 END)", upto || range.end]
+          )
           totals = Ledger::Isolation.guard(LlmCostTracker::CallTag) do
             LlmCostTracker::CallTag
-              .where(key: key, value: values, TIME_COLUMN => window_range(window, bucket))
-              .group(:value).sum(COST_COLUMN)
+              .where(key: key, value: values, TIME_COLUMN => range)
+              .group(:value)
+              .pluck(:value, Arel.sql("SUM(#{COST_COLUMN})"), Arel.sql(upto_sum))
+              .to_h { |value, *sums| [value, sums.map(&:to_d)] }
           end
           warn_slow_read(key, window, Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at)
           totals

data/lib/llm_cost_tracker/budget.rb

@@ -37,25 +37,40 @@ module LlmCostTracker
         config = LlmCostTracker.configuration
         return unless event.total_cost
 
-        check_per_call_budget(event, config, behavior_override)
+        errors = [check_per_call_budget(event, config, behavior_override)]
         check_windowed({ daily: config.budgets.daily, monthly: config.budgets.monthly }.compact,
                        time: event.tracked_at) do |budget_type, total, budget|
-          handle_exceeded(budget_type: budget_type,
-                          total: total,
-                          budget: budget,
-                          previous_total: total - event.total_cost,
-                          last_event: event,
-                          behavior: behavior_override)
+          errors << handle_exceeded(budget_type: budget_type,
+                                    total: total,
+                                    budget: budget,
+                                    previous_total: total - event.total_cost,
+                                    last_event: event,
+                                    behavior: behavior_override)
         end
+        errors.concat(persisted_errors([event], behavior_override)) unless Ingestion.async?
+        raise_first(errors)
       end
 
       def check_persisted!(events, behavior_override: nil)
-        by_rule = PerTag.rules_for_events(events.select(&:total_cost))
+        raise_first(persisted_errors(events, behavior_override))
+      end
+
+      def notify_persisted_safely!(events)
+        check_persisted!(events, behavior_override: :notify)
+      rescue StandardError => e
+        Logging.warn("Per-tag budget check failed after ingest: #{e.class}: #{e.message}")
+      end
+
+      private
+
+      def persisted_errors(events, behavior_override)
+        by_rule = PerTag.rules_for_events(events)
         by_rule = by_rule.reject { |rule, _| rule.on_exceeded.nil? } if behavior_override == :notify
-        window_buckets(by_rule).each do |(key, window, bucket), scored|
-          totals = PerTag.spend_by_value(key, scored.keys.map(&:value), window, bucket)
-          scored.each do |rule, recorded|
-            total = totals.fetch(rule.value, 0).to_d
+        window_buckets(by_rule).flat_map do |(key, window, bucket), scored|
+          upto = scored.values.flatten.map(&:tracked_at).max unless Ingestion.async?
+          totals = PerTag.spend_by_value(key, scored.keys.map(&:value), window, bucket, upto)
+          scored.filter_map do |rule, recorded|
+            total, upto_total = totals.fetch(rule.value, [0, 0])
             limit = rule.windows.fetch(window)
             next if total < limit
 
@@ -63,7 +78,7 @@ module LlmCostTracker
               budget_type: window,
               total: total,
               budget: limit,
-              previous_total: total - recorded.sum(&:total_cost),
+              previous_total: (upto_total >= limit ? upto_total : total) - recorded.sum(&:total_cost),
               last_event: recorded.last,
               scope: scope_for(rule),
               behavior: behavior_override || rule.behavior,
@@ -73,14 +88,11 @@ module LlmCostTracker
         end
       end
 
-      def notify_persisted_safely!(events)
-        check_persisted!(events, behavior_override: :notify)
-      rescue StandardError => e
-        Logging.warn("Per-tag budget check failed after ingest: #{e.class}: #{e.message}")
+      def raise_first(errors)
+        error = errors.compact.first
+        raise error if error
       end
 
-      private
-
       def window_buckets(by_rule)
         by_rule.each_with_object({}) do |(rule, events), grouped|
           rule.windows.each_key do |window|
@@ -98,14 +110,6 @@ module LlmCostTracker
         Pricing::Estimator.call(provider: provider, model: model, request: request) || BigDecimal("0")
       end
 
-      def raise_per_call_pre_send(estimate, budget)
-        return unless estimate >= budget
-
-        raise BudgetExceededError.new(**budget_payload(
-          budget_type: :per_call, total: estimate, budget: budget, last_event: nil, stage: :pre_send
-        ))
-      end
-
       def check_per_call_budget(event, config, behavior_override)
         budget = config.budgets.per_call
         return unless budget
@@ -122,7 +126,10 @@ module LlmCostTracker
       end
 
       def enforce_globally(config, estimate:, time:)
-        raise_per_call_pre_send(estimate, config.budgets.per_call) if config.budgets.per_call && estimate.positive?
+        per_call = config.budgets.per_call
+        if per_call && estimate.positive? && estimate >= per_call
+          raise_pre_send(budget_type: :per_call, total: estimate, budget: per_call)
+        end
 
         check_windowed({ monthly: config.budgets.monthly, daily: config.budgets.daily }.compact,
                        time: time,
@@ -132,9 +139,9 @@ module LlmCostTracker
       end
 
       def raise_pre_send(budget_type:, total:, budget:, scope: nil)
-        raise BudgetExceededError.new(**budget_payload(
-          budget_type: budget_type, total: total, budget: budget, last_event: nil, stage: :pre_send, scope: scope
-        ))
+        raise BudgetExceededError.new(
+          budget_type: budget_type, total: total, budget: budget, stage: :pre_send, scope: scope
+        )
       end
 
       def check_per_tag(tags, time:, estimate: BigDecimal("0"), blocking_only: false)
@@ -153,21 +160,14 @@ module LlmCostTracker
       def check_windowed(budgets, time:, estimate: BigDecimal("0"))
         return if budgets.empty?
 
-        totals = totals_for(budgets.keys, time: time)
+        periods = budgets.keys.map { |type| BUDGET_TYPE_TO_PERIOD.fetch(type) }
+        totals = LlmCostTracker::Ledger::Period::Totals.call(periods, time: time)
         budgets.each do |budget_type, budget|
-          total = totals.fetch(budget_type) + estimate
+          total = totals.fetch(BUDGET_TYPE_TO_PERIOD.fetch(budget_type)) + estimate
           yield(budget_type, total, budget) if total >= budget
         end
       end
 
-      def totals_for(budget_types, time:)
-        return {} if budget_types.empty?
-
-        period_for = budget_types.to_h { |type| [type, BUDGET_TYPE_TO_PERIOD.fetch(type)] }
-        period_totals = LlmCostTracker::Ledger::Period::Totals.call(period_for.values, time: time)
-        period_for.transform_values { |period| period_totals.fetch(period) }
-      end
-
       def handle_exceeded(budget_type:,
                           total:,
                           budget:,
@@ -179,28 +179,17 @@ module LlmCostTracker
         config = LlmCostTracker.configuration
         behavior ||= config.budgets.exceeded_behavior
         on_exceeded ||= config.budgets.on_exceeded
-        payload = budget_payload(
+        payload = {
           budget_type: budget_type,
           total: total,
           budget: budget,
           last_event: last_event,
           stage: :post_spend,
           scope: scope
-        )
+        }
 
         on_exceeded.call(payload) if on_exceeded && (previous_total.nil? || previous_total < budget)
-        raise BudgetExceededError.new(**payload) if %i[raise block_requests].include?(behavior)
-      end
-
-      def budget_payload(budget_type:, total:, budget:, last_event:, stage:, scope: nil)
-        {
-          budget_type: budget_type,
-          total: total,
-          budget: budget,
-          last_event: last_event,
-          stage: stage,
-          scope: scope
-        }
+        BudgetExceededError.new(**payload) if %i[raise block_requests].include?(behavior)
       end
     end
   end

data/lib/llm_cost_tracker/capture/event_window.rb

@@ -8,7 +8,7 @@ module LlmCostTracker
       HEAD_EVENTS = 16
       TAIL_EVENTS = 32
       HEAVY_STRING_BYTES = 8 * 1024
-      IGNORED_PAYLOAD_KEYS = %w[b64_json partial_image_b64].freeze
+      IGNORED_PAYLOAD_KEYS = %w[b64_json partial_image_b64 snapshot logprobs].freeze
 
       def initialize(notable: nil)
         @notable = notable
@@ -20,7 +20,8 @@ module LlmCostTracker
       end
 
       def push(data, type: nil)
-        return if @overflowed
+        # The openai gem's chat stream helper resends all logprobs so far in logprobs.* events; none carry usage.
+        return if @overflowed || type&.start_with?("logprobs.")
 
         event = { event: type, data: strip_heavy_payload(data) }
         size = approximate_bytesize(event)

data/lib/llm_cost_tracker/capture/sdk_payload.rb

@@ -3,6 +3,8 @@
 require "active_support/core_ext/object/deep_dup"
 require "active_support/core_ext/object/try"
 
+require_relative "event_window"
+
 module LlmCostTracker
   module Capture
     module SdkPayload
@@ -11,7 +13,9 @@ module LlmCostTracker
       def normalize(value)
         case value
         when Hash
-          value.each_with_object({}) { |(key, nested), out| out[key.to_s] = normalize(nested) }
+          value.each_with_object({}) do |(key, nested), out|
+            out[key.to_s] = normalize(nested) unless EventWindow::IGNORED_PAYLOAD_KEYS.include?(key.to_s)
+          end
         when Array
           value.map { |nested| normalize(nested) }
         when Symbol

data/lib/llm_cost_tracker/capture/stream_collector.rb

@@ -5,6 +5,7 @@ require "active_support/core_ext/object/deep_dup"
 require "json"
 
 require_relative "event_window"
+require_relative "sdk_payload"
 require_relative "../timing"
 
 module LlmCostTracker
@@ -21,7 +22,6 @@ module LlmCostTracker
                      provider_workspace_id: nil,
                      pricing_mode: nil,
                      metadata: {},
-                     context_tags: nil,
                      request: nil)
         @provider = provider.to_s
         @model = model
@@ -32,7 +32,7 @@ module LlmCostTracker
         @provider_workspace_id = provider_workspace_id
         @pricing_mode = pricing_mode
         @metadata = (metadata || {}).deep_dup
-        @context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).deep_dup
+        @context_tags = LlmCostTracker::Tags::Context.tags.deep_dup
         @request = request
         @window = EventWindow.new(notable: Parsers.find_for_provider(@provider)&.method(:retain_stream_event?))
         @explicit_usage = nil
@@ -57,9 +57,10 @@ module LlmCostTracker
       end
 
       def event(data, type: nil)
+        data = SdkPayload.normalize(data) if data.is_a?(Hash)
         @mutex.synchronize do
           ensure_open!
-          @window.push(data, type: type) unless data.nil?
+          @window.push(data, type: type&.to_s) unless data.nil?
         end
       end
 
@@ -91,6 +92,8 @@ module LlmCostTracker
         record_snapshot(snapshot, errored: errored)
       rescue TransactionAbortedError
         raise
+      rescue ActiveRecord::RecordNotUnique
+        nil
       rescue StandardError => e
         raise unless errored
 
@@ -165,20 +168,27 @@ module LlmCostTracker
       end
 
       def build_event(snapshot)
-        return build_from_explicit_usage(snapshot) if snapshot[:explicit_usage]
-        return build_unknown_usage(snapshot) if snapshot[:overflowed]
+        return build_unparsed_event(snapshot) if snapshot[:explicit_usage]
 
-        event = Parsers.find_for_provider(@provider)&.parse_stream(
-          response_status: 200,
-          events: snapshot[:events],
-          request_body: request_body_for(snapshot[:request])
-        )
+        if snapshot[:overflowed]
+          Logging.warn("#{@provider} stream events exceeded #{SSE::LIMIT_BYTES} bytes; " \
+                       "recording usage_source=#{Usage::Source::UNKNOWN}.")
+          return build_unparsed_event(snapshot)
+        end
+
+        request_body = request_body_for(snapshot[:request])
+        events = Parsers.all_for_provider(@provider).filter_map do |parser|
+          parser.parse_stream(
+            response_status: 200, events: snapshot[:events], request_body: request_body, model: snapshot[:model]
+          )
+        end
+        event = events.find { |parsed| parsed.usage_source != Usage::Source::UNKNOWN } || events.first
         if event
           model = present_model(event.model) || present_model(snapshot[:model]) || Event::UNKNOWN_MODEL
           return event.with(provider: @provider, model: model, **snapshot.fetch(:capture_dimensions))
         end
 
-        build_unknown_usage(snapshot)
+        build_unparsed_event(snapshot)
       end
 
       def request_body_for(request)
@@ -194,25 +204,14 @@ module LlmCostTracker
         string unless string == Event::UNKNOWN_MODEL
       end
 
-      def build_from_explicit_usage(snapshot)
-        Event.build(
-          provider: @provider,
-          model: snapshot[:model] || Event::UNKNOWN_MODEL,
-          token_usage: snapshot[:explicit_usage],
-          stream: true,
-          usage_source: Usage::Source::MANUAL,
-          pricing_mode: snapshot[:pricing_mode],
-          **snapshot.fetch(:capture_dimensions)
-        )
-      end
-
-      def build_unknown_usage(snapshot)
+      def build_unparsed_event(snapshot)
+        explicit_usage = snapshot[:explicit_usage]
         Event.build(
           provider: @provider,
           model: snapshot[:model] || Event::UNKNOWN_MODEL,
-          token_usage: Usage::TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
+          token_usage: explicit_usage || Usage::TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
-          usage_source: Usage::Source::UNKNOWN,
+          usage_source: explicit_usage ? Usage::Source::MANUAL : Usage::Source::UNKNOWN,
           pricing_mode: snapshot[:pricing_mode],
           **snapshot.fetch(:capture_dimensions)
         )

data/lib/llm_cost_tracker/capture/stream_tracker.rb

@@ -1,10 +1,7 @@
 # frozen_string_literal: true
 
-require "active_support/core_ext/object/deep_dup"
 require "active_support/core_ext/object/try"
 
-require_relative "sdk_payload"
-
 module LlmCostTracker
   module Capture
     class StreamTracker
@@ -21,23 +18,15 @@ module LlmCostTracker
       def wrap
         return @stream unless @stream
 
-        iterator_wrapped = false
-        if @stream.instance_variable_defined?(:@iterator)
-          iterator = @stream.instance_variable_get(:@iterator)
-          if iterator.respond_to?(:each)
-            @stream.instance_variable_set(:@iterator,
-                                          Enumerator.new do |yielder|
-                                            each_from(iterator) { |event| yielder << event }
-                                          end)
-            iterator_wrapped = true
-          end
-        end
-        each_wrapped = false
-        if !iterator_wrapped && @stream.respond_to?(:each)
+        iterator = @stream.instance_variable_get(:@iterator)
+        if iterator.respond_to?(:each)
+          @stream.instance_variable_set(:@iterator,
+                                        Enumerator.new do |yielder|
+                                          each_from(iterator) { |event| yielder << event }
+                                        end)
+        elsif @stream.respond_to?(:each)
           wrap_each
-          each_wrapped = true
-        end
-        unless iterator_wrapped || each_wrapped
+        else
           Logging.warn(
             "stream integration found no wrappable iterator on #{@stream.class} " \
             "(missing both `@iterator` ivar and `#each`); usage will not be captured"
@@ -83,10 +72,7 @@ module LlmCostTracker
       end
 
       def capture(event)
-        raw_payload = event.try(:deep_to_h) || event.try(:to_h) || {}
-        payload = SdkPayload.normalize(raw_payload)
-        type = event.try(:type) || payload["type"]
-        @collector.event(payload, type: type&.to_s)
+        @collector.event(event.try(:to_h) || {}, type: event.try(:type)&.to_s)
       rescue StandardError => e
         warn_capture_failure(e)
       end

data/lib/llm_cost_tracker/capture_verifier.rb

@@ -25,7 +25,7 @@ module LlmCostTracker
       [
         enabled_check,
         *integration_checks,
-        *storage_checks
+        *LlmCostTracker::Ingestion.verify
       ].compact
     end
 
@@ -49,11 +49,5 @@ module LlmCostTracker
         check.with(name: "sdk integration #{check.name}")
       end
     end
-
-    def storage_checks
-      LlmCostTracker::Ingestion.verify
-    rescue LlmCostTracker::Error => e
-      [Check.new(:error, "storage", e.message)]
-    end
   end
 end