llm_cost_tracker 0.8.0 → 0.10.0

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

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require_relative "adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module IngestionLeases
+        REQUIRED_COLUMNS = %w[
+          name
+          locked_by
+          locked_until
+          created_at
+          updated_at
+        ].freeze
+
+        UNIQUE_COLUMNS = %i[name].freeze
+
+        class << self
+          def current_schema_errors
+            connection = LlmCostTracker::Ingestion::Lease.connection
+            Adapter.ensure_supported!(connection)
+            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 - columns.keys
+            errors << "missing columns: #{missing.join(', ')}" if missing.any?
+            errors << "missing unique index: name" unless name_unique_index?(connection, table_name)
+            errors
+          end
+
+          def name_unique_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/schema/provider_invoice_imports.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require_relative "adapter"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module ProviderInvoiceImports
+        REQUIRED_COLUMNS = %w[
+          source provider cursor window_start window_end state last_error
+          rows_imported started_at finished_at
+        ].freeze
+        SOURCE_PROVIDER_STARTED_AT_INDEX = %i[source provider started_at].freeze
+
+        class << self
+          def current_schema_errors
+            connection = LlmCostTracker::Call.connection
+            Adapter.ensure_supported!(connection)
+            table_name = LlmCostTracker::ProviderInvoiceImport.table_name
+            return ["#{table_name} table is missing"] unless connection.data_source_exists?(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 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_PROVIDER_STARTED_AT_INDEX)
+
+            ["missing index: source, provider, started_at"]
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -10,7 +10,7 @@ module LlmCostTracker
           source period_start period_end external_id billed_amount currency metadata imported_at
         ].freeze
         UNIQUE_INDEX_COLUMNS = %i[external_id].freeze
-        SOURCE_PERIOD_INDEX_COLUMNS = %i[source period_start].freeze
+        SOURCE_PERIOD_INDEX_COLUMNS = %i[source currency period_start].freeze
 
         class << self
           def current_schema_errors
@@ -19,37 +19,52 @@ 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)
               errors << "missing unique index: external_id"
             end
             unless connection.index_exists?(table_name, SOURCE_PERIOD_INDEX_COLUMNS)
-              errors << "missing index: source, period_start"
+              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

@@ -5,32 +5,40 @@ require "json"
 require_relative "../pricing"
 require_relative "../billing/line_item"
 require_relative "rollups"
+require_relative "tags/encoding"
 
 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?
-
-          insertable = insertable_events(events)
-
-          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)
-              Ledger::Rollups.increment_many!(insertable)
-            end
+          return if events.empty?
+
+          insertable = skip_existence_check ? events : insertable_events(events)
+          return unless insertable.any?
+
+          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,
@@ -109,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
@@ -118,15 +126,18 @@ module LlmCostTracker
           LlmCostTracker::CallTag.insert_all!(rows, record_timestamps: false, returning: false)
         end
 
-        def tag_row_value(value)
-          case value
-          when Hash, Array then JSON.generate(stored_tag_value(value))
-          else value.to_s
-          end
+        def stored_details(details)
+          (details || {}).transform_keys(&:to_s).transform_values { |value| Tags::Encoding.normalize_value(value) }
         end
 
-        def stored_details(details)
-          (details || {}).transform_keys(&:to_s).transform_values { |value| stored_tag_value(value) }
+        def increment_rollups_safely(events)
+          Ledger::Rollups.increment_many!(events)
+        rescue StandardError => e
+          raise if LlmCostTracker::Call.connection.open_transactions.positive?
+
+          LlmCostTracker::Logging.warn(
+            "Rollup increment failed for #{events.size} events after ledger commit: #{e.class}: #{e.message}"
+          )
         end
 
         def insertable_events(events)
@@ -138,14 +149,6 @@ module LlmCostTracker
             !existing_ids.include?(event_id) && seen_ids.add?(event_id)
           end
         end
-
-        def stored_tag_value(value)
-          if value.is_a?(Hash)
-            return value.transform_keys(&:to_s).transform_values { |nested| stored_tag_value(nested) }
-          end
-
-          value.to_s
-        end
       end
     end
   end

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

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+require "json"
+
+module LlmCostTracker
+  module Ledger
+    module Tags
+      module Encoding
+        module_function
+
+        def encode(value)
+          case value
+          when Hash then JSON.generate(normalize_hash(value))
+          when Array then JSON.generate(normalize_array(value))
+          else value.to_s
+          end
+        end
+
+        def normalize_hash(hash)
+          hash.transform_keys(&:to_s).sort.to_h.transform_values { |v| normalize_value(v) }
+        end
+
+        def normalize_array(array)
+          array.map { |v| normalize_value(v) }
+        end
+
+        def normalize_value(value)
+          case value
+          when Hash then normalize_hash(value)
+          when Array then normalize_array(value)
+          else value.to_s
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require_relative "../schema/adapter"
+require_relative "encoding"
 
 module LlmCostTracker
   module Ledger
@@ -8,7 +8,7 @@ module LlmCostTracker
       module Query
         class << self
           def apply(tags)
-            normalized_tags = (tags || {}).to_h.transform_keys(&:to_s).transform_values(&:to_s)
+            normalized_tags = (tags || {}).to_h.transform_keys(&:to_s).transform_values { |v| Encoding.encode(v) }
             return LlmCostTracker::Call.all if normalized_tags.empty?
 
             normalized_tags.inject(LlmCostTracker::Call.all) do |relation, (key, value)|

data/lib/llm_cost_tracker/ledger.rb

@@ -5,7 +5,8 @@ require_relative "ledger/schema/calls"
 require_relative "ledger/schema/call_rollups"
 require_relative "ledger/schema/call_line_items"
 require_relative "ledger/schema/call_tags"
-require_relative "ledger/schema/provider_invoices"
+require_relative "ledger/schema/ingestion_inbox_entries"
+require_relative "ledger/schema/ingestion_leases"
 require_relative "ledger/tags/query"
 require_relative "ledger/tags/sql"
 require_relative "ledger/period"

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/masking.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Masking
+    SENSITIVE_KEYS = %i[
+      provider_api_key_id provider_workspace_id provider_organization_id provider_project_id
+    ].to_set.freeze
+    MASK_TAIL_LENGTH = 4
+
+    module_function
+
+    def mask_value(key, value)
+      string = value.to_s
+      return string unless SENSITIVE_KEYS.include?(key.to_sym)
+      return string if string.length <= MASK_TAIL_LENGTH
+
+      "***#{string[-MASK_TAIL_LENGTH, MASK_TAIL_LENGTH]}"
+    end
+
+    def format_attribution(attribution, separator: ", ")
+      return "" if attribution.nil? || attribution.empty?
+
+      attribution.map { |key, value| "#{key}=#{mask_value(key, value)}" }.join(separator)
+    end
+
+    def mask_hash(hash)
+      return hash unless hash.is_a?(Hash)
+
+      hash.each_with_object({}) do |(key, value), masked|
+        masked[key] = case value
+                      when Hash then mask_hash(value)
+                      when Array then value.map { |entry| entry.is_a?(Hash) ? mask_hash(entry) : entry }
+                      else
+                        mask_value(key, value)
+                      end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -2,6 +2,7 @@
 
 require "faraday"
 require "json"
+require "stringio"
 require "uri"
 
 require_relative "../logging"
@@ -17,34 +18,116 @@ 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_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
 
+        invoke_app_with_capture(
+          request_env: request_env, parser: parser, request_url: request_url,
+          request_body: request_body, streaming: streaming, stream_buffer: stream_buffer,
+          context_tags: context_tags, metadata: metadata, started_at: started_at
+        )
+      end
+
+      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
         @app.call(request_env).on_complete do |response_env|
+          response_received = true
           process(
-            parser: parser,
-            request_url: request_url,
-            request_body: request_body,
-            response_env: response_env,
+            parser: parser, request_url: request_url, request_body: request_body,
+            response_env: response_env, latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at),
+            streaming: streaming, stream_buffer: stream_buffer,
+            context_tags: context_tags, metadata: metadata
+          )
+        end
+      rescue StandardError => e
+        if streaming && parser && !response_received
+          process_interrupted_stream(
+            parser: parser, request_url: request_url, request_body: request_body,
             latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at),
-            streaming: streaming,
-            stream_buffer: stream_buffer,
-            context_tags: context_tags,
-            metadata: metadata
+            context_tags: context_tags, metadata: metadata, error: e
           )
         end
+        raise
       end
 
-      private
+      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)
+
+        stream_options = request_parsed["stream_options"]
+        return nil if stream_options.is_a?(Hash) && stream_options.key?("include_usage")
+
+        request_parsed["stream_options"] = (stream_options || {}).merge("include_usage" => true)
+        new_body = request_parsed.to_json
+        request_env.body = new_body
+        new_body
+      end
+
+      def process_interrupted_stream(parser:, request_url:, request_body:, latency_ms:,
+                                     context_tags:, metadata:, error:)
+        request = parser.safe_json_parse(request_body)
+        event = Event.build(
+          provider: parser.provider_for(request_url),
+          model: request["model"] || Event::UNKNOWN_MODEL,
+          token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
+          stream: true,
+          usage_source: :unknown
+        )
+        merged_metadata = (metadata || {}).merge(
+          stream_interrupted: true,
+          stream_interrupted_error: "#{error.class}: #{error.message}"
+        )
+        Tracker.record(
+          event: event,
+          latency_ms: latency_ms,
+          metadata: merged_metadata,
+          context_tags: context_tags
+        )
+      rescue StandardError => e
+        Logging.warn("Error recording interrupted stream: #{e.class}: #{e.message}")
+      end
 
       def process(parser:, request_url:, request_body:, response_env:,
                   latency_ms:, streaming:, stream_buffer:, context_tags:, metadata:)
@@ -70,7 +153,7 @@ module LlmCostTracker
         return unless parsed
 
         Tracker.record(
-          capture: parsed,
+          event: parsed,
           latency_ms: latency_ms,
           metadata: metadata,
           context_tags: context_tags
@@ -102,21 +185,14 @@ module LlmCostTracker
       end
 
       def parse_stream(parser:, request_url:, request_body:, response_env:, stream_buffer:)
-        if stream_buffer&.dig(:overflowed)
-          Logging.warn(capture_warning(request_url, stream_buffer))
-          return parser.parse_stream(
-            request_url: request_url,
-            request_body: request_body,
-            response_status: response_env.status,
-            response_headers: response_env.response_headers
-          )
-        end
+        overflowed = stream_buffer&.dig(:overflowed) == true
+        Logging.warn(capture_warning(request_url, stream_buffer)) if overflowed
 
         body = stream_buffer&.dig(:buffer)&.string
         body = read_body(response_env.body) if body.blank?
 
         if body.blank?
-          Logging.warn(capture_warning(request_url, stream_buffer))
+          Logging.warn(capture_warning(request_url, stream_buffer)) unless overflowed
           return parser.parse_stream(
             request_url: request_url,
             request_body: request_body,
@@ -125,7 +201,7 @@ module LlmCostTracker
           )
         end
 
-        events = Parsers::SSE.parse(body)
+        events = overflowed ? [] : Parsers::SSE.parse(body)
         parser.parse_stream(
           request_url: request_url,
           request_body: request_body,
@@ -135,6 +211,17 @@ module LlmCostTracker
         )
       end
 
+      def forward_on_data_chunk(callable, chunk, size, env)
+        arity = callable.arity
+        return callable.call(chunk, size, env) if arity.negative?
+
+        case arity
+        when 0, 1 then callable.call(chunk)
+        when 2 then callable.call(chunk, size)
+        else callable.call(chunk, size, env)
+        end
+      end
+
       def install_stream_tap(request_env)
         request = request_env.request
         return nil unless request
@@ -145,16 +232,16 @@ module LlmCostTracker
         state = { buffer: StringIO.new, bytes: 0, overflowed: false }
         request.on_data = proc do |chunk, size, env|
           chunk = chunk.to_s
-          unless state[:overflowed]
-            if state[:bytes] + chunk.bytesize <= Capture::Stream::LIMIT_BYTES
-              state[:buffer] << chunk
-              state[:bytes] += chunk.bytesize
-            else
-              state[:overflowed] = true
-              state[:buffer] = nil
-            end
+          remaining = Capture::Stream::LIMIT_BYTES - state[:bytes]
+          if chunk.bytesize <= remaining
+            state[:buffer] << chunk
+            state[:bytes] += chunk.bytesize
+          else
+            state[:buffer] << chunk.byteslice(0, remaining) if remaining.positive?
+            state[:bytes] += [remaining, 0].max
+            state[:overflowed] = true
           end
-          original.call(chunk, size, env)
+          forward_on_data_chunk(original, chunk, size, env)
         end
         state
       rescue StandardError => e