llm_cost_tracker 0.7.0 → 0.7.1

data/lib/llm_cost_tracker/{storage/active_record_periods.rb → ledger/period.rb}

@@ -1,8 +1,8 @@
 # frozen_string_literal: true
 
 module LlmCostTracker
-  module Storage
-    module ActiveRecordPeriods
+  module Ledger
+    module Period
       PERIODS = {
         monthly: "month",
         daily:   "day"

data/lib/llm_cost_tracker/ledger/rollups/batch.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+require_relative "../period"
+
+module LlmCostTracker
+  module Ledger
+    class Rollups
+      class Batch
+        def self.rows(events)
+          new(events).rows
+        end
+
+        def initialize(events)
+          @events = events
+        end
+
+        def rows
+          totals.map do |(period, period_start), total_cost|
+            {
+              period: period,
+              period_start: period_start,
+              total_cost: total_cost
+            }
+          end
+        end
+
+        private
+
+        attr_reader :events
+
+        def totals
+          events.each_with_object(Hash.new { |hash, key| hash[key] = BigDecimal("0") }) do |event, rows|
+            Period::PERIODS.each do |period, name|
+              rows[[name, Period.bucket(period, event.tracked_at)]] += BigDecimal(event.total_cost.to_s)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/rollups/upsert_sql.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require_relative "../schema/adapter"
+
+module LlmCostTracker
+  module Ledger
+    class Rollups
+      class UpsertSql
+        def self.call(model)
+          new(model).call
+        end
+
+        def initialize(model)
+          @model = model
+        end
+
+        def call
+          return Arel.sql(mysql_sql) if Ledger::Schema::Adapter.mysql?(connection)
+          return Arel.sql(postgres_sql) if Ledger::Schema::Adapter.postgresql?(connection)
+
+          Ledger::Schema::Adapter.ensure_supported!(connection)
+        end
+
+        private
+
+        attr_reader :model
+
+        def postgres_sql
+          total_cost = connection.quote_column_name("total_cost")
+          updated_at = connection.quote_column_name("updated_at")
+
+          "#{total_cost} = #{model.quoted_table_name}.#{total_cost} + excluded.#{total_cost}, " \
+            "#{updated_at} = excluded.#{updated_at}"
+        end
+
+        def mysql_sql
+          "total_cost = total_cost + VALUES(total_cost), updated_at = VALUES(updated_at)"
+        end
+
+        def connection
+          model.connection
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+require "bigdecimal"
+
+require_relative "period"
+require_relative "rollups/batch"
+require_relative "rollups/upsert_sql"
+
+module LlmCostTracker
+  module Ledger
+    class Rollups
+      class << self
+        def increment!(event)
+          return unless event.total_cost
+
+          Period::Total.upsert_all(
+            period_rows(event),
+            on_duplicate: Ledger::Rollups::UpsertSql.call(Period::Total),
+            record_timestamps: true,
+            unique_by: unique_by(Period::Total, %i[period period_start])
+          )
+        end
+
+        def increment_many!(events)
+          events = Array(events).select(&:total_cost)
+          return if events.empty?
+
+          Period::Total.upsert_all(
+            Ledger::Rollups::Batch.rows(events),
+            on_duplicate: Ledger::Rollups::UpsertSql.call(Period::Total),
+            record_timestamps: true,
+            unique_by: unique_by(Period::Total, %i[period period_start])
+          )
+        end
+
+        def decrement!(call_rows)
+          totals = period_decrement_totals(call_rows)
+          return if totals.empty?
+
+          apply_decrements(totals)
+        end
+
+        private
+
+        def period_rows(event)
+          Period::PERIODS.map do |period, name|
+            {
+              period: name,
+              period_start: Period.bucket(period, event.tracked_at),
+              total_cost: event.total_cost
+            }
+          end
+        end
+
+        def period_decrement_totals(call_rows)
+          call_rows.each_with_object(Hash.new { |totals, key| totals[key] = BigDecimal("0") }) do |row, totals|
+            _id, tracked_at, total_cost = row
+            next unless total_cost
+
+            Period::PERIODS.each_key do |period|
+              totals[[period, Period.bucket(period, tracked_at)]] += BigDecimal(total_cost.to_s)
+            end
+          end
+        end
+
+        def apply_decrements(totals)
+          now = Time.now.utc
+
+          totals.each do |(period, period_start), amount|
+            row = Period::Total.lock.find_by(period: Period::PERIODS.fetch(period),
+                                             period_start: period_start)
+            next unless row
+
+            row.update_columns(total_cost: [BigDecimal(row.total_cost.to_s) - amount, BigDecimal("0")].max,
+                               updated_at: now)
+          end
+        end
+
+        def unique_by(model, column)
+          return unless model.connection.supports_insert_conflict_target?
+
+          column
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/schema/adapter.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require_relative "../../errors"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module Adapter
+        MYSQL_ADAPTERS = %w[
+          ActiveRecord::ConnectionAdapters::AbstractMysqlAdapter
+          ActiveRecord::ConnectionAdapters::Mysql2Adapter
+          ActiveRecord::ConnectionAdapters::TrilogyAdapter
+        ].freeze
+        POSTGRESQL_ADAPTERS = %w[
+          ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
+        ].freeze
+        MYSQL_PATTERN = /mysql|trilogy|mariadb/i
+        POSTGRESQL_PATTERN = /postgres/i
+
+        class << self
+          def mysql?(value)
+            adapter_instance?(value, MYSQL_ADAPTERS) || adapter_name(value).match?(MYSQL_PATTERN)
+          end
+
+          def postgresql?(value)
+            adapter_instance?(value, POSTGRESQL_ADAPTERS) || adapter_name(value).match?(POSTGRESQL_PATTERN)
+          end
+
+          def ensure_supported!(value)
+            return if mysql?(value) || postgresql?(value)
+
+            raise Error, "Unsupported database adapter: #{adapter_name(value)}. Use PostgreSQL or MySQL."
+          end
+
+          private
+
+          def adapter_instance?(value, class_names)
+            class_names.any? do |class_name|
+              adapter_class = class_name.safe_constantize
+              adapter_class && value.is_a?(adapter_class)
+            end
+          end
+
+          def adapter_name(value)
+            value.try(:adapter_name).presence || value.to_s
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/schema/calls.rb

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+require "llm_cost_tracker/ledger/schema/adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module Calls
+        CURRENT_SCHEMA_COLUMNS = %w[
+          event_id
+          provider
+          model
+          input_tokens
+          output_tokens
+          total_tokens
+          cache_read_input_tokens
+          cache_write_input_tokens
+          cache_write_1h_input_tokens
+          hidden_output_tokens
+          input_cost
+          output_cost
+          total_cost
+          cache_read_input_cost
+          cache_write_input_cost
+          cache_write_1h_input_cost
+          latency_ms
+          stream
+          usage_source
+          provider_response_id
+          pricing_mode
+          tags
+          tracked_at
+        ].freeze
+
+        class << self
+          def current_schema?
+            current_schema_errors.empty?
+          end
+
+          def current_schema_errors
+            schema_capabilities.fetch(:current_schema_errors)
+          end
+
+          def missing_current_schema_columns
+            schema_capabilities.fetch(:missing_current_schema_columns)
+          end
+
+          private
+
+          def schema_capabilities
+            columns = Ledger::Call.columns_hash
+            adapter_name = Ledger::Call.connection.adapter_name
+            cache = @schema_capabilities
+
+            return cache.fetch(:values) if cache && cache.fetch(:columns).equal?(columns) &&
+                                           cache.fetch(:adapter_name) == adapter_name
+
+            values = build_schema_capabilities(columns, adapter_name)
+            @schema_capabilities = { columns: columns, adapter_name: adapter_name, values: values }
+            values
+          end
+
+          def build_schema_capabilities(columns, adapter_name)
+            Ledger::Schema::Adapter.ensure_supported!(adapter_name)
+
+            {
+              missing_current_schema_columns: missing_columns_for(columns),
+              current_schema_errors: schema_errors_for(columns, adapter_name)
+            }
+          end
+
+          def schema_errors_for(columns, adapter_name)
+            errors = []
+            missing = missing_columns_for(columns)
+            errors << "missing columns: #{missing.join(', ')}" if missing.any?
+
+            tag_column = columns["tags"]
+            if tag_column
+              postgresql = Ledger::Schema::Adapter.postgresql?(adapter_name)
+              expected_type = postgresql ? "jsonb" : "json"
+              valid_type =
+                if postgresql
+                  tag_column.type == :jsonb || tag_column.sql_type.to_s.downcase == "jsonb"
+                else
+                  tag_column.type == :json
+                end
+
+              errors << "tags column must use #{expected_type}" unless valid_type
+            end
+
+            errors
+          end
+
+          def missing_columns_for(columns)
+            CURRENT_SCHEMA_COLUMNS - columns.keys
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/schema/period_totals.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module PeriodTotals
+        REQUIRED_COLUMNS = %w[period period_start total_cost].freeze
+        UNIQUE_COLUMNS = %i[period period_start].freeze
+
+        class << self
+          def current_schema_errors
+            connection = Ledger::Call.connection
+            table_name = Ledger::Period::Total.table_name
+            return ["llm_cost_tracker_period_totals table is missing"] unless connection.data_source_exists?(table_name)
+
+            errors = []
+            missing = REQUIRED_COLUMNS - Ledger::Period::Total.columns_hash.keys
+            errors << "missing columns: #{missing.join(', ')}" if missing.any?
+            errors << "missing unique index: period, period_start" unless unique_period_index?(connection, table_name)
+            errors
+          end
+
+          private
+
+          def unique_period_index?(connection, table_name)
+            connection.index_exists?(table_name, UNIQUE_COLUMNS, unique: true)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/store.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+require_relative "../pricing"
+require_relative "rollups"
+
+module LlmCostTracker
+  module Ledger
+    class Store
+      class << self
+        def insert_many(events)
+          events = Array(events)
+          return [] if events.empty?
+
+          model = LlmCostTracker::Ledger::Call
+          insertable = new_events(model, events)
+
+          if insertable.any?
+            rows = insertable.map { |event| attributes_for(event) }
+            model.insert_all!(rows, record_timestamps: true, returning: false)
+            Ledger::Rollups.increment_many!(insertable)
+          end
+          events
+        end
+
+        private
+
+        def attributes_for(event)
+          tags = (event.tags || {}).transform_keys(&:to_s).transform_values { |value| stringify_tag_value(value) }
+          usage = event.token_usage.stored_attributes
+
+          attributes = {
+            event_id: event.event_id,
+            provider: event.provider,
+            model: event.model,
+            tags: tags,
+            tracked_at: event.tracked_at,
+            pricing_mode: event.pricing_mode,
+            latency_ms: event.latency_ms,
+            stream: event.stream,
+            usage_source: event.usage_source,
+            provider_response_id: event.provider_response_id
+          }
+
+          attributes.merge(usage).merge(Pricing.stored_cost_attributes(event.cost || {}))
+        end
+
+        def new_events(model, events)
+          existing_ids = model.where(event_id: events.map(&:event_id)).pluck(:event_id).to_set
+          events.reject { |event| existing_ids.include?(event.event_id) }
+        end
+
+        def stringify_tag_value(value)
+          return value.transform_values { |nested| stringify_tag_value(nested) } if value.is_a?(Hash)
+
+          value.to_s
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/tags/query.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require "json"
+
+require_relative "../schema/adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Tags
+      module Query
+        class << self
+          def apply(model, tags)
+            normalized_tags = (tags || {}).to_h.transform_keys(&:to_s).transform_values(&:to_s)
+            return model.all if normalized_tags.empty?
+
+            connection = model.connection
+            json = normalized_tags.to_json
+
+            if Schema::Adapter.postgresql?(connection)
+              model.where("tags @> ?::jsonb", json)
+            else
+              model.where("JSON_CONTAINS(tags, ?)", json)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/tags/sql.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require_relative "../schema/adapter"
+require_relative "../../tags/key"
+
+module LlmCostTracker
+  module Ledger
+    module Tags
+      module Sql
+        class << self
+          def value_expression(model, key, table_name:)
+            key = LlmCostTracker::Tags::Key.validate!(key)
+            column = "#{table_name}.#{model.connection.quote_column_name('tags')}"
+
+            if Ledger::Schema::Adapter.postgresql?(model.connection)
+              "#{column}->>#{model.connection.quote(key)}"
+            elsif Ledger::Schema::Adapter.mysql?(model.connection)
+              "JSON_UNQUOTE(JSON_EXTRACT(#{column}, #{model.connection.quote(json_path(key))}))"
+            else
+              Ledger::Schema::Adapter.ensure_supported!(model.connection)
+            end
+          end
+
+          private
+
+          def json_path(key)
+            "$.\"#{key}\""
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+require_relative "ledger/schema/adapter"
+require_relative "ledger/schema/calls"
+require_relative "ledger/schema/period_totals"
+require_relative "ledger/tags/query"
+require_relative "ledger/tags/sql"
+require_relative "ledger/period"
+require_relative "ledger/rollups/batch"
+require_relative "ledger/rollups/upsert_sql"
+require_relative "ledger/rollups"
+require_relative "ledger/store"
+require_relative "ledger/period/totals"

data/lib/llm_cost_tracker/logging.rb

@@ -19,9 +19,10 @@ module LlmCostTracker
 
       def log(level, message)
         message = prefixed(message)
+        logger = Rails.logger
 
-        if rails_logger
-          rails_logger.public_send(level, message)
+        if logger
+          logger.try(level, message)
         else
           Kernel.warn(message)
         end
@@ -35,10 +36,6 @@ module LlmCostTracker
 
         "#{PREFIX} #{message}"
       end
-
-      def rails_logger
-        Rails.logger if defined?(Rails) && Rails.respond_to?(:logger) && Rails.logger
-      end
     end
   end
 end

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -2,10 +2,10 @@
 
 require "faraday"
 require "json"
+require "uri"
 
 require_relative "../logging"
-require_relative "../request_url"
-require_relative "../stream_capture"
+require_relative "../capture/stream"
 
 module LlmCostTracker
   module Middleware
@@ -20,12 +20,12 @@ module LlmCostTracker
 
         request_url  = request_env.url.to_s
         request_body = read_body(request_env.body) || ""
-        parser       = Parsers::Registry.find_for(request_url)
+        parser       = Parsers.find_for(request_url)
         streaming    = parser&.streaming_request?(request_url, request_body)
         stream_buffer = install_stream_tap(request_env) if streaming
 
         Tracker.enforce_budget! if parser
-        started_at = monotonic_time
+        started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
 
         @app.call(request_env).on_complete do |response_env|
           process(
@@ -34,7 +34,7 @@ module LlmCostTracker
             request_url: request_url,
             request_body: request_body,
             response_env: response_env,
-            latency_ms: elapsed_ms(started_at),
+            latency_ms: ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at) * 1000).round,
             streaming: streaming,
             stream_buffer: stream_buffer
           )
@@ -56,15 +56,9 @@ module LlmCostTracker
         return unless parsed
 
         Tracker.record(
-          provider: parsed.provider,
-          model: parsed.model,
-          input_tokens: parsed.input_tokens,
-          output_tokens: parsed.output_tokens,
+          capture: parsed,
           latency_ms: latency_ms,
-          stream: parsed.stream,
-          usage_source: parsed.usage_source,
-          provider_response_id: parsed.provider_response_id,
-          metadata: resolved_tags(request_env).merge(parsed.metadata)
+          metadata: resolved_tags(request_env)
         )
       rescue LlmCostTracker::Error
         raise
@@ -76,7 +70,7 @@ module LlmCostTracker
         response_body = read_body(response_env.body)
         unless response_body
           Logging.warn(
-            "Unable to read response body for #{RequestUrl.label(request_url)}; " \
+            "Unable to read response body for #{request_url_label(request_url)}; " \
             "known streaming responses are captured automatically, or via LlmCostTracker.track_stream " \
             "for custom clients."
           )
@@ -93,9 +87,9 @@ module LlmCostTracker
         end
 
         body = stream_buffer&.dig(:buffer)&.string
-        body = read_body(response_env.body) if body.nil? || body.empty?
+        body = read_body(response_env.body) if body.blank?
 
-        if body.nil? || body.empty?
+        if body.blank?
           Logging.warn(capture_warning(request_url, stream_buffer))
           return parser.parse_stream(request_url, request_body, response_env.status, [])
         end
@@ -105,16 +99,17 @@ module LlmCostTracker
       end
 
       def install_stream_tap(request_env)
-        return nil unless request_env.respond_to?(:request) && request_env.request
+        request = request_env.try(:request)
+        return nil unless request
 
-        original = request_env.request.on_data
+        original = request.on_data
         return nil unless original
 
         state = { buffer: StringIO.new, bytes: 0, overflowed: false }
-        request_env.request.on_data = proc do |chunk, size, env|
+        request.on_data = proc do |chunk, size, env|
           chunk = chunk.to_s
           unless state[:overflowed]
-            if state[:bytes] + chunk.bytesize <= StreamCapture::LIMIT_BYTES
+            if state[:bytes] + chunk.bytesize <= Capture::Stream::LIMIT_BYTES
               state[:buffer] << chunk
               state[:bytes] += chunk.bytesize
             else
@@ -136,38 +131,42 @@ module LlmCostTracker
         when nil then ""
         when Hash, Array then body.to_json
         else
-          body.respond_to?(:to_str) ? body.to_str : nil
+          body.try(:to_str)
         end
       end
 
       def resolved_tags(request_env)
-        tags = @tags.respond_to?(:call) ? call_tags(request_env) : @tags
+        tags =
+          if @tags.respond_to?(:call)
+            @tags.arity.zero? ? @tags.call : @tags.call(request_env)
+          else
+            @tags
+          end
         return {} if tags.nil?
 
         tags.to_h
       end
 
-      def call_tags(request_env)
-        @tags.arity.zero? ? @tags.call : @tags.call(request_env)
-      end
-
-      def monotonic_time
-        Process.clock_gettime(Process::CLOCK_MONOTONIC)
-      end
-
-      def elapsed_ms(started_at)
-        ((monotonic_time - started_at) * 1000).round
-      end
-
       def capture_warning(request_url, stream_buffer)
         unless stream_buffer&.dig(:overflowed)
-          return "Unable to capture streaming response for #{RequestUrl.label(request_url)}; " \
+          return "Unable to capture streaming response for #{request_url_label(request_url)}; " \
                  "recording usage_source=unknown. Use LlmCostTracker.track_stream for manual capture."
         end
 
-        "Streaming response for #{RequestUrl.label(request_url)} exceeded #{StreamCapture::LIMIT_BYTES} bytes; " \
+        "Streaming response for #{request_url_label(request_url)} exceeded #{Capture::Stream::LIMIT_BYTES} bytes; " \
           "recording usage_source=unknown. Use LlmCostTracker.track_stream for manual capture."
       end
+
+      def request_url_label(value)
+        uri = URI.parse(value.to_s)
+        uri.query = nil
+        uri.fragment = nil
+        uri.try(:user=, nil)
+        uri.try(:password=, nil)
+        uri.to_s
+      rescue URI::InvalidURIError
+        value.to_s.split("?", 2).first
+      end
     end
   end
 end