llm_cost_tracker 0.9.0 → 0.11.0

data/lib/llm_cost_tracker/ledger/rollups.rb

@@ -34,8 +34,8 @@ module LlmCostTracker
         private
 
         def period_rows(event)
-          currency = currency_for(event)
-          provider = provider_for(event)
+          currency = currency_from_snapshot(event.pricing_snapshot)
+          provider = event.provider.to_s
           Period::PERIODS.map do |period, name|
             {
               period: name,
@@ -61,8 +61,8 @@ module LlmCostTracker
 
         def call_rollups(events)
           events.each_with_object(Hash.new { |totals, key| totals[key] = BigDecimal("0") }) do |event, totals|
-            currency = currency_for(event)
-            provider = provider_for(event)
+            currency = currency_from_snapshot(event.pricing_snapshot)
+            provider = event.provider.to_s
             Period::PERIODS.each do |period, name|
               key = [name, Period.bucket(period, event.tracked_at), currency, provider]
               totals[key] += BigDecimal(event.total_cost.to_s)
@@ -118,20 +118,11 @@ module LlmCostTracker
           end
         end
 
-        def currency_for(event)
-          snapshot = event.respond_to?(:pricing_snapshot) ? event.pricing_snapshot : nil
-          currency_from_snapshot(snapshot)
-        end
-
         def currency_from_snapshot(snapshot)
           value = (snapshot.is_a?(Hash) && (snapshot["currency"] || snapshot[:currency])) || DEFAULT_CURRENCY
           value.to_s.upcase
         end
 
-        def provider_for(event)
-          (event.respond_to?(:provider) ? event.provider : nil).to_s
-        end
-
         def upsert_call_rollups(rows)
           LlmCostTracker::CallRollup.upsert_all(
             rows,

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

@@ -42,6 +42,17 @@ module LlmCostTracker
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
             columns = LlmCostTracker::CallLineItem.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
+            errors
+          end
+
+          private
+
+          def compute_errors(connection, table_name, columns)
             errors = []
             missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?

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

@@ -16,8 +16,20 @@ module LlmCostTracker
             table_name = LlmCostTracker::CallRollup.table_name
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
+            columns = LlmCostTracker::CallRollup.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
+            errors
+          end
+
+          private
+
+          def compute_errors(connection, table_name, columns)
             errors = []
-            missing = REQUIRED_COLUMNS - LlmCostTracker::CallRollup.columns_hash.keys
+            missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?
             unless unique_period_index?(connection, table_name)
               errors << "missing unique index: period, period_start, currency, provider"
@@ -25,8 +37,6 @@ module LlmCostTracker
             errors
           end
 
-          private
-
           def unique_period_index?(connection, table_name)
             connection.index_exists?(table_name, UNIQUE_COLUMNS, unique: true)
           end

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

@@ -19,6 +19,17 @@ module LlmCostTracker
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
             columns = LlmCostTracker::CallTag.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
+            errors
+          end
+
+          private
+
+          def compute_errors(connection, table_name, columns)
             errors = []
             missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?

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

@@ -47,10 +47,6 @@ module LlmCostTracker
         private_constant :REQUIRED_INDEXES
 
         class << self
-          def current_schema?
-            current_schema_errors.empty?
-          end
-
           def current_schema_errors
             schema_capabilities.fetch(:current_schema_errors)
           end

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

@@ -28,15 +28,25 @@ module LlmCostTracker
             table_name = LlmCostTracker::Ingestion::InboxEntry.table_name
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
+            columns = LlmCostTracker::Ingestion::InboxEntry.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
+            errors
+          end
+
+          private
+
+          def compute_errors(connection, table_name, columns)
             errors = []
-            missing = REQUIRED_COLUMNS - LlmCostTracker::Ingestion::InboxEntry.columns_hash.keys
+            missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?
             errors << "missing unique index: event_id" unless event_id_unique_index?(connection, table_name)
             errors
           end
 
-          private
-
           def event_id_unique_index?(connection, table_name)
             connection.index_exists?(table_name, UNIQUE_COLUMNS, unique: true)
           end

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

@@ -23,15 +23,25 @@ module LlmCostTracker
             table_name = LlmCostTracker::Ingestion::Lease.table_name
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
+            columns = LlmCostTracker::Ingestion::Lease.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
+            errors
+          end
+
+          private
+
+          def compute_errors(connection, table_name, columns)
             errors = []
-            missing = REQUIRED_COLUMNS - LlmCostTracker::Ingestion::Lease.columns_hash.keys
+            missing = REQUIRED_COLUMNS - columns.keys
             errors << "missing columns: #{missing.join(', ')}" if missing.any?
             errors << "missing unique index: name" unless name_unique_index?(connection, table_name)
             errors
           end
 
-          private
-
           def name_unique_index?(connection, table_name)
             connection.index_exists?(table_name, UNIQUE_COLUMNS, unique: true)
           end

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

@@ -7,10 +7,10 @@ module LlmCostTracker
     module Schema
       module ProviderInvoiceImports
         REQUIRED_COLUMNS = %w[
-          source cursor window_start window_end state last_error
+          source provider cursor window_start window_end state last_error
           rows_imported started_at finished_at
         ].freeze
-        SOURCE_STARTED_AT_INDEX = %i[source started_at].freeze
+        SOURCE_PROVIDER_STARTED_AT_INDEX = %i[source provider started_at].freeze
 
         class << self
           def current_schema_errors
@@ -19,25 +19,35 @@ module LlmCostTracker
             table_name = LlmCostTracker::ProviderInvoiceImport.table_name
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
-            errors = []
-            errors.concat(column_errors)
-            errors.concat(index_errors(connection, table_name))
+            columns = LlmCostTracker::ProviderInvoiceImport.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
             errors
           end
 
           private
 
-          def column_errors
-            missing = REQUIRED_COLUMNS - LlmCostTracker::ProviderInvoiceImport.columns_hash.keys
+          def compute_errors(connection, table_name, columns)
+            errors = []
+            errors.concat(column_errors(columns))
+            errors.concat(index_errors(connection, table_name))
+            errors
+          end
+
+          def column_errors(columns)
+            missing = REQUIRED_COLUMNS - columns.keys
             return [] if missing.empty?
 
             ["missing columns: #{missing.join(', ')}"]
           end
 
           def index_errors(connection, table_name)
-            return [] if connection.index_exists?(table_name, SOURCE_STARTED_AT_INDEX)
+            return [] if connection.index_exists?(table_name, SOURCE_PROVIDER_STARTED_AT_INDEX)
 
-            ["missing index: source, started_at"]
+            ["missing index: source, provider, started_at"]
           end
         end
       end

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

@@ -19,27 +19,32 @@ module LlmCostTracker
             table_name = LlmCostTracker::ProviderInvoice.table_name
             return ["#{table_name} table is missing"] unless connection.data_source_exists?(table_name)
 
-            errors = []
-            errors.concat(column_errors)
-            errors.concat(metadata_type_errors(connection))
-            errors.concat(index_errors(connection, table_name))
+            columns = LlmCostTracker::ProviderInvoice.columns_hash
+            cache = @schema_capabilities
+            return cache.fetch(:errors) if cache && cache.fetch(:columns).equal?(columns)
+
+            errors = compute_errors(connection, table_name, columns)
+            @schema_capabilities = { columns: columns, errors: errors }
             errors
           end
 
           private
 
-          def column_errors
-            missing = REQUIRED_COLUMNS - LlmCostTracker::ProviderInvoice.columns_hash.keys
+          def compute_errors(connection, table_name, columns)
+            errors = []
+            errors.concat(column_errors(columns))
+            errors.concat(Adapter.json_column_errors(columns["metadata"], connection, "metadata"))
+            errors.concat(index_errors(connection, table_name))
+            errors
+          end
+
+          def column_errors(columns)
+            missing = REQUIRED_COLUMNS - columns.keys
             return [] if missing.empty?
 
             ["missing columns: #{missing.join(', ')}"]
           end
 
-          def metadata_type_errors(connection)
-            metadata = LlmCostTracker::ProviderInvoice.columns_hash["metadata"]
-            Adapter.json_column_errors(metadata, connection, "metadata")
-          end
-
           def index_errors(connection, table_name)
             errors = []
             unless connection.index_exists?(table_name, UNIQUE_INDEX_COLUMNS, unique: true)
@@ -48,8 +53,18 @@ module LlmCostTracker
             unless connection.index_exists?(table_name, SOURCE_PERIOD_INDEX_COLUMNS)
               errors << "missing index: source, currency, period_start"
             end
+            if Adapter.postgresql?(connection) && !gin_metadata_index?(connection, table_name)
+              errors << "missing GIN index on metadata " \
+                        "(run bin/rails generate llm_cost_tracker:upgrade_provider_invoices_metadata_index)"
+            end
             errors
           end
+
+          def gin_metadata_index?(connection, table_name)
+            connection.indexes(table_name).any? do |index|
+              index.columns == ["metadata"] && index.using.to_s == "gin"
+            end
+          end
         end
       end
     end

data/lib/llm_cost_tracker/ledger/store.rb

@@ -11,27 +11,34 @@ module LlmCostTracker
   module Ledger
     class Store
       class << self
-        def insert_many(events)
+        def insert(events, skip_existence_check: false)
           events = Array(events)
-          return [] if events.empty?
+          return if events.empty?
 
-          insertable = insertable_events(events)
+          insertable = skip_existence_check ? events : insertable_events(events)
+          return unless insertable.any?
 
-          if insertable.any?
-            LlmCostTracker::Call.transaction do
-              rows = insertable.map { |event| attributes_for(event) }
-              LlmCostTracker::Call.insert_all!(rows, record_timestamps: true, returning: false)
-              call_ids = call_ids_for(insertable)
-              insert_line_items(insertable, call_ids)
-              insert_call_tags(insertable, call_ids)
-            end
-            increment_rollups_safely(insertable) if LlmCostTracker.configuration.cache_rollups
+          LlmCostTracker::Call.transaction do
+            rows = insertable.map { |event| attributes_for(event) }
+            call_ids = insert_calls_returning_ids(rows, insertable)
+            insert_line_items(insertable, call_ids)
+            insert_call_tags(insertable, call_ids)
           end
-          events
+          increment_rollups_safely(insertable) if LlmCostTracker.configuration.cache_rollups
         end
 
         private
 
+        def insert_calls_returning_ids(rows, insertable)
+          if LlmCostTracker::Call.connection.supports_insert_returning?
+            result = LlmCostTracker::Call.insert_all!(rows, record_timestamps: true, returning: %i[id event_id])
+            result.rows.to_h { |id, event_id| [event_id, id] }
+          else
+            LlmCostTracker::Call.insert_all!(rows, record_timestamps: true, returning: false)
+            call_ids_for(insertable)
+          end
+        end
+
         def attributes_for(event)
           attributes = {
             event_id: event.event_id,
@@ -110,7 +117,7 @@ module LlmCostTracker
               {
                 llm_cost_tracker_call_id: call_ids.fetch(event.event_id),
                 key: key.to_s,
-                value: tag_row_value(value)
+                value: Tags::Encoding.encode(value)
               }
             end
           end
@@ -119,10 +126,6 @@ module LlmCostTracker
           LlmCostTracker::CallTag.insert_all!(rows, record_timestamps: false, returning: false)
         end
 
-        def tag_row_value(value)
-          Tags::Encoding.encode(value)
-        end
-
         def stored_details(details)
           (details || {}).transform_keys(&:to_s).transform_values { |value| Tags::Encoding.normalize_value(value) }
         end

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

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-require_relative "../schema/adapter"
 require_relative "encoding"
 
 module LlmCostTracker

data/lib/llm_cost_tracker/ledger.rb

@@ -14,3 +14,16 @@ require_relative "ledger/rollups/upsert_sql"
 require_relative "ledger/rollups"
 require_relative "ledger/store"
 require_relative "ledger/period/totals"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      CORE_SCHEMAS = [
+        [Calls, "llm_cost_tracker_calls"],
+        [CallLineItems, "llm_cost_tracker_call_line_items"],
+        [CallTags, "llm_cost_tracker_call_tags"]
+      ].freeze
+      CACHE_ROLLUPS_SCHEMA = [CallRollups, "llm_cost_tracker_call_rollups"].freeze
+    end
+  end
+end

data/lib/llm_cost_tracker/logging.rb

@@ -9,10 +9,6 @@ module LlmCostTracker
         log(:debug, message)
       end
 
-      def info(message)
-        log(:info, message)
-      end
-
       def warn(message)
         log(:warn, message)
       end

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -6,6 +6,7 @@ require "stringio"
 require "uri"
 
 require_relative "../logging"
+require_relative "../capture/sse"
 require_relative "../capture/stream"
 require_relative "../timing"
 
@@ -18,16 +19,25 @@ module LlmCostTracker
       end
 
       def call(request_env)
-        return @app.call(request_env) unless LlmCostTracker.configuration.enabled
+        return @app.call(request_env) unless enabled?
 
         request_url  = request_env.url.to_s
         request_body = read_body(request_env.body)
         parser       = Parsers.find_for(request_url)
-        streaming    = parser&.streaming_request?(request_url, request_body)
-        request_body = inject_stream_usage_flag(request_env, parser, request_url) if streaming
+        request_parsed = parser ? safe_json_parse(request_body) : nil
+        streaming = parser&.streaming_request?(request_url, request_parsed)
+        if streaming
+          request_body = inject_stream_usage_flag(request_env, parser, request_url, request_parsed) || request_body
+        end
         stream_buffer = install_stream_tap(request_env) if streaming
 
-        Tracker.enforce_budget! if parser
+        if parser
+          Tracker.enforce_budget!(
+            provider: parser.provider_for(request_url),
+            model: parser.model_for(request_url, request_parsed),
+            request: request_parsed
+          )
+        end
         context_tags, metadata = tag_snapshot(request_env) if parser
         started_at = LlmCostTracker::Timing.now_monotonic
 
@@ -40,6 +50,26 @@ module LlmCostTracker
 
       private
 
+      def enabled?
+        return @enabled if defined?(@enabled)
+
+        @enabled = LlmCostTracker.configuration.enabled
+      end
+
+      def safe_json_parse(body)
+        return {} if body.nil? || body.empty?
+
+        JSON.parse(body)
+      rescue JSON::ParserError
+        {}
+      end
+
+      def auto_enable_stream_usage?
+        return @auto_enable_stream_usage if defined?(@auto_enable_stream_usage)
+
+        @auto_enable_stream_usage = LlmCostTracker.configuration.auto_enable_stream_usage
+      end
+
       def invoke_app_with_capture(request_env:, parser:, request_url:, request_body:, streaming:,
                                   stream_buffer:, context_tags:, metadata:, started_at:)
         response_received = false
@@ -63,16 +93,15 @@ module LlmCostTracker
         raise
       end
 
-      def inject_stream_usage_flag(request_env, parser, request_url)
-        body_string = read_body(request_env.body)
-        return body_string unless LlmCostTracker.configuration.auto_enable_stream_usage
-        return body_string unless parser&.auto_enable_stream_usage?(request_url)
+      def inject_stream_usage_flag(request_env, parser, request_url, request_parsed)
+        return nil unless auto_enable_stream_usage?
+        return nil unless parser&.auto_enable_stream_usage?(request_url)
 
-        body = JSON.parse(body_string)
-        return body_string if body["stream_options"].is_a?(Hash) && body["stream_options"].key?("include_usage")
+        stream_options = request_parsed["stream_options"]
+        return nil if stream_options.is_a?(Hash) && stream_options.key?("include_usage")
 
-        body["stream_options"] = (body["stream_options"] || {}).merge("include_usage" => true)
-        new_body = body.to_json
+        request_parsed["stream_options"] = (stream_options || {}).merge("include_usage" => true)
+        new_body = request_parsed.to_json
         request_env.body = new_body
         new_body
       end
@@ -80,9 +109,9 @@ module LlmCostTracker
       def process_interrupted_stream(parser:, request_url:, request_body:, latency_ms:,
                                      context_tags:, metadata:, error:)
         request = parser.safe_json_parse(request_body)
-        capture = UsageCapture.build(
+        event = Event.build(
           provider: parser.provider_for(request_url),
-          model: request["model"] || UsageCapture::UNKNOWN_MODEL,
+          model: request["model"] || Event::UNKNOWN_MODEL,
           token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
           usage_source: :unknown
@@ -92,7 +121,7 @@ module LlmCostTracker
           stream_interrupted_error: "#{error.class}: #{error.message}"
         )
         Tracker.record(
-          capture: capture,
+          event: event,
           latency_ms: latency_ms,
           metadata: merged_metadata,
           context_tags: context_tags
@@ -125,7 +154,7 @@ module LlmCostTracker
         return unless parsed
 
         Tracker.record(
-          capture: parsed,
+          event: parsed,
           latency_ms: latency_ms,
           metadata: metadata,
           context_tags: context_tags
@@ -173,7 +202,7 @@ module LlmCostTracker
           )
         end
 
-        events = overflowed ? [] : Parsers::SSE.parse(body)
+        events = overflowed ? [] : Capture::SSE.parse(body)
         parser.parse_stream(
           request_url: request_url,
           request_body: request_body,