llm_cost_tracker 0.2.0 → 0.3.1

data/lib/llm_cost_tracker/price_sync.rb

@@ -0,0 +1,142 @@
+# frozen_string_literal: true
+
+require "date"
+
+require_relative "price_sync/fetcher"
+require_relative "price_sync/raw_price"
+require_relative "price_sync/source"
+require_relative "price_sync/source_result"
+require_relative "price_sync/registry_loader"
+require_relative "price_sync/registry_writer"
+require_relative "price_sync/refresh_plan_builder"
+require_relative "price_sync/model_catalog"
+require_relative "price_sync/merger"
+require_relative "price_sync/validator"
+require_relative "price_sync/sources/litellm"
+require_relative "price_sync/sources/open_router"
+
+module LlmCostTracker
+  module PriceSync
+    DEFAULT_OUTPUT_PATH = PriceRegistry::DEFAULT_PRICES_PATH
+
+    SourceUsage = Data.define(:prices_count, :source_version)
+    SyncResult = Data.define(
+      :path,
+      :updated_models,
+      :changes,
+      :orphaned_models,
+      :failed_sources,
+      :discrepancies,
+      :rejected,
+      :flagged,
+      :sources_used,
+      :written
+    )
+    CheckResult = Data.define(
+      :path,
+      :changes,
+      :orphaned_models,
+      :failed_sources,
+      :discrepancies,
+      :rejected,
+      :flagged,
+      :sources_used,
+      :up_to_date
+    )
+    RefreshPlan = Data.define(
+      :path,
+      :registry,
+      :updated_registry,
+      :accepted,
+      :changes,
+      :orphaned_models,
+      :failed_sources,
+      :discrepancies,
+      :rejected,
+      :flagged,
+      :sources_used,
+      :source_results
+    ) do
+      def refresh_succeeded?
+        source_results.any? { |_source, result| result.prices.any? }
+      end
+
+      def up_to_date?
+        changes.empty? && failed_sources.empty? && rejected.empty?
+      end
+    end
+
+    class << self
+      def sync(path: DEFAULT_OUTPUT_PATH, seed_path: DEFAULT_OUTPUT_PATH, preview: false, strict: false,
+               fetcher: Fetcher.new, today: Date.today)
+        plan = RefreshPlanBuilder.new(sources: sources).call(
+          path: path,
+          seed_path: seed_path,
+          fetcher: fetcher,
+          today: today
+        )
+        raise Error, strict_failure_message(plan) if strict_sync_failure?(plan, strict: strict)
+
+        written = !preview && plan.refresh_succeeded?
+        RegistryWriter.new.call(path: plan.path, registry: plan.updated_registry) if written
+
+        SyncResult.new(
+          path: plan.path,
+          updated_models: plan.changes.keys.sort,
+          changes: plan.changes,
+          orphaned_models: plan.orphaned_models,
+          failed_sources: plan.failed_sources,
+          discrepancies: plan.discrepancies,
+          rejected: plan.rejected,
+          flagged: plan.flagged,
+          sources_used: plan.sources_used,
+          written: written
+        )
+      end
+
+      def check(path: DEFAULT_OUTPUT_PATH, seed_path: DEFAULT_OUTPUT_PATH, fetcher: Fetcher.new, today: Date.today)
+        plan = RefreshPlanBuilder.new(sources: sources).call(
+          path: path,
+          seed_path: seed_path,
+          fetcher: fetcher,
+          today: today
+        )
+
+        CheckResult.new(
+          path: plan.path,
+          changes: plan.changes,
+          orphaned_models: plan.orphaned_models,
+          failed_sources: plan.failed_sources,
+          discrepancies: plan.discrepancies,
+          rejected: plan.rejected,
+          flagged: plan.flagged,
+          sources_used: plan.sources_used,
+          up_to_date: plan.up_to_date?
+        )
+      end
+
+      private
+
+      def sources
+        [Sources::Litellm.new, Sources::OpenRouter.new]
+      end
+
+      def strict_sync_failure?(plan, strict:)
+        strict && (plan.failed_sources.any? || plan.rejected.any?)
+      end
+
+      def strict_failure_message(plan)
+        messages = []
+        if plan.failed_sources.any?
+          details = plan.failed_sources.map { |source, message| "#{source}: #{message}" }.join(", ")
+          messages << "source failures: #{details}"
+        end
+        if plan.rejected.any?
+          details = plan.rejected.map { |issue| "#{issue.model} (#{issue.reason})" }.join(", ")
+          messages << "validator rejections: #{details}"
+        end
+        "Price sync failed in strict mode: #{messages.join('; ')}"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/pricing.rb

@@ -3,21 +3,11 @@
 require "monitor"
 
 module LlmCostTracker
-  # Calculates costs from price entries expressed in USD per 1M tokens.
   module Pricing
     PRICES = PriceRegistry.builtin_prices
     MUTEX = Monitor.new
 
     class << self
-      # Estimate model cost from token counts.
-      #
-      # @param model [String] Provider model identifier.
-      # @param input_tokens [Integer] Input token count, including cached tokens if reported that way.
-      # @param output_tokens [Integer] Output token count.
-      # @param cached_input_tokens [Integer] OpenAI-style cached input tokens.
-      # @param cache_read_input_tokens [Integer] Anthropic-style cache read tokens.
-      # @param cache_creation_input_tokens [Integer] Anthropic-style cache creation tokens.
-      # @return [LlmCostTracker::Cost, nil] nil when no price is configured for the model.
       def cost_for(model:, input_tokens:, output_tokens:, cached_input_tokens: 0,
                    cache_read_input_tokens: 0, cache_creation_input_tokens: 0)
         prices = lookup(model)
@@ -111,7 +101,6 @@ module LlmCostTracker
         model.to_s.split("/").last
       end
 
-      # Try to match model names like "gpt-4o-2024-08-06" to "gpt-4o".
       def fuzzy_match(model, normalized_model, table)
         sorted_price_keys(table).each do |key|
           return table[key] if model.start_with?(key) || normalized_model.start_with?(key)

data/lib/llm_cost_tracker/railtie.rb

@@ -15,7 +15,6 @@ module LlmCostTracker
     end
 
     initializer "llm_cost_tracker.configure" do
-      # Auto-require ActiveRecord storage if configured
       ActiveSupport.on_load(:active_record) do
         if LlmCostTracker.configuration.active_record?
           require_relative "llm_api_call"

data/lib/llm_cost_tracker/report.rb

@@ -8,11 +8,6 @@ module LlmCostTracker
     DEFAULT_DAYS = ReportData::DEFAULT_DAYS
 
     class << self
-      # Render a terminal-friendly cost report from ActiveRecord storage.
-      #
-      # @param days [Integer] Number of trailing days to include.
-      # @param now [Time] Report end time.
-      # @return [String]
       def generate(days: DEFAULT_DAYS, now: Time.now.utc, tag_breakdowns: nil)
         ReportFormatter.new(data(days: days, now: now, tag_breakdowns: tag_breakdowns)).to_s
       rescue LoadError => e

data/lib/llm_cost_tracker/storage/active_record_store.rb

@@ -19,22 +19,23 @@ module LlmCostTracker
             tags:          tags_for_storage(tags),
             tracked_at:    event.tracked_at
           }
-          attributes[:latency_ms] = event.latency_ms if model_class.latency_column?
-
-          model_class.create!(attributes)
+          attributes[:latency_ms] = event.latency_ms if LlmCostTracker::LlmApiCall.latency_column?
+          attributes[:stream] = event.stream if LlmCostTracker::LlmApiCall.stream_column?
+          attributes[:usage_source] = event.usage_source if LlmCostTracker::LlmApiCall.usage_source_column?
+          if LlmCostTracker::LlmApiCall.provider_response_id_column?
+            attributes[:provider_response_id] = event.provider_response_id
+          end
+
+          LlmCostTracker::LlmApiCall.create!(attributes)
         end
 
         def monthly_total(time: Time.now.utc)
-          model_class
+          LlmCostTracker::LlmApiCall
             .where(tracked_at: time.beginning_of_month..time)
             .sum(:total_cost)
             .to_f
         end
 
-        def model_class
-          LlmCostTracker::LlmApiCall
-        end
-
         private
 
         def stringify_tags(tags)
@@ -42,7 +43,7 @@ module LlmCostTracker
         end
 
         def tags_for_storage(tags)
-          model_class.tags_json_column? ? tags : tags.to_json
+          LlmCostTracker::LlmApiCall.tags_json_column? ? tags : tags.to_json
         end
 
         def stringify_tag_value(value)

data/lib/llm_cost_tracker/stream_collector.rb

@@ -0,0 +1,162 @@
+# frozen_string_literal: true
+
+require "monitor"
+
+require_relative "value_helpers"
+
+module LlmCostTracker
+  class StreamCollector
+    attr_reader :provider
+
+    def initialize(provider:, model:, latency_ms: nil, provider_response_id: nil, metadata: {})
+      @provider = provider.to_s
+      @model = model
+      @latency_ms = latency_ms
+      @provider_response_id = provider_response_id
+      @metadata = ValueHelpers.deep_dup(metadata || {})
+      @events = []
+      @explicit_usage = nil
+      @started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      @finished = false
+      @monitor = Monitor.new
+    end
+
+    def model = @monitor.synchronize { @model }
+
+    def metadata = @monitor.synchronize { ValueHelpers.deep_dup(@metadata) }
+
+    def provider_response_id = @monitor.synchronize { @provider_response_id }
+
+    def model=(value)
+      @monitor.synchronize do
+        ensure_open!
+        @model = value
+      end
+    end
+
+    def provider_response_id=(value)
+      @monitor.synchronize do
+        ensure_open!
+        @provider_response_id = value
+      end
+    end
+
+    def event(data, type: nil)
+      @monitor.synchronize do
+        ensure_open!
+        @events << { event: type, data: ValueHelpers.deep_dup(data) } unless data.nil?
+      end
+      self
+    end
+    alias chunk event
+
+    def usage(input_tokens:, output_tokens:, **extra)
+      @monitor.synchronize do
+        ensure_open!
+        @explicit_usage = ValueHelpers.deep_dup(
+          extra.merge(
+            input_tokens: input_tokens.to_i,
+            output_tokens: output_tokens.to_i
+          )
+        )
+      end
+      self
+    end
+
+    def finish!(errored: false)
+      snapshot = @monitor.synchronize do
+        return if @finished
+
+        @finished = true
+        {
+          events: ValueHelpers.deep_dup(@events),
+          explicit_usage: ValueHelpers.deep_dup(@explicit_usage),
+          model: @model,
+          latency_ms: @latency_ms,
+          provider_response_id: @provider_response_id,
+          metadata: ValueHelpers.deep_dup(@metadata)
+        }
+      end
+
+      parsed = build_parsed_usage(snapshot)
+      Tracker.record(
+        provider: parsed.provider,
+        model: parsed.model,
+        input_tokens: parsed.input_tokens,
+        output_tokens: parsed.output_tokens,
+        latency_ms: snapshot[:latency_ms] || elapsed_ms,
+        stream: true,
+        usage_source: parsed.usage_source,
+        provider_response_id: parsed.provider_response_id || snapshot[:provider_response_id],
+        metadata: error_metadata(errored).merge(snapshot[:metadata]).merge(parsed.metadata)
+      )
+    end
+
+    private
+
+    def ensure_open!
+      return unless @finished
+
+      raise FrozenError, "can't modify finished LlmCostTracker::StreamCollector"
+    end
+
+    def build_parsed_usage(snapshot)
+      return build_from_explicit_usage(snapshot) if snapshot[:explicit_usage]
+
+      parsed = Parsers::Registry.find_for_provider(@provider)&.parse_stream(nil, nil, 200, snapshot[:events])
+      return finalize(parsed, snapshot) if parsed
+
+      build_unknown_usage(snapshot)
+    end
+
+    def finalize(parsed, snapshot)
+      parsed.with(
+        provider: @provider,
+        model: present_model(parsed.model) || snapshot[:model]
+      )
+    end
+
+    def present_model(value)
+      return nil if value.nil?
+
+      string = value.to_s
+      return nil if string.empty? || string == "unknown"
+
+      string
+    end
+
+    def build_from_explicit_usage(snapshot)
+      explicit = snapshot[:explicit_usage]
+      input = explicit[:input_tokens]
+      output = explicit[:output_tokens]
+      extras = explicit.except(:input_tokens, :output_tokens)
+
+      ParsedUsage.build(
+        provider: @provider,
+        model: snapshot[:model],
+        input_tokens: input,
+        output_tokens: output,
+        total_tokens: input + output,
+        stream: true,
+        usage_source: :manual,
+        **extras
+      )
+    end
+
+    def build_unknown_usage(snapshot)
+      ParsedUsage.build(
+        provider: @provider,
+        model: snapshot[:model],
+        input_tokens: 0,
+        output_tokens: 0,
+        total_tokens: 0,
+        stream: true,
+        usage_source: :unknown
+      )
+    end
+
+    def error_metadata(errored) = errored ? { stream_errored: true } : {}
+
+    def elapsed_ms = ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - @started_at) * 1000).round
+  end
+end

data/lib/llm_cost_tracker/tags_column.rb

@@ -24,5 +24,17 @@ module LlmCostTracker
     def latency_column?
       columns_hash.key?("latency_ms")
     end
+
+    def stream_column?
+      columns_hash.key?("stream")
+    end
+
+    def usage_source_column?
+      columns_hash.key?("usage_source")
+    end
+
+    def provider_response_id_column?
+      columns_hash.key?("provider_response_id")
+    end
   end
 end

data/lib/llm_cost_tracker/tracker.rb

@@ -6,21 +6,15 @@ module LlmCostTracker
   class Tracker
     EVENT_NAME = "llm_request.llm_cost_tracker"
 
+    USAGE_SOURCES = %i[response stream_final manual unknown].freeze
+
     class << self
       def enforce_budget!
         Budget.enforce!
       end
 
-      # Build, notify, persist, and budget-check a single LLM usage event.
-      #
-      # @param provider [String] Provider name.
-      # @param model [String] Model identifier.
-      # @param input_tokens [Integer] Input token count.
-      # @param output_tokens [Integer] Output token count.
-      # @param metadata [Hash] Attribution tags plus provider-specific usage metadata.
-      # @param latency_ms [Integer, nil] Optional latency in milliseconds.
-      # @return [LlmCostTracker::Event]
-      def record(provider:, model:, input_tokens:, output_tokens:, metadata: {}, latency_ms: nil)
+      def record(provider:, model:, input_tokens:, output_tokens:, latency_ms: nil, stream: false,
+                 usage_source: nil, provider_response_id: nil, metadata: {})
         usage = EventMetadata.usage_data(input_tokens, output_tokens, metadata)
 
         cost_data = Pricing.cost_for(
@@ -43,13 +37,14 @@ module LlmCostTracker
           cost: cost_data,
           tags: LlmCostTracker.configuration.default_tags.merge(EventMetadata.tags(metadata)).freeze,
           latency_ms: normalized_latency_ms(latency_ms),
+          stream: stream ? true : false,
+          usage_source: normalized_usage_source(usage_source),
+          provider_response_id: normalized_provider_response_id(provider_response_id),
           tracked_at: Time.now.utc
         )
 
-        # Emit ActiveSupport::Notifications event
         ActiveSupport::Notifications.instrument(EVENT_NAME, event.to_h)
 
-        # Store based on backend
         stored = store(event)
         Budget.check!(event) unless stored == false
 
@@ -77,6 +72,8 @@ module LlmCostTracker
                   "tokens=#{event.input_tokens}+#{event.output_tokens} " \
                   "cost=#{log_cost_label(event)}"
         message += " latency=#{event.latency_ms}ms" if event.latency_ms
+        message += " stream=#{event.stream}" if event.stream
+        message += " source=#{event.usage_source}" if event.usage_source
         message += " tags=#{event.tags}" unless event.tags.empty?
 
         Logging.log(config.log_level, message)
@@ -119,6 +116,20 @@ module LlmCostTracker
 
         [latency_ms.to_i, 0].max
       end
+
+      def normalized_usage_source(value)
+        return nil if value.nil?
+
+        symbol = value.to_sym
+        USAGE_SOURCES.include?(symbol) ? symbol.to_s : nil
+      end
+
+      def normalized_provider_response_id(value)
+        return nil if value.nil?
+
+        string = value.to_s
+        string.empty? ? nil : string
+      end
     end
   end
 end

data/lib/llm_cost_tracker/value_helpers.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module ValueHelpers
+    class << self
+      def deep_dup(value)
+        case value
+        when Hash
+          value.each_with_object({}) do |(key, nested_value), duplicated|
+            duplicated[deep_dup(key)] = deep_dup(nested_value)
+          end
+        when Array
+          value.map { |nested_value| deep_dup(nested_value) }
+        when String
+          value.dup
+        else
+          value
+        end
+      end
+
+      def deep_freeze(value)
+        case value
+        when Hash
+          value.each do |key, nested_value|
+            deep_freeze(key)
+            deep_freeze(nested_value)
+          end
+          value.frozen? ? value : value.freeze
+        when Array
+          value.each { |nested_value| deep_freeze(nested_value) }
+          value.frozen? ? value : value.freeze
+        when String
+          value.frozen? ? value : value.freeze
+        else
+          value
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/version.rb

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

data/lib/llm_cost_tracker.rb

@@ -2,15 +2,18 @@
 
 require "active_support"
 require "active_support/notifications"
+require "monitor"
 
 require_relative "llm_cost_tracker/version"
 require_relative "llm_cost_tracker/configuration"
 require_relative "llm_cost_tracker/errors"
 require_relative "llm_cost_tracker/logging"
+require_relative "llm_cost_tracker/parameter_hash"
 require_relative "llm_cost_tracker/cost"
 require_relative "llm_cost_tracker/event"
 require_relative "llm_cost_tracker/parsed_usage"
 require_relative "llm_cost_tracker/price_registry"
+require_relative "llm_cost_tracker/price_sync"
 require_relative "llm_cost_tracker/pricing"
 require_relative "llm_cost_tracker/parsers/base"
 require_relative "llm_cost_tracker/parsers/openai_usage"
@@ -18,6 +21,7 @@ require_relative "llm_cost_tracker/parsers/openai"
 require_relative "llm_cost_tracker/parsers/openai_compatible"
 require_relative "llm_cost_tracker/parsers/anthropic"
 require_relative "llm_cost_tracker/parsers/gemini"
+require_relative "llm_cost_tracker/parsers/sse"
 require_relative "llm_cost_tracker/parsers/registry"
 require_relative "llm_cost_tracker/middleware/faraday"
 require_relative "llm_cost_tracker/budget"
@@ -34,71 +38,80 @@ require_relative "llm_cost_tracker/report_formatter"
 require_relative "llm_cost_tracker/report"
 
 module LlmCostTracker
-  class << self
-    attr_writer :configuration
+  CONFIGURATION_MUTEX = Monitor.new
 
+  class << self
     def configuration
-      @configuration ||= Configuration.new
+      CONFIGURATION_MUTEX.synchronize { @configuration ||= Configuration.new }
     end
 
-    # Configure the gem once during application boot.
-    #
-    # @yieldparam configuration [LlmCostTracker::Configuration]
-    # @return [void]
     def configure
-      yield(configuration)
-      configuration.normalize_openai_compatible_providers!
-      warn_for_configuration!
+      config = CONFIGURATION_MUTEX.synchronize do
+        current = @configuration || Configuration.new
+        current = current.dup_for_configuration if current.finalized?
+        @configuration = current
+        yield(current)
+        current.normalize_openai_compatible_providers!
+        current.finalize!
+        current
+      end
+      warn_for_configuration!(config)
     end
 
     def reset_configuration!
-      @configuration = Configuration.new
+      CONFIGURATION_MUTEX.synchronize { @configuration = Configuration.new }
+    end
+
+    def enforce_budget!
+      Tracker.enforce_budget!
     end
 
-    # Track an LLM request manually for non-Faraday clients.
-    #
-    #   LlmCostTracker.track(
-    #     provider: :openai,
-    #     model: "gpt-4o",
-    #     input_tokens: 150,
-    #     output_tokens: 50,
-    #     feature: "chat",
-    #     user_id: current_user.id
-    #   )
-    #
-    # @param provider [String, Symbol] Provider name, such as :openai or :anthropic.
-    # @param model [String] Provider model identifier.
-    # @param input_tokens [Integer] Billed input token count.
-    # @param output_tokens [Integer] Billed output token count.
-    # @param latency_ms [Integer, nil] Optional request latency in milliseconds.
-    # @param metadata [Hash] Attribution tags and provider-specific usage metadata.
-    # @return [LlmCostTracker::Event] The tracked event.
-    def track(provider:, model:, input_tokens:, output_tokens:, latency_ms: nil, **metadata)
+    def track(provider:, model:, input_tokens:, output_tokens:, latency_ms: nil, stream: false, usage_source: :manual,
+              enforce_budget: false, provider_response_id: nil, **metadata)
+      enforce_budget! if enforce_budget
       Tracker.record(
         provider: provider.to_s,
         model: model,
         input_tokens: input_tokens,
         output_tokens: output_tokens,
         latency_ms: latency_ms,
+        stream: stream,
+        usage_source: usage_source,
+        provider_response_id: provider_response_id,
+        metadata: metadata
+      )
+    end
+
+    def track_stream(provider:, model:, latency_ms: nil, enforce_budget: false, provider_response_id: nil, **metadata)
+      require_relative "llm_cost_tracker/stream_collector"
+      enforce_budget! if enforce_budget
+      collector = StreamCollector.new(
+        provider: provider.to_s,
+        model: model,
+        latency_ms: latency_ms,
+        provider_response_id: provider_response_id,
         metadata: metadata
       )
+      yield collector
+      collector.finish!
+    rescue StandardError
+      collector&.finish!(errored: true)
+      raise
     end
 
     private
 
-    def warn_for_configuration!
-      return unless configuration.budget_exceeded_behavior == :block_requests
-      return if configuration.active_record?
+    def warn_for_configuration!(config = configuration)
+      return unless config.budget_exceeded_behavior == :block_requests
+      return if config.active_record?
 
       Logging.warn(":block_requests requires storage_backend = :active_record; preflight blocking will be skipped.")
     end
   end
 end
 
-# Load Railtie if Rails is present
 require_relative "llm_cost_tracker/railtie" if defined?(Rails::Railtie)
 
-# Auto-register Faraday middleware
 if defined?(Faraday)
   Faraday::Middleware.register_middleware(
     llm_cost_tracker: LlmCostTracker::Middleware::Faraday