llm_cost_tracker 0.4.0 → 0.4.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: ccb9a8365f4a06026a4352385efa1318ac59ce403cb848e0c9aff992fc80f64c
-  data.tar.gz: f21503cd322e923dc5bde0139cc61bc1547cef01eac59fe7a3861e1ab33e9860
+  metadata.gz: d2cdd5f30c6fbd8c0168549b0853e9d8bc54586e60921733ce11a89a1d86078c
+  data.tar.gz: c91384579df6acdeb04d24b62f8bf916040f98156fd2bc882c94afc534f7dba5
 SHA512:
-  metadata.gz: 304ab6de6404f070b21b1dd72ce9eae2b44fb2fc7845eae8831a04971ed2b8ec2b6f740bc082fb36cfa42d90f0be59ab5800d43d72b68f04918a113b6d7d8cbd
-  data.tar.gz: afa2e92a99062bb1e0b4a00ab1d0762ca688f1890e0d76a29801881e2319e68db217c036d8f8c5d99558b580d5d4c039f8b3334283631763ee093fd12d369329
+  metadata.gz: 88d61d6714101ee9e8162814f5527bde487eced83663d86a1f938b77bcee1e4fcb4db2c4dde763720a828368a779a8677da57ecf10b2fadec78a959e6fdce6a7
+  data.tar.gz: d2d2bb097058507c06c1ea330a0c7d5a63d2e92b824fd8e4e7640a052dff100038eadb2e097edb4500457c382368080414c5261a2cc0a4f69436ae2234fd420d

data/CHANGELOG.md

@@ -4,6 +4,22 @@ Format: [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). Versioning: [S
 
 ## [Unreleased]
 
+## [0.4.1] - 2026-04-24
+
+### Changed
+
+- Batched ActiveRecord period rollup writes and budget total reads.
+- Memoized schema capability checks and refreshed them on `reset_column_information`.
+- Install migration adds `[:model, :tracked_at]` composite index and drops redundant single-column `:provider` / `:model` indexes.
+- Data Quality now reads counters and usage sums through one aggregate query.
+- Parser URL matching, stream-event extraction, and custom parser registration now share a smaller base/registry extension surface.
+- Added cookbook recipes for `ruby-openai`, `anthropic-sdk-ruby`, `gemini-ai`, `langchainrb`, Azure OpenAI, and LiteLLM proxy setups.
+
+### Fixed
+
+- `llm_cost_tracker:add_period_totals` now imports legacy monthly rollups and backfills before adding the unique index.
+- Budget docs now describe `:notify` across monthly, daily, and per-call budgets.
+
 ## [0.4.0] - 2026-04-24
 
 ### Changed

data/README.md

@@ -161,6 +161,8 @@ end
 
 Run `bin/rails g llm_cost_tracker:add_streaming` once on existing installs to add the `stream` and `usage_source` columns. Run `bin/rails g llm_cost_tracker:add_provider_response_id` to persist provider-issued response IDs. Run `bin/rails g llm_cost_tracker:add_usage_breakdown` to add cache-read, cache-write, hidden-output, and pricing-mode columns.
 
+More client-specific snippets live in [`docs/cookbook.md`](docs/cookbook.md).
+
 ### Manual tracking
 
 ```ruby
@@ -269,7 +271,7 @@ config.per_call_budget = 1.00
 config.budget_exceeded_behavior = :block_requests
 ```
 
-- `:notify` — fire `on_budget_exceeded` after an event pushes the month over budget.
+- `:notify` — fire `on_budget_exceeded` after an event pushes the monthly, daily, or per-call budget over the limit.
 - `:raise` — record the event, then raise `BudgetExceededError`.
 - `:block_requests` — block preflight when the stored monthly or daily total is already over budget; still raises post-response on the event that crosses the line. Needs `:active_record` storage for preflight.
 
@@ -456,21 +458,23 @@ Configured hosts are parsed using the OpenAI-compatible usage shape (`prompt_tok
 For providers with a non-OpenAI usage shape:
 
 ```ruby
-require "uri"
-
 class AcmeParser < LlmCostTracker::Parsers::Base
+  HOSTS = %w[api.acme-llm.example].freeze
+  TRACKED_PATHS = %w[/v1/generate].freeze
+
+  def provider_names
+    %w[acme]
+  end
+
   def match?(url)
-    uri = URI.parse(url.to_s)
-    uri.host == "api.acme-llm.example" && uri.path == "/v1/generate"
-  rescue URI::InvalidURIError
-    false
+    match_uri?(url, hosts: HOSTS, exact_paths: TRACKED_PATHS)
   end
 
-  def parse(request_url, request_body, response_status, response_body)
+  def parse(_request_url, _request_body, response_status, response_body)
     return nil unless response_status == 200
 
     payload = safe_json_parse(response_body)
-    usage = payload&.dig("usage")
+    usage = payload.dig("usage")
     return nil unless usage
 
     LlmCostTracker::ParsedUsage.build(
@@ -482,7 +486,7 @@ class AcmeParser < LlmCostTracker::Parsers::Base
   end
 end
 
-LlmCostTracker::Parsers::Registry.register(AcmeParser.new)
+LlmCostTracker::Parsers::Registry.register(AcmeParser)
 ```
 
 ## Supported providers
@@ -531,7 +535,7 @@ The gem is designed for multi-threaded hosts — Puma with `max_threads > 1` and
 
 ## Development
 
-Architecture rules for future changes live in [`docs/architecture.md`](docs/architecture.md).
+Architecture rules for future changes live in [`docs/architecture.md`](docs/architecture.md). Integration recipes live in [`docs/cookbook.md`](docs/cookbook.md).
 
 ```bash
 bundle install

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

@@ -29,84 +29,85 @@ module LlmCostTracker
     class DataQuality
       class << self
         def call(scope: LlmCostTracker::LlmApiCall.all)
-          total = scope.count
+          model = scope.klass
+          aggregates = DataQualityAggregate.call(scope: scope)
+          total = aggregates.fetch(:total_calls).to_i
 
           DataQualityStats.new(
             total_calls: total,
-            unknown_pricing_count: scope.unknown_pricing.count,
-            untagged_calls_count: total - scope.with_json_tags.count,
-            **latency_stats(scope),
-            **stream_stats(scope),
-            **provider_response_id_stats(scope),
-            **usage_stats(scope),
+            unknown_pricing_count: aggregates.fetch(:unknown_pricing_count).to_i,
+            untagged_calls_count: total - aggregates.fetch(:tagged_calls_count).to_i,
+            **latency_stats(aggregates, model:),
+            **stream_stats(aggregates, model:),
+            **provider_response_id_stats(aggregates, model:),
+            **usage_stats(aggregates, model:),
             unknown_pricing_by_model: unknown_pricing_by_model(scope)
           )
         end
 
         private
 
-        def latency_stats(scope)
-          latency_present = LlmCostTracker::LlmApiCall.latency_column?
+        def latency_stats(aggregates, model:)
+          latency_present = model.latency_column?
 
           {
-            missing_latency_count: latency_present ? scope.where(latency_ms: nil).count : nil,
+            missing_latency_count: latency_present ? aggregates.fetch(:missing_latency_count).to_i : nil,
             latency_column_present: latency_present
           }
         end
 
-        def stream_stats(scope)
-          stream_present = LlmCostTracker::LlmApiCall.stream_column?
+        def stream_stats(aggregates, model:)
+          stream_present = model.stream_column?
+          usage_source_present = model.usage_source_column?
+          streaming_missing_usage_count = nil
+          if stream_present && usage_source_present
+            streaming_missing_usage_count = aggregates.fetch(:streaming_missing_usage_count).to_i
+          end
 
           {
-            streaming_count: stream_present ? scope.streaming.count : nil,
-            streaming_missing_usage_count: streaming_missing_usage_count(scope, stream_present),
+            streaming_count: stream_present ? aggregates.fetch(:streaming_count).to_i : nil,
+            streaming_missing_usage_count: streaming_missing_usage_count,
             stream_column_present: stream_present
           }
         end
 
-        def provider_response_id_stats(scope)
-          column_present = LlmCostTracker::LlmApiCall.provider_response_id_column?
+        def provider_response_id_stats(aggregates, model:)
+          column_present = model.provider_response_id_column?
+          missing_provider_response_id_count = nil
+          if column_present
+            missing_provider_response_id_count = aggregates.fetch(:missing_provider_response_id_count).to_i
+          end
 
           {
-            missing_provider_response_id_count: column_present ? scope.missing_provider_response_id.count : nil,
+            missing_provider_response_id_count: missing_provider_response_id_count,
             provider_response_id_column_present: column_present
           }
         end
 
-        def usage_stats(scope)
-          usage_breakdown_present = LlmCostTracker::LlmApiCall.usage_breakdown_columns?
-          usage_breakdown_cost_present = LlmCostTracker::LlmApiCall.usage_breakdown_cost_columns?
-          sums = sum_columns(scope, usage_sum_columns(usage_breakdown_present, usage_breakdown_cost_present))
+        def usage_stats(aggregates, model:)
+          usage_breakdown_present = model.usage_breakdown_columns?
+          usage_breakdown_cost_present = model.usage_breakdown_cost_columns?
+          cache_read_input_cost = nil
+          cache_write_input_cost = nil
+          if usage_breakdown_cost_present
+            cache_read_input_cost = decimal_sum(aggregates.fetch(:cache_read_input_cost))
+            cache_write_input_cost = decimal_sum(aggregates.fetch(:cache_write_input_cost))
+          end
 
           {
             usage_breakdown_column_present: usage_breakdown_present,
-            input_tokens: sums[:input_tokens].to_i,
-            cache_read_input_tokens: usage_breakdown_present ? sums[:cache_read_input_tokens].to_i : nil,
-            cache_write_input_tokens: usage_breakdown_present ? sums[:cache_write_input_tokens].to_i : nil,
-            output_tokens: sums[:output_tokens].to_i,
-            hidden_output_tokens: usage_breakdown_present ? sums[:hidden_output_tokens].to_i : nil,
-            input_cost: decimal_sum(sums[:input_cost]),
-            cache_read_input_cost: usage_breakdown_cost_present ? decimal_sum(sums[:cache_read_input_cost]) : nil,
-            cache_write_input_cost: usage_breakdown_cost_present ? decimal_sum(sums[:cache_write_input_cost]) : nil,
-            output_cost: decimal_sum(sums[:output_cost])
+            input_tokens: aggregates.fetch(:input_tokens).to_i,
+            cache_read_input_tokens: usage_breakdown_present ? aggregates.fetch(:cache_read_input_tokens).to_i : nil,
+            cache_write_input_tokens: usage_breakdown_present ? aggregates.fetch(:cache_write_input_tokens).to_i : nil,
+            output_tokens: aggregates.fetch(:output_tokens).to_i,
+            hidden_output_tokens: usage_breakdown_present ? aggregates.fetch(:hidden_output_tokens).to_i : nil,
+            input_cost: decimal_sum(aggregates.fetch(:input_cost)),
+            cache_read_input_cost: cache_read_input_cost,
+            cache_write_input_cost: cache_write_input_cost,
+            output_cost: decimal_sum(aggregates.fetch(:output_cost))
           }
         end
 
-        def usage_sum_columns(usage_breakdown_present, usage_breakdown_cost_present)
-          columns = %i[input_tokens output_tokens input_cost output_cost]
-          if usage_breakdown_present
-            columns += %i[cache_read_input_tokens cache_write_input_tokens hidden_output_tokens]
-          end
-          columns += %i[cache_read_input_cost cache_write_input_cost] if usage_breakdown_cost_present
-          columns
-        end
-
-        def streaming_missing_usage_count(scope, stream_present)
-          return unless stream_present && LlmCostTracker::LlmApiCall.usage_source_column?
-
-          scope.streaming_missing_usage.count
-        end
-
         def unknown_pricing_by_model(scope)
           scope.unknown_pricing
                .group(:model)
@@ -116,16 +117,6 @@ module LlmCostTracker
                .to_h
         end
 
-        def sum_columns(scope, columns)
-          values = scope.unscope(:order).pick(*columns.map { |column| sum_expression(scope, column) })
-
-          columns.zip(values).to_h
-        end
-
-        def sum_expression(scope, column)
-          Arel.sql("COALESCE(SUM(#{scope.connection.quote_column_name(column)}), 0)")
-        end
-
         def decimal_sum(value)
           value.to_f.round(8)
         end

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

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Dashboard
+    class DataQualityAggregate
+      class << self
+        def call(scope:)
+          model = scope.klass
+          expressions = aggregate_expressions(scope, model:)
+          values = Array(scope.unscope(:order).pick(*expressions.values))
+
+          expressions.keys.zip(values).to_h
+        end
+
+        private
+
+        def aggregate_expressions(scope, model:)
+          usage_breakdown_present = model.usage_breakdown_columns?
+          usage_breakdown_cost_present = model.usage_breakdown_cost_columns?
+
+          expressions = {
+            total_calls: Arel.sql("COUNT(*)"),
+            unknown_pricing_count: conditional_count_expression("total_cost IS NULL"),
+            tagged_calls_count: tagged_calls_expression(model)
+          }
+
+          if model.latency_column?
+            expressions[:missing_latency_count] = conditional_count_expression("latency_ms IS NULL")
+          end
+          expressions[:streaming_count] = conditional_count_expression("stream") if model.stream_column?
+          if model.stream_column? && model.usage_source_column?
+            expressions[:streaming_missing_usage_count] =
+              conditional_count_expression("stream AND (usage_source = 'unknown' OR usage_source IS NULL)")
+          end
+          if model.provider_response_id_column?
+            expressions[:missing_provider_response_id_count] =
+              conditional_count_expression("provider_response_id IS NULL OR provider_response_id = ''")
+          end
+
+          usage_sum_columns(usage_breakdown_present, usage_breakdown_cost_present).each do |column|
+            expressions[column] = sum_expression(scope, column)
+          end
+
+          expressions
+        end
+
+        def usage_sum_columns(usage_breakdown_present, usage_breakdown_cost_present)
+          columns = %i[input_tokens output_tokens input_cost output_cost]
+          if usage_breakdown_present
+            columns += %i[cache_read_input_tokens cache_write_input_tokens hidden_output_tokens]
+          end
+          columns += %i[cache_read_input_cost cache_write_input_cost] if usage_breakdown_cost_present
+          columns
+        end
+
+        def conditional_count_expression(predicate)
+          Arel.sql("COALESCE(SUM(CASE WHEN #{predicate} THEN 1 ELSE 0 END), 0)")
+        end
+
+        def tagged_calls_expression(model)
+          table = model.quoted_table_name
+          column = "#{table}.#{model.connection.quote_column_name('tags')}"
+
+          Arel.sql(case
+                   when model.tags_jsonb_column?
+                     "COALESCE(SUM(CASE WHEN #{column} <> '{}'::jsonb THEN 1 ELSE 0 END), 0)"
+                   when model.tags_mysql_json_column?
+                     "COALESCE(SUM(CASE WHEN JSON_LENGTH(#{column}) > 0 THEN 1 ELSE 0 END), 0)"
+                   else
+                     "COALESCE(SUM(CASE WHEN #{column} IS NOT NULL AND #{column} <> '' " \
+                     "AND #{column} <> '{}' THEN 1 ELSE 0 END), 0)"
+                   end)
+        end
+
+        def sum_expression(scope, column)
+          Arel.sql("COALESCE(SUM(#{scope.connection.quote_column_name(column)}), 0)")
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/budget.rb

@@ -10,8 +10,16 @@ module LlmCostTracker
         return unless config.budget_exceeded_behavior == :block_requests
         return unless config.active_record?
 
-        enforce_period_budget(:monthly, config.monthly_budget)
-        enforce_period_budget(:daily, config.daily_budget)
+        budgets = enforce_period_budgets(config)
+        return if budgets.empty?
+
+        totals = active_record_totals(budgets.keys, time: Time.now.utc)
+
+        budgets.each do |period, budget|
+          total = totals.fetch(period)
+
+          handle_exceeded(budget_type: period, total: total, budget: budget) if total >= budget
+        end
       end
 
       def check!(event)
@@ -19,21 +27,18 @@ module LlmCostTracker
         return unless event.cost
 
         check_per_call_budget(event, config)
-        check_period_budget(event, config, :daily, config.daily_budget)
-        check_period_budget(event, config, :monthly, config.monthly_budget)
-      end
+        budgets = check_period_budgets(config)
+        totals = totals_for_check(event, config, budgets)
 
-      private
-
-      def enforce_period_budget(period, budget)
-        return unless budget
+        budgets.each do |period, budget|
+          total = totals.fetch(period)
 
-        total = active_record_total(period, time: Time.now.utc)
-        return unless total >= budget
-
-        handle_exceeded(budget_type: period, total: total, budget: budget)
+          handle_exceeded(budget_type: period, total: total, budget: budget, last_event: event) if total >= budget
+        end
       end
 
+      private
+
       def check_per_call_budget(event, config)
         budget = config.per_call_budget
         return unless budget
@@ -44,40 +49,32 @@ module LlmCostTracker
         handle_exceeded(budget_type: :per_call, total: call_cost, budget: budget, last_event: event)
       end
 
-      def check_period_budget(event, config, period, budget)
-        return unless budget
-
-        total = if config.active_record?
-                  active_record_total(period, time: event.tracked_at)
-                else
-                  event.cost.total_cost
-                end
-        return unless total >= budget
-
-        handle_exceeded(budget_type: period, total: total, budget: budget, last_event: event)
+      def enforce_period_budgets(config)
+        {
+          monthly: config.monthly_budget,
+          daily: config.daily_budget
+        }.compact
       end
 
-      def active_record_total(period, time:)
-        case period
-        when :monthly then active_record_monthly_total(time: time)
-        when :daily   then active_record_daily_total(time: time)
-        end
+      def check_period_budgets(config)
+        {
+          daily: config.daily_budget,
+          monthly: config.monthly_budget
+        }.compact
       end
 
-      def active_record_monthly_total(time: Time.now.utc)
-        require_relative "llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
-        require_relative "storage/active_record_store" unless defined?(LlmCostTracker::Storage::ActiveRecordStore)
+      def totals_for_check(event, config, budgets)
+        return {} if budgets.empty?
+        return active_record_totals(budgets.keys, time: event.tracked_at) if config.active_record?
 
-        LlmCostTracker::Storage::ActiveRecordStore.monthly_total(time: time)
-      rescue LoadError => e
-        raise Error, "ActiveRecord storage requires the active_record gem: #{e.message}"
+        budgets.to_h { |period, _budget| [period, event.cost.total_cost] }
       end
 
-      def active_record_daily_total(time: Time.now.utc)
+      def active_record_totals(periods, time:)
         require_relative "llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
         require_relative "storage/active_record_store" unless defined?(LlmCostTracker::Storage::ActiveRecordStore)
 
-        LlmCostTracker::Storage::ActiveRecordStore.daily_total(time: time)
+        LlmCostTracker::Storage::ActiveRecordStore.period_totals(periods, time: time)
       rescue LoadError => e
         raise Error, "ActiveRecord storage requires the active_record gem: #{e.message}"
       end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/add_period_totals_to_llm_cost_tracker.rb.erb

@@ -8,10 +8,10 @@ class AddPeriodTotalsToLlmCostTracker < ActiveRecord::Migration<%= migration_ver
       t.timestamps
     end unless table_exists?(:llm_cost_tracker_period_totals)
 
+    backfill_period_totals
+
     add_index :llm_cost_tracker_period_totals, [:period, :period_start],
               unique: true unless index_exists?(:llm_cost_tracker_period_totals, [:period, :period_start])
-
-    backfill_period_totals
   end
 
   def down
@@ -22,23 +22,53 @@ class AddPeriodTotalsToLlmCostTracker < ActiveRecord::Migration<%= migration_ver
   private
 
   def backfill_period_totals
+    backfill_legacy_monthly_totals if table_exists?(:llm_cost_tracker_monthly_totals)
     return unless table_exists?(:llm_api_calls)
 
     backfill_period_total("day", day_bucket_sql)
     backfill_period_total("month", month_bucket_sql)
   end
 
+  def backfill_legacy_monthly_totals
+    execute <<~SQL
+      INSERT INTO llm_cost_tracker_period_totals (period, period_start, total_cost, created_at, updated_at)
+      SELECT #{connection.quote("month")} AS period,
+             month AS period_start,
+             total_cost,
+             CURRENT_TIMESTAMP,
+             CURRENT_TIMESTAMP
+      FROM llm_cost_tracker_monthly_totals legacy
+      WHERE NOT EXISTS (
+        SELECT 1
+        FROM llm_cost_tracker_period_totals existing
+        WHERE existing.period = #{connection.quote("month")}
+          AND existing.period_start = legacy.month
+      )
+    SQL
+  end
+
   def backfill_period_total(period, bucket_sql)
     execute <<~SQL
       INSERT INTO llm_cost_tracker_period_totals (period, period_start, total_cost, created_at, updated_at)
-      SELECT #{connection.quote(period)} AS period,
-             #{bucket_sql} AS period_start,
-             SUM(total_cost) AS total_cost,
+      SELECT aggregated.period,
+             aggregated.period_start,
+             aggregated.total_cost,
              CURRENT_TIMESTAMP,
              CURRENT_TIMESTAMP
-      FROM llm_api_calls
-      WHERE total_cost IS NOT NULL
-      GROUP BY #{bucket_sql}
+      FROM (
+        SELECT #{connection.quote(period)} AS period,
+               #{bucket_sql} AS period_start,
+               SUM(total_cost) AS total_cost
+        FROM llm_api_calls
+        WHERE total_cost IS NOT NULL
+        GROUP BY #{bucket_sql}
+      ) aggregated
+      WHERE NOT EXISTS (
+        SELECT 1
+        FROM llm_cost_tracker_period_totals existing
+        WHERE existing.period = aggregated.period
+          AND existing.period_start = aggregated.period_start
+      )
     SQL
   end
 

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/create_llm_api_calls.rb.erb

@@ -37,10 +37,9 @@ class CreateLlmApiCalls < ActiveRecord::Migration<%= migration_version %>
       t.timestamps
     end
 
-    add_index :llm_api_calls, :provider
-    add_index :llm_api_calls, :model
     add_index :llm_api_calls, :tracked_at
     add_index :llm_api_calls, [:provider, :tracked_at]
+    add_index :llm_api_calls, [:model, :tracked_at]
     add_index :llm_api_calls, :stream
     add_index :llm_api_calls, :usage_source
     add_index :llm_api_calls, :provider_response_id

data/lib/llm_cost_tracker/parsers/anthropic.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require "uri"
-
 require_relative "base"
 
 module LlmCostTracker
@@ -10,10 +8,7 @@ module LlmCostTracker
       HOSTS = %w[api.anthropic.com].freeze
 
       def match?(url)
-        uri = URI.parse(url.to_s)
-        HOSTS.include?(uri.host.to_s.downcase) && uri.path.include?("/v1/messages")
-      rescue URI::InvalidURIError
-        false
+        match_uri?(url, hosts: HOSTS, path_includes: "/v1/messages")
       end
 
       def provider_names
@@ -52,25 +47,25 @@ module LlmCostTracker
         usage = stream_usage(events)
         response_id = stream_response_id(events)
 
-        usage ? build_stream_result(model, usage, response_id) : build_unknown_stream_result(model, response_id)
+        if usage
+          build_stream_result(model, usage, response_id)
+        else
+          build_unknown_stream_usage(
+            provider: "anthropic",
+            model: model,
+            provider_response_id: response_id
+          )
+        end
       end
 
       private
 
       def stream_usage(events)
-        start_usage = nil
-        latest_delta = nil
-
-        events.each do |event|
-          data = event[:data]
-          next unless data.is_a?(Hash)
-
-          case data["type"]
-          when "message_start"
-            start_usage = data.dig("message", "usage")
-          when "message_delta"
-            latest_delta = data["usage"] if data["usage"].is_a?(Hash)
-          end
+        start_usage = find_event_value(events, reverse: true) do |data|
+          data.dig("message", "usage") if data["type"] == "message_start"
+        end
+        latest_delta = find_event_value(events, reverse: true) do |data|
+          data["usage"] if data["type"] == "message_delta" && data["usage"].is_a?(Hash)
         end
 
         return nil unless start_usage || latest_delta
@@ -81,25 +76,11 @@ module LlmCostTracker
       end
 
       def stream_model(events)
-        events.each do |event|
-          data = event[:data]
-          next unless data.is_a?(Hash)
-
-          model = data.dig("message", "model")
-          return model if model && !model.empty?
-        end
-        nil
+        find_event_value(events) { |data| data.dig("message", "model") }
       end
 
       def stream_response_id(events)
-        events.each do |event|
-          data = event[:data]
-          next unless data.is_a?(Hash)
-
-          id = data.dig("message", "id") || data["id"]
-          return id if id && !id.to_s.empty?
-        end
-        nil
+        find_event_value(events) { |data| data.dig("message", "id") || data["id"] }
       end
 
       def build_stream_result(model, usage, response_id)
@@ -121,19 +102,6 @@ module LlmCostTracker
           usage_source: :stream_final
         )
       end
-
-      def build_unknown_stream_result(model, response_id)
-        ParsedUsage.build(
-          provider: "anthropic",
-          provider_response_id: response_id,
-          model: model,
-          input_tokens: 0,
-          output_tokens: 0,
-          total_tokens: 0,
-          stream: true,
-          usage_source: :unknown
-        )
-      end
     end
   end
 end

data/lib/llm_cost_tracker/parsers/base.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "json"
+require "uri"
 
 module LlmCostTracker
   module Parsers
@@ -40,6 +41,85 @@ module LlmCostTracker
       rescue JSON::ParserError
         {}
       end
+
+      def uri_matches?(url)
+        uri = parsed_uri(url)
+        uri ? yield(uri) : false
+      end
+
+      def match_uri?(url, hosts: nil, exact_paths: nil, path_includes: nil, path_suffixes: nil, path_pattern: nil)
+        uri_matches?(url) do |uri|
+          host_match = hosts.nil? || host_matches?(uri, hosts)
+          path_match = path_matches?(
+            uri,
+            exact_paths: exact_paths,
+            path_includes: path_includes,
+            path_suffixes: path_suffixes,
+            path_pattern: path_pattern
+          )
+          extra_match = block_given? ? yield(uri) : true
+
+          host_match && path_match && extra_match ? true : false
+        end
+      end
+
+      def parsed_uri(url)
+        URI.parse(url.to_s)
+      rescue URI::InvalidURIError
+        nil
+      end
+
+      def host_matches?(uri, hosts)
+        hosts.include?(uri.host.to_s.downcase)
+      end
+
+      def path_matches?(uri, exact_paths: nil, path_includes: nil, path_suffixes: nil, path_pattern: nil)
+        path = uri.path.to_s
+        matches = true
+
+        matches &&= exact_paths.include?(path) if exact_paths
+        matches &&= Array(path_includes).all? { |fragment| path.include?(fragment) } if path_includes
+        matches &&= path.match?(path_pattern) if path_pattern
+
+        matches &&= path_suffixes.any? { |suffix| path == suffix || path.end_with?(suffix) } if path_suffixes
+
+        matches
+      end
+
+      def each_event_data(events, reverse: false)
+        enumerator = reverse ? events.reverse_each : events.each
+
+        enumerator.each do |event|
+          data = event[:data]
+          yield data if data.is_a?(Hash)
+        end
+      end
+
+      def find_event_value(events, reverse: false)
+        each_event_data(events, reverse:) do |data|
+          value = yield(data)
+          return value if event_value_present?(value)
+        end
+
+        nil
+      end
+
+      def build_unknown_stream_usage(provider:, model:, provider_response_id:)
+        ParsedUsage.build(
+          provider: provider,
+          provider_response_id: provider_response_id,
+          model: model,
+          input_tokens: 0,
+          output_tokens: 0,
+          total_tokens: 0,
+          stream: true,
+          usage_source: :unknown
+        )
+      end
+
+      def event_value_present?(value)
+        !value.nil? && (!value.respond_to?(:empty?) || !value.empty?)
+      end
     end
   end
 end

data/lib/llm_cost_tracker/parsers/gemini.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require "uri"
-
 require_relative "base"
 
 module LlmCostTracker
@@ -12,10 +10,7 @@ module LlmCostTracker
       STREAM_PATH_PATTERN  = /:streamGenerateContent\z/
 
       def match?(url)
-        uri = URI.parse(url.to_s)
-        HOSTS.include?(uri.host.to_s.downcase) && uri.path.match?(TRACKED_PATH_PATTERN)
-      rescue URI::InvalidURIError
-        false
+        match_uri?(url, hosts: HOSTS, path_pattern: TRACKED_PATH_PATTERN)
       end
 
       def provider_names
@@ -48,6 +43,7 @@ module LlmCostTracker
 
         usage = merged_stream_usage(events)
         model = extract_model_from_url(request_url)
+        response_id = stream_response_id(events)
 
         if usage
           build_parsed_usage(
@@ -55,18 +51,13 @@ module LlmCostTracker
             usage,
             stream: true,
             usage_source: :stream_final,
-            provider_response_id: stream_response_id(events)
+            provider_response_id: response_id
           )
         else
-          ParsedUsage.build(
+          build_unknown_stream_usage(
             provider: "gemini",
-            provider_response_id: stream_response_id(events),
             model: model,
-            input_tokens: 0,
-            output_tokens: 0,
-            total_tokens: 0,
-            stream: true,
-            usage_source: :unknown
+            provider_response_id: response_id
           )
         end
       end
@@ -91,15 +82,10 @@ module LlmCostTracker
       end
 
       def merged_stream_usage(events)
-        latest = nil
-        events.each do |event|
-          data = event[:data]
-          next unless data.is_a?(Hash)
-
+        find_event_value(events, reverse: true) do |data|
           meta = data["usageMetadata"]
-          latest = meta if meta.is_a?(Hash)
+          meta if meta.is_a?(Hash)
         end
-        latest
       end
 
       def output_tokens(usage)
@@ -107,28 +93,19 @@ module LlmCostTracker
       end
 
       def stream_response_id(events)
-        events.each do |event|
-          data = event[:data]
-          next unless data.is_a?(Hash)
-
-          id = data["responseId"]
-          return id if id && !id.to_s.empty?
-        end
-        nil
+        find_event_value(events) { |data| data["responseId"] }
       end
 
       def streaming_url?(request_url)
-        URI.parse(request_url.to_s).path.match?(STREAM_PATH_PATTERN)
-      rescue URI::InvalidURIError
-        false
+        match_uri?(request_url, path_pattern: STREAM_PATH_PATTERN)
       end
 
       def extract_model_from_url(url)
-        uri = URI.parse(url.to_s)
+        uri = parsed_uri(url)
+        return nil unless uri
+
         match = uri.path.match(%r{/models/([^/:]+)})
         match && match[1]
-      rescue URI::InvalidURIError
-        nil
       end
     end
   end

data/lib/llm_cost_tracker/parsers/openai.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require "uri"
-
 require_relative "base"
 require_relative "openai_usage"
 
@@ -14,10 +12,7 @@ module LlmCostTracker
       TRACKED_PATHS = %w[/v1/chat/completions /v1/completions /v1/embeddings /v1/responses].freeze
 
       def match?(url)
-        uri = URI.parse(url.to_s)
-        HOSTS.include?(uri.host.to_s.downcase) && TRACKED_PATHS.include?(uri.path)
-      rescue URI::InvalidURIError
-        false
+        match_uri?(url, hosts: HOSTS, exact_paths: TRACKED_PATHS)
       end
 
       def provider_names

data/lib/llm_cost_tracker/parsers/openai_compatible.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require "uri"
-
 require_relative "base"
 require_relative "openai_usage"
 
@@ -13,10 +11,7 @@ module LlmCostTracker
       TRACKED_PATH_SUFFIXES = %w[/chat/completions /completions /embeddings /responses].freeze
 
       def match?(url)
-        uri = URI.parse(url.to_s)
-        !provider_for_host(uri.host).nil? && tracked_path?(uri.path)
-      rescue URI::InvalidURIError
-        false
+        match_uri?(url, path_suffixes: TRACKED_PATH_SUFFIXES) { |uri| provider_for_uri(uri) }
       end
 
       def provider_names
@@ -37,18 +32,14 @@ module LlmCostTracker
       private
 
       def provider_for(request_url)
-        uri = URI.parse(request_url.to_s)
-        provider_for_host(uri.host) || "openai_compatible"
-      rescue URI::InvalidURIError
-        "openai_compatible"
+        uri = parsed_uri(request_url)
+        provider_for_uri(uri) || "openai_compatible"
       end
 
-      def provider_for_host(host)
-        LlmCostTracker.configuration.openai_compatible_providers[host.to_s.downcase]&.to_s
-      end
+      def provider_for_uri(uri)
+        return nil unless uri
 
-      def tracked_path?(path)
-        TRACKED_PATH_SUFFIXES.any? { |suffix| path == suffix || path.end_with?(suffix) }
+        LlmCostTracker.configuration.openai_compatible_providers[uri.host.to_s.downcase]&.to_s
       end
     end
   end

data/lib/llm_cost_tracker/parsers/openai_usage.rb

@@ -34,12 +34,13 @@ module LlmCostTracker
         request = safe_json_parse(request_body)
         model = detect_stream_model(events) || request["model"]
         usage = detect_stream_usage(events)
+        response_id = detect_stream_response_id(events)
 
         if usage
           cache_read = cache_read_input_tokens(usage)
           ParsedUsage.build(
             provider: provider_for(request_url),
-            provider_response_id: detect_stream_response_id(events),
+            provider_response_id: response_id,
             model: model,
             input_tokens: regular_input_tokens(usage, cache_read),
             output_tokens: (usage["completion_tokens"] || usage["output_tokens"]).to_i,
@@ -50,50 +51,27 @@ module LlmCostTracker
             usage_source: :stream_final
           )
         else
-          ParsedUsage.build(
+          build_unknown_stream_usage(
             provider: provider_for(request_url),
-            provider_response_id: detect_stream_response_id(events),
             model: model,
-            input_tokens: 0,
-            output_tokens: 0,
-            total_tokens: 0,
-            stream: true,
-            usage_source: :unknown
+            provider_response_id: response_id
           )
         end
       end
 
       def detect_stream_usage(events)
-        events.reverse_each do |event|
-          data = event[:data]
-          next unless data.is_a?(Hash)
-
+        find_event_value(events, reverse: true) do |data|
           usage = data["usage"]
-          return usage if usage.is_a?(Hash) && !usage.empty?
+          usage if usage.is_a?(Hash)
         end
-        nil
       end
 
       def detect_stream_model(events)
-        events.each do |event|
-          data = event[:data]
-          next unless data.is_a?(Hash)
-
-          model = data["model"]
-          return model if model && !model.to_s.empty?
-        end
-        nil
+        find_event_value(events) { |data| data["model"] }
       end
 
       def detect_stream_response_id(events)
-        events.each do |event|
-          data = event[:data]
-          next unless data.is_a?(Hash)
-
-          id = data["id"] || data.dig("response", "id")
-          return id if id && !id.to_s.empty?
-        end
-        nil
+        find_event_value(events) { |data| data["id"] || data.dig("response", "id") }
       end
 
       def regular_input_tokens(usage, cache_read)

data/lib/llm_cost_tracker/parsers/registry.rb

@@ -13,10 +13,14 @@ module LlmCostTracker
         end
 
         def register(parser)
+          parser = coerce_parser(parser)
+
           MUTEX.synchronize do
             current = @parsers || default_parsers.freeze
             @parsers = ([parser] + current).freeze
           end
+
+          parser
         end
 
         def find_for(url)
@@ -24,8 +28,8 @@ module LlmCostTracker
         end
 
         def find_for_provider(provider)
-          provider_name = provider.to_s
-          parsers.find { |parser| parser.provider_names.include?(provider_name) }
+          provider_name = provider.to_s.downcase
+          parsers.find { |parser| provider_names_for(parser).include?(provider_name) }
         end
 
         def reset!
@@ -34,6 +38,17 @@ module LlmCostTracker
 
         private
 
+        def coerce_parser(parser)
+          return parser.new if parser.is_a?(Class) && parser <= Base
+          return parser if parser.is_a?(Base)
+
+          raise ArgumentError, "parser must be a LlmCostTracker::Parsers::Base instance or class"
+        end
+
+        def provider_names_for(parser)
+          Array(parser.provider_names).map { |name| name.to_s.downcase }
+        end
+
         def default_parsers
           [Openai.new, OpenaiCompatible.new, Anthropic.new, Gemini.new]
         end

data/lib/llm_cost_tracker/storage/active_record_rollups.rb

@@ -17,47 +17,67 @@ module LlmCostTracker
           return unless event.cost&.total_cost
           return unless period_totals_enabled?
 
-          PERIODS.each_key { |period| increment_period_total(period, event) }
+          model = period_total_model
+          model.upsert_all(
+            period_rows(event),
+            on_duplicate: total_upsert_sql(model),
+            record_timestamps: true,
+            unique_by: unique_by(model, %i[period period_start])
+          )
         end
 
         def monthly_total(time: Time.now.utc)
-          period_total(:monthly, time)
+          period_totals(%i[monthly], time: time).fetch(:monthly)
         end
 
         def daily_total(time: Time.now.utc)
-          period_total(:daily, time)
+          period_totals(%i[daily], time: time).fetch(:daily)
         end
 
-        private
+        def period_totals(periods, time: Time.now.utc)
+          periods = periods.map(&:to_sym).select { |period| PERIODS.key?(period) }
+          return {} if periods.empty?
 
-        def period_total(period, time)
           if period_totals_enabled?
-            period_total_model
-              .where(period: PERIODS.fetch(period), period_start: bucket_for(period, time))
-              .pick(:total_cost)
-              .to_f
+            rollup_period_totals(periods, time)
           else
-            LlmCostTracker::LlmApiCall
-              .where(tracked_at: range_start_for(period, time)..time)
-              .sum(:total_cost)
-              .to_f
+            periods.to_h { |period| [period, fallback_period_total(period, time)] }
           end
         end
 
-        def increment_period_total(period, event)
-          model = period_total_model
-          model.upsert_all(
-            [
-              {
-                period: PERIODS.fetch(period),
-                period_start: bucket_for(period, event.tracked_at),
-                total_cost: event.cost.total_cost
-              }
-            ],
-            on_duplicate: total_upsert_sql(model),
-            record_timestamps: true,
-            unique_by: unique_by(model, %i[period period_start])
-          )
+        private
+
+        def period_rows(event)
+          PERIODS.map do |period, name|
+            {
+              period: name,
+              period_start: bucket_for(period, event.tracked_at),
+              total_cost: event.cost.total_cost
+            }
+          end
+        end
+
+        def rollup_period_totals(periods, time)
+          buckets = periods.to_h { |period| [period, bucket_for(period, time)] }
+          index = buckets.to_h { |period, bucket| [[PERIODS.fetch(period), bucket], period] }
+          totals = periods.to_h { |period| [period, 0.0] }
+
+          period_total_model
+            .where(period: periods.map { |period| PERIODS.fetch(period) }, period_start: buckets.values)
+            .pluck(:period, :period_start, :total_cost)
+            .each do |name, start, total|
+              period = index[[name, start.to_date]]
+              totals[period] = total.to_f if period
+            end
+
+          totals
+        end
+
+        def fallback_period_total(period, time)
+          LlmCostTracker::LlmApiCall
+            .where(tracked_at: range_start_for(period, time)..time)
+            .sum(:total_cost)
+            .to_f
         end
 
         def period_totals_enabled?

data/lib/llm_cost_tracker/storage/active_record_store.rb

@@ -50,6 +50,10 @@ module LlmCostTracker
           ActiveRecordRollups.daily_total(time: time)
         end
 
+        def period_totals(periods, time: Time.now.utc)
+          ActiveRecordRollups.period_totals(periods, time: time)
+        end
+
         private
 
         def stringify_tags(tags)

data/lib/llm_cost_tracker/tags_column.rb

@@ -2,58 +2,96 @@
 
 module LlmCostTracker
   module TagsColumn
+    USAGE_BREAKDOWN_COLUMNS = %w[
+      cache_read_input_tokens
+      cache_write_input_tokens
+      hidden_output_tokens
+    ].freeze
+
+    USAGE_BREAKDOWN_COST_COLUMNS = %w[
+      cache_read_input_cost
+      cache_write_input_cost
+    ].freeze
+
+    def reset_column_information
+      remove_instance_variable(:@lct_schema_capabilities) if instance_variable_defined?(:@lct_schema_capabilities)
+
+      super
+    end
+
     def tags_json_column?
-      tags_jsonb_column? || tags_mysql_json_column?
+      capabilities = lct_schema_capabilities
+
+      capabilities.fetch(:tags_jsonb) || capabilities.fetch(:tags_mysql_json)
     end
 
     def tags_jsonb_column?
-      column = columns_hash["tags"]
-      return false unless column
-
-      column.type == :jsonb || column.sql_type.to_s.downcase == "jsonb"
+      lct_schema_capabilities.fetch(:tags_jsonb)
     end
 
     def tags_mysql_json_column?
-      column = columns_hash["tags"]
-      return false unless column
-      return false if tags_jsonb_column?
-
-      column.type == :json && connection.adapter_name.match?(/mysql/i)
+      lct_schema_capabilities.fetch(:tags_mysql_json)
     end
 
     def latency_column?
-      columns_hash.key?("latency_ms")
+      lct_schema_capabilities.fetch(:latency)
     end
 
     def stream_column?
-      columns_hash.key?("stream")
+      lct_schema_capabilities.fetch(:stream)
     end
 
     def usage_source_column?
-      columns_hash.key?("usage_source")
+      lct_schema_capabilities.fetch(:usage_source)
     end
 
     def provider_response_id_column?
-      columns_hash.key?("provider_response_id")
+      lct_schema_capabilities.fetch(:provider_response_id)
     end
 
     def pricing_mode_column?
-      columns_hash.key?("pricing_mode")
+      lct_schema_capabilities.fetch(:pricing_mode)
     end
 
     def usage_breakdown_columns?
-      %w[
-        cache_read_input_tokens
-        cache_write_input_tokens
-        hidden_output_tokens
-      ].all? { |column| columns_hash.key?(column) }
+      lct_schema_capabilities.fetch(:usage_breakdown)
     end
 
     def usage_breakdown_cost_columns?
-      %w[
-        cache_read_input_cost
-        cache_write_input_cost
-      ].all? { |column| columns_hash.key?(column) }
+      lct_schema_capabilities.fetch(:usage_breakdown_cost)
+    end
+
+    private
+
+    def lct_schema_capabilities
+      columns = columns_hash
+      adapter_name = connection.adapter_name
+      cache = @lct_schema_capabilities
+
+      return cache.fetch(:values) if cache && cache.fetch(:columns).equal?(columns) &&
+                                     cache.fetch(:adapter_name) == adapter_name
+
+      values = build_lct_schema_capabilities(columns, adapter_name)
+      @lct_schema_capabilities = { columns: columns, adapter_name: adapter_name, values: values }
+      values
+    end
+
+    def build_lct_schema_capabilities(columns, adapter_name)
+      tag_column = columns["tags"]
+      tags_jsonb = tag_column && (tag_column.type == :jsonb || tag_column.sql_type.to_s.downcase == "jsonb")
+      tags_mysql_json = tag_column && !tags_jsonb && tag_column.type == :json && adapter_name.match?(/mysql/i)
+
+      {
+        tags_jsonb: tags_jsonb ? true : false,
+        tags_mysql_json: tags_mysql_json ? true : false,
+        latency: columns.key?("latency_ms"),
+        stream: columns.key?("stream"),
+        usage_source: columns.key?("usage_source"),
+        provider_response_id: columns.key?("provider_response_id"),
+        pricing_mode: columns.key?("pricing_mode"),
+        usage_breakdown: USAGE_BREAKDOWN_COLUMNS.all? { |column| columns.key?(column) },
+        usage_breakdown_cost: USAGE_BREAKDOWN_COST_COLUMNS.all? { |column| columns.key?(column) }
+      }
     end
   end
 end

data/lib/llm_cost_tracker/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module LlmCostTracker
-  VERSION = "0.4.0"
+  VERSION = "0.4.1"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: llm_cost_tracker
 version: !ruby/object:Gem::Version
-  version: 0.4.0
+  version: 0.4.1
 platform: ruby
 authors:
 - Sergii Khomenko
@@ -240,6 +240,7 @@ files:
 - app/helpers/llm_cost_tracker/dashboard_query_helper.rb
 - app/helpers/llm_cost_tracker/pagination_helper.rb
 - app/services/llm_cost_tracker/dashboard/data_quality.rb
+- app/services/llm_cost_tracker/dashboard/data_quality_aggregate.rb
 - app/services/llm_cost_tracker/dashboard/filter.rb
 - app/services/llm_cost_tracker/dashboard/overview_stats.rb
 - app/services/llm_cost_tracker/dashboard/provider_breakdown.rb