llm_cost_tracker 0.11.0 → 0.13.0

data/lib/llm_cost_tracker/ledger/schema/ingestion/leases.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+require_relative "../base"
+
+module LlmCostTracker
+  module Ledger
+    module Schema
+      module Ingestion
+        module Leases
+          extend Base
+
+          columns :name, :locked_by, :locked_until, :created_at, :updated_at
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/schema.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require_relative "schema/adapter"
+require_relative "schema/calls"
+require_relative "schema/call_rollups"
+require_relative "schema/call_line_items"
+require_relative "schema/call_tags"
+require_relative "schema/ingestion/inbox_entries"
+require_relative "schema/ingestion/leases"
+
+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
+      ASYNC_SCHEMAS = [
+        [Ingestion::InboxEntries, "llm_cost_tracker_ingestion_inbox_entries"],
+        [Ingestion::Leases, "llm_cost_tracker_ingestion_leases"]
+      ].freeze
+    end
+  end
+end

data/lib/llm_cost_tracker/ledger/store.rb

@@ -3,28 +3,31 @@
 require "json"
 
 require_relative "../pricing"
-require_relative "../billing/line_item"
 require_relative "rollups"
 require_relative "tags/encoding"
 
 module LlmCostTracker
   module Ledger
-    class Store
+    module Store
       class << self
-        def insert(events, skip_existence_check: false)
+        def insert(events)
           events = Array(events)
           return if events.empty?
 
-          insertable = skip_existence_check ? events : insertable_events(events)
-          return unless insertable.any?
+          persist_records(events)
+          Ledger::Rollups.increment_safely!(events) if LlmCostTracker.configuration.cache_rollups
+        end
+
+        def persist_records(events)
+          events = Array(events)
+          return if events.empty?
 
           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)
+            rows = events.map { |event| attributes_for(event) }
+            call_ids = insert_calls_returning_ids(rows, events)
+            insert_line_items(events, call_ids)
+            insert_call_tags(events, call_ids)
           end
-          increment_rollups_safely(insertable) if LlmCostTracker.configuration.cache_rollups
         end
 
         private
@@ -45,22 +48,22 @@ module LlmCostTracker
             provider: event.provider,
             model: event.model,
             tracked_at: event.tracked_at,
-            pricing_mode: event.pricing_mode&.name,
+            pricing_mode: event.pricing_mode,
             latency_ms: event.latency_ms,
             stream: event.stream,
-            usage_source: event.usage_source&.name,
+            usage_source: event.usage_source,
             provider_response_id: event.provider_response_id,
             provider_project_id: event.provider_project_id,
             provider_api_key_id: event.provider_api_key_id,
             provider_workspace_id: event.provider_workspace_id,
-            batch: event.batch,
+            batch: event.batch?,
             cost_status: event.cost_status,
             pricing_snapshot: event.pricing_snapshot
           }
 
           attributes
             .merge(event.token_usage.to_h)
-            .merge(Pricing.stored_cost_attributes(event.cost || {}))
+            .merge(total_cost: event.cost&.total)
         end
 
         def call_ids_for(events)
@@ -89,20 +92,20 @@ module LlmCostTracker
           {
             llm_cost_tracker_call_id: call_id,
             position: position,
-            kind: line_item.kind&.to_s,
-            direction: line_item.direction&.to_s,
-            modality: line_item.modality&.to_s,
-            cache_state: line_item.cache_state&.to_s || "none",
+            kind: line_item.kind,
+            direction: line_item.direction,
+            modality: line_item.modality,
+            cache_state: line_item.cache_state,
             quantity: line_item.quantity,
-            unit: line_item.unit&.to_s,
+            unit: line_item.unit,
             rate_amount: line_item.rate_amount,
             rate_quantity: line_item.rate_quantity,
             cost: line_item.cost,
             currency: line_item.currency,
             cost_status: line_item.cost_status,
-            pricing_basis: line_item.pricing_basis&.to_s,
+            pricing_basis: line_item.pricing_basis,
             price_key: line_item.price_key,
-            price_source: line_item.price_source&.to_s,
+            price_source: line_item.price_source,
             price_source_version: line_item.price_source_version,
             provider_field: line_item.provider_field,
             provider_item_id: line_item.provider_item_id,
@@ -129,26 +132,6 @@ module LlmCostTracker
         def stored_details(details)
           (details || {}).transform_keys(&:to_s).transform_values { |value| Tags::Encoding.normalize_value(value) }
         end
-
-        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)
-          existing_ids = LlmCostTracker::Call.where(event_id: events.map(&:event_id)).pluck(:event_id).to_set
-          seen_ids = Set.new
-
-          events.select do |event|
-            event_id = event.event_id
-            !existing_ids.include?(event_id) && seen_ids.add?(event_id)
-          end
-        end
       end
     end
   end

data/lib/llm_cost_tracker/ledger/tags/{sql.rb → breakdown.rb}

@@ -5,7 +5,7 @@ require_relative "../../tags/key"
 module LlmCostTracker
   module Ledger
     module Tags
-      module Sql
+      module Breakdown
         UNTAGGED_LABEL = "(untagged)"
 
         class << self

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

@@ -6,9 +6,7 @@ module LlmCostTracker
   module Ledger
     module Tags
       module Encoding
-        module_function
-
-        def encode(value)
+        def self.encode(value)
           case value
           when Hash then JSON.generate(normalize_hash(value))
           when Array then JSON.generate(normalize_array(value))
@@ -16,15 +14,15 @@ module LlmCostTracker
           end
         end
 
-        def normalize_hash(hash)
+        def self.normalize_hash(hash)
           hash.transform_keys(&:to_s).sort.to_h.transform_values { |v| normalize_value(v) }
         end
 
-        def normalize_array(array)
+        def self.normalize_array(array)
           array.map { |v| normalize_value(v) }
         end
 
-        def normalize_value(value)
+        def self.normalize_value(value)
           case value
           when Hash then normalize_hash(value)
           when Array then normalize_array(value)

data/lib/llm_cost_tracker/ledger.rb

@@ -1,29 +1,19 @@
 # frozen_string_literal: true
 
-require_relative "ledger/schema/adapter"
-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/ingestion_inbox_entries"
-require_relative "ledger/schema/ingestion_leases"
-require_relative "ledger/tags/query"
-require_relative "ledger/tags/sql"
+require_relative "ledger/schema"
 require_relative "ledger/period"
-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
+    module Tags
+      autoload :Query, "llm_cost_tracker/ledger/tags/query"
+      autoload :Breakdown, "llm_cost_tracker/ledger/tags/breakdown"
+    end
+
+    module Period
+      autoload :Totals, "llm_cost_tracker/ledger/period/totals"
     end
   end
 end

data/lib/llm_cost_tracker/logging.rb

@@ -2,32 +2,15 @@
 
 module LlmCostTracker
   module Logging
-    PREFIX = "[LlmCostTracker]"
-
     class << self
-      def debug(message)
-        log(:debug, message)
-      end
-
-      def warn(message)
-        log(:warn, message)
-      end
+      def debug(message) = tagged.debug(message)
 
-      def log(level, message)
-        message = prefixed(message)
-        logger = Rails.logger
-        return Kernel.warn(message) unless logger
-
-        logger.public_send(level, message)
-      end
+      def warn(message) = tagged.warn(message)
 
       private
 
-      def prefixed(message)
-        message = message.to_s
-        return message if message.start_with?(PREFIX)
-
-        "#{PREFIX} #{message}"
+      def tagged
+        Rails.logger.tagged(LlmCostTracker.name)
       end
     end
   end

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -5,9 +5,7 @@ require "json"
 require "stringio"
 require "uri"
 
-require_relative "../logging"
 require_relative "../capture/sse"
-require_relative "../capture/stream"
 require_relative "../timing"
 
 module LlmCostTracker
@@ -19,12 +17,12 @@ module LlmCostTracker
       end
 
       def call(request_env)
-        return @app.call(request_env) unless enabled?
+        return @app.call(request_env) unless LlmCostTracker.configuration.enabled
 
         request_url  = request_env.url.to_s
         request_body = read_body(request_env.body)
         parser       = Parsers.find_for(request_url)
-        request_parsed = parser ? safe_json_parse(request_body) : nil
+        request_parsed = parser&.safe_json_parse(request_body)
         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
@@ -32,7 +30,7 @@ module LlmCostTracker
         stream_buffer = install_stream_tap(request_env) if streaming
 
         if parser
-          Tracker.enforce_budget!(
+          Budget.enforce!(
             provider: parser.provider_for(request_url),
             model: parser.model_for(request_url, request_parsed),
             request: request_parsed
@@ -42,59 +40,61 @@ module LlmCostTracker
         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
+          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:)
+      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, latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at),
-            streaming: streaming, stream_buffer: stream_buffer,
-            context_tags: context_tags, metadata: metadata
+            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,
+            parser: parser,
+            request_url: request_url,
+            request_body: request_body,
             latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at),
-            context_tags: context_tags, metadata: metadata, error: e
+            context_tags: context_tags,
+            metadata: metadata,
+            error: e
           )
         end
         raise
       end
 
       def inject_stream_usage_flag(request_env, parser, request_url, request_parsed)
-        return nil unless auto_enable_stream_usage?
+        return nil unless LlmCostTracker.configuration.auto_enable_stream_usage
         return nil unless parser&.auto_enable_stream_usage?(request_url)
 
         stream_options = request_parsed["stream_options"]
@@ -106,15 +106,20 @@ module LlmCostTracker
         new_body
       end
 
-      def process_interrupted_stream(parser:, request_url:, request_body:, latency_ms:,
-                                     context_tags:, metadata:, error:)
+      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),
+          token_usage: Usage::TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
-          usage_source: :unknown
+          usage_source: Usage::Source::UNKNOWN
         )
         merged_metadata = (metadata || {}).merge(
           stream_interrupted: true,
@@ -130,8 +135,15 @@ module LlmCostTracker
         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:)
+      def process(parser:,
+                  request_url:,
+                  request_body:,
+                  response_env:,
+                  latency_ms:,
+                  streaming:,
+                  stream_buffer:,
+                  context_tags:,
+                  metadata:)
         return unless parser
 
         parsed =
@@ -233,7 +245,7 @@ module LlmCostTracker
         state = { buffer: StringIO.new, bytes: 0, overflowed: false }
         request.on_data = proc do |chunk, size, env|
           chunk = chunk.to_s
-          remaining = Capture::Stream::LIMIT_BYTES - state[:bytes]
+          remaining = Capture::SSE::LIMIT_BYTES - state[:bytes]
           if chunk.bytesize <= remaining
             state[:buffer] << chunk
             state[:bytes] += chunk.bytesize
@@ -280,13 +292,12 @@ module LlmCostTracker
       end
 
       def capture_warning(request_url, stream_buffer)
-        unless stream_buffer&.dig(:overflowed)
-          return "Unable to capture streaming response for #{request_url_label(request_url)}; " \
-                 "recording usage_source=unknown. Use LlmCostTracker.track_stream for manual capture."
-        end
+        suffix = "recording usage_source=#{Usage::Source::UNKNOWN}. " \
+                 "Use LlmCostTracker.track_stream for manual capture."
+        label = request_url_label(request_url)
+        return "Unable to capture streaming response for #{label}; #{suffix}" unless stream_buffer&.dig(:overflowed)
 
-        "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."
+        "Streaming response for #{label} exceeded #{Capture::SSE::LIMIT_BYTES} bytes; #{suffix}"
       end
 
       def request_url_label(value)

data/lib/llm_cost_tracker/parsers.rb

@@ -1,45 +1,158 @@
 # frozen_string_literal: true
 
-module LlmCostTracker
-  module Parsers
-    autoload :Base,                 "llm_cost_tracker/parsers/base"
-    autoload :OpenaiUsage,          "llm_cost_tracker/parsers/openai_usage"
-    autoload :Openai,               "llm_cost_tracker/parsers/openai"
-    autoload :Azure,                "llm_cost_tracker/parsers/azure"
-    autoload :OpenaiCompatible,     "llm_cost_tracker/parsers/openai_compatible"
-    autoload :Anthropic,            "llm_cost_tracker/parsers/anthropic"
-    autoload :Gemini,               "llm_cost_tracker/parsers/gemini"
+require "active_support/core_ext/object/blank"
+require "json"
+require "uri"
 
-    MUTEX = Mutex.new
-    PARSER_CONSTANTS = %i[Openai Azure OpenaiCompatible Anthropic Gemini].freeze
+require_relative "providers"
 
-    module_function
+module LlmCostTracker
+  module Parsers
+    PARSER_PROVIDERS = %i[Openai Azure OpenaiCompatible Anthropic Gemini].freeze
 
-    def find_for(url)
-      PARSER_CONSTANTS.each do |name|
-        klass = const_get(name)
-        return instance_for(klass) if klass.match?(url)
+    def self.find_for(url)
+      instances.each do |klass, instance|
+        return instance if klass.match?(url)
       end
       nil
     end
 
-    def find_for_provider(provider)
+    def self.find_for_provider(provider)
       provider_name = provider.to_s.downcase
-      PARSER_CONSTANTS.each do |name|
-        klass = const_get(name)
-        return instance_for(klass) if klass.provider_names.include?(provider_name)
+      instances.each do |klass, instance|
+        return instance if klass.provider_names.include?(provider_name)
       end
       nil
     end
 
-    def instance_for(klass)
-      cached = (@instances ||= {})[klass]
-      return cached if cached
+    def self.parser_classes
+      PARSER_PROVIDERS.map { |name| Providers.const_get(name)::Parser }
+    end
+    private_class_method :parser_classes
+
+    def self.instances
+      @instances ||= parser_classes.to_h { |klass| [klass, klass.new] }.freeze
+    end
+    private_class_method :instances
+
+    module UrlMatchers
+      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? || hosts.include?(uri.host.to_s.downcase)
+          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)
+        end
+      end
+
+      def uri_matches?(url)
+        uri = parsed_uri(url)
+        uri ? yield(uri) : false
+      end
+
+      def parsed_uri(url)
+        URI.parse(url.to_s)
+      rescue URI::InvalidURIError
+        nil
+      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
+    end
+
+    class Base
+      extend UrlMatchers
+      include UrlMatchers
+
+      class << self
+        def match?(_url)
+          raise NotImplementedError
+        end
+
+        def provider_names
+          []
+        end
+      end
+
+      def parse(**)
+        raise NotImplementedError
+      end
+
+      def streaming_request?(_request_url, request_parsed)
+        request_parsed["stream"] == true
+      end
+
+      def model_for(_request_url, request_parsed)
+        request_parsed["model"]
+      end
+
+      def parse_stream(**)
+        nil
+      end
+
+      def auto_enable_stream_usage?(_request_url)
+        false
+      end
+
+      def safe_json_parse(body)
+        return {} if body.blank?
+
+        parsed = JSON.parse(body)
+        parsed.is_a?(Hash) ? parsed : {}
+      rescue JSON::ParserError
+        {}
+      end
+
+      private
+
+      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 value.present?
+        end
+
+        nil
+      end
 
-      MUTEX.synchronize do
-        @instances[klass] ||= klass.new
+      def build_unknown_stream_usage(provider:,
+                                     model:,
+                                     provider_response_id:,
+                                     pricing_mode: nil,
+                                     service_line_items: nil)
+        Event.build(
+          provider: provider,
+          provider_response_id: provider_response_id,
+          pricing_mode: pricing_mode,
+          model: model || Event::UNKNOWN_MODEL,
+          token_usage: Usage::TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
+          stream: true,
+          usage_source: Usage::Source::UNKNOWN,
+          service_line_items: service_line_items
+        )
       end
     end
-    private_class_method :instance_for
   end
 end