llm_cost_tracker 0.2.0 → 0.3.0

data/lib/llm_cost_tracker/configuration.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "errors"
+require_relative "value_helpers"
 
 module LlmCostTracker
   class Configuration
@@ -14,22 +15,32 @@ module LlmCostTracker
     STORAGE_ERROR_BEHAVIORS = %i[ignore warn raise].freeze
     STORAGE_BACKENDS = %i[log active_record custom].freeze
     UNKNOWN_PRICING_BEHAVIORS = %i[ignore warn raise].freeze
+    SHARED_SCALAR_ATTRIBUTES = %i[
+      enabled
+      custom_storage
+      on_budget_exceeded
+      monthly_budget
+      log_level
+      prices_file
+    ].freeze
+    SHARED_ENUM_ATTRIBUTES = {
+      storage_backend: [STORAGE_BACKENDS, :log],
+      budget_exceeded_behavior: [BUDGET_EXCEEDED_BEHAVIORS, :notify],
+      storage_error_behavior: [STORAGE_ERROR_BEHAVIORS, :warn],
+      unknown_pricing_behavior: [UNKNOWN_PRICING_BEHAVIORS, :warn]
+    }.freeze
 
-    attr_accessor :enabled,
-                  :custom_storage,     # callable object for :custom backend
-                  :default_tags,       # Hash of default tags added to every event
-                  :on_budget_exceeded, # callable, receives event hash
-                  :monthly_budget,     # Float, in USD — nil means no limit
-                  :log_level,          # :debug, :info, :warn
-                  :prices_file,        # JSON/YAML file that overrides built-in prices
-                  :pricing_overrides,  # Hash to override built-in pricing
-                  :report_tag_breakdowns # Array of tag keys to break down in the rake report
-
-    attr_reader :budget_exceeded_behavior, # :notify, :raise, :block_requests
-                :storage_backend, # :log, :active_record, :custom
-                :storage_error_behavior, # :ignore, :warn, :raise
-                :unknown_pricing_behavior, # :ignore, :warn, :raise
-                :openai_compatible_providers
+    attr_reader(
+      *SHARED_SCALAR_ATTRIBUTES,
+      :budget_exceeded_behavior,
+      :default_tags,
+      :pricing_overrides,
+      :report_tag_breakdowns,
+      :storage_backend,
+      :storage_error_behavior,
+      :unknown_pricing_behavior,
+      :openai_compatible_providers
+    )
 
     def initialize
       @enabled = true
@@ -46,55 +57,74 @@ module LlmCostTracker
       @pricing_overrides  = {}
       @report_tag_breakdowns = []
       self.openai_compatible_providers = OPENAI_COMPATIBLE_PROVIDERS
+      @finalized = false
+    end
+
+    def default_tags=(value)
+      ensure_shared_configuration_mutable!
+      @default_tags = value
     end
 
     def openai_compatible_providers=(providers)
+      ensure_shared_configuration_mutable!
       @openai_compatible_providers = normalize_openai_compatible_providers(providers)
     end
 
-    def storage_backend=(value)
-      @storage_backend = normalize_enum(:storage_backend, value, STORAGE_BACKENDS, default: :log)
+    def pricing_overrides=(value)
+      ensure_shared_configuration_mutable!
+      @pricing_overrides = value
     end
 
-    def budget_exceeded_behavior=(value)
-      @budget_exceeded_behavior = normalize_enum(
-        :budget_exceeded_behavior,
-        value,
-        BUDGET_EXCEEDED_BEHAVIORS,
-        default: :notify
-      )
+    def report_tag_breakdowns=(value)
+      ensure_shared_configuration_mutable!
+      @report_tag_breakdowns = value
     end
 
-    def storage_error_behavior=(value)
-      @storage_error_behavior = normalize_enum(
-        :storage_error_behavior,
-        value,
-        STORAGE_ERROR_BEHAVIORS,
-        default: :warn
-      )
+    SHARED_SCALAR_ATTRIBUTES.each do |name|
+      define_method("#{name}=") do |value|
+        ensure_shared_configuration_mutable!
+        instance_variable_set(:"@#{name}", value)
+      end
     end
 
-    def unknown_pricing_behavior=(value)
-      @unknown_pricing_behavior = normalize_enum(
-        :unknown_pricing_behavior,
-        value,
-        UNKNOWN_PRICING_BEHAVIORS,
-        default: :warn
-      )
+    SHARED_ENUM_ATTRIBUTES.each do |name, (allowed, default)|
+      define_method("#{name}=") do |value|
+        ensure_shared_configuration_mutable!
+        instance_variable_set(:"@#{name}", normalize_enum(name, value, allowed, default: default))
+      end
     end
 
     def normalize_openai_compatible_providers!
       self.openai_compatible_providers = openai_compatible_providers
     end
 
-    def active_record?
-      storage_backend == :active_record
+    def finalize!
+      @default_tags = ValueHelpers.deep_freeze(@default_tags || {})
+      @pricing_overrides = ValueHelpers.deep_freeze(@pricing_overrides || {})
+      @report_tag_breakdowns = ValueHelpers.deep_freeze(Array(@report_tag_breakdowns))
+      @openai_compatible_providers = ValueHelpers.deep_freeze(@openai_compatible_providers || {})
+      @finalized = true
+      self
     end
 
-    def log?
-      storage_backend == :log
+    def finalized? = @finalized
+
+    def dup_for_configuration
+      copy = dup
+      copy.instance_variable_set(:@default_tags, ValueHelpers.deep_dup(@default_tags || {}))
+      copy.instance_variable_set(:@pricing_overrides, ValueHelpers.deep_dup(@pricing_overrides || {}))
+      copy.instance_variable_set(:@report_tag_breakdowns, ValueHelpers.deep_dup(@report_tag_breakdowns || []))
+      copy.instance_variable_set(
+        :@openai_compatible_providers,
+        ValueHelpers.deep_dup(@openai_compatible_providers || {})
+      )
+      copy.instance_variable_set(:@finalized, false)
+      copy
     end
 
+    def active_record? = storage_backend == :active_record
+    def log? = storage_backend == :log
+
     private
 
     def normalize_enum(name, value, allowed, default:)
@@ -110,5 +140,11 @@ module LlmCostTracker
         normalized[host.to_s.downcase] = provider.to_s
       end
     end
+
+    def ensure_shared_configuration_mutable!
+      return unless finalized?
+
+      raise FrozenError, "can't modify frozen LlmCostTracker::Configuration"
+    end
   end
 end

data/lib/llm_cost_tracker/event.rb

@@ -10,6 +10,8 @@ module LlmCostTracker
     :cost,
     :tags,
     :latency_ms,
+    :stream,
+    :usage_source,
     :tracked_at
   ) do
     def to_h

data/lib/llm_cost_tracker/generators/llm_cost_tracker/add_streaming_generator.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+require "rails/generators"
+require "rails/generators/active_record"
+
+module LlmCostTracker
+  module Generators
+    class AddStreamingGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+
+      source_root File.expand_path("templates", __dir__)
+
+      desc "Creates a migration to add llm_api_calls.stream and llm_api_calls.usage_source"
+
+      def create_migration_file
+        migration_template(
+          "add_streaming_to_llm_api_calls.rb.erb",
+          "db/migrate/add_streaming_to_llm_api_calls.rb"
+        )
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/add_streaming_to_llm_api_calls.rb.erb

@@ -0,0 +1,25 @@
+class AddStreamingToLlmApiCalls < ActiveRecord::Migration<%= migration_version %>
+  def up
+    unless column_exists?(:llm_api_calls, :stream)
+      add_column :llm_api_calls, :stream, :boolean, null: false, default: false
+      add_index  :llm_api_calls, :stream
+    end
+
+    unless column_exists?(:llm_api_calls, :usage_source)
+      add_column :llm_api_calls, :usage_source, :string
+      add_index  :llm_api_calls, :usage_source
+    end
+  end
+
+  def down
+    if column_exists?(:llm_api_calls, :usage_source)
+      remove_index  :llm_api_calls, :usage_source if index_exists?(:llm_api_calls, :usage_source)
+      remove_column :llm_api_calls, :usage_source
+    end
+
+    if column_exists?(:llm_api_calls, :stream)
+      remove_index  :llm_api_calls, :stream if index_exists?(:llm_api_calls, :stream)
+      remove_column :llm_api_calls, :stream
+    end
+  end
+end

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/create_llm_api_calls.rb.erb

@@ -10,6 +10,8 @@ class CreateLlmApiCalls < ActiveRecord::Migration<%= migration_version %>
       t.decimal :output_cost,   precision: 20, scale: 8
       t.decimal :total_cost,    precision: 20, scale: 8
       t.integer :latency_ms
+      t.boolean :stream,        null: false, default: false
+      t.string  :usage_source
       if postgresql?
         t.jsonb :tags, null: false, default: {}
       else
@@ -24,6 +26,8 @@ class CreateLlmApiCalls < ActiveRecord::Migration<%= migration_version %>
     add_index :llm_api_calls, :model
     add_index :llm_api_calls, :tracked_at
     add_index :llm_api_calls, [:provider, :tracked_at]
+    add_index :llm_api_calls, :stream
+    add_index :llm_api_calls, :usage_source
     add_index :llm_api_calls, :tags, using: :gin if postgresql?
   end
 

data/lib/llm_cost_tracker/generators/llm_cost_tracker/templates/llm_cost_tracker_prices.yml.erb

@@ -14,8 +14,11 @@
 #
 # Optional metadata keys, ignored by cost calculation:
 # - _source
+# - _source_version
+# - _fetched_at
 # - _updated
 # - _notes
+# - _validator_override
 #
 # Example: custom fine-tune
 # models:
@@ -30,7 +33,11 @@
 #   "gpt-4o":
 #     input: 2.00
 #     output: 8.00
-#     _source: "Enterprise agreement"
+#     _source: "manual"
 #     _updated: "2026-04-18"
+#
+# Use _source: "manual" for custom or orphaned entries you never want sync to touch.
+# Use _validator_override: ["skip_relative_change"] if a negotiated price would
+# otherwise trip the >3x sync warning.
 
 models:

data/lib/llm_cost_tracker/llm_api_call.rb

@@ -21,6 +21,14 @@ module LlmCostTracker
     scope :without_cost, -> { where(total_cost: nil) }
     scope :unknown_pricing, -> { without_cost }
     scope :with_latency, -> { latency_column? ? where.not(latency_ms: nil) : none }
+    scope :streaming,     -> { stream_column? ? where(stream: true) : none }
+    scope :non_streaming, -> { stream_column? ? where(stream: [false, nil]) : all }
+    scope :by_usage_source, ->(source) { usage_source_column? ? where(usage_source: source.to_s) : none }
+    scope :streaming_missing_usage, lambda {
+      return none unless stream_column? && usage_source_column?
+
+      where(stream: true).where(usage_source: ["unknown", nil])
+    }
 
     scope :with_json_tags, lambda {
       if tags_json_column?

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -18,22 +18,39 @@ module LlmCostTracker
 
         request_url  = request_env.url.to_s
         request_body = read_body(request_env.body) || ""
+        parser       = Parsers::Registry.find_for(request_url)
+        streaming    = parser&.streaming_request?(request_url, request_body)
+        stream_buffer = install_stream_tap(request_env) if streaming
 
-        enforce_budget_before_request(request_url)
+        Tracker.enforce_budget! if parser
         started_at = monotonic_time
 
         @app.call(request_env).on_complete do |response_env|
-          process(request_env, request_url, request_body, response_env, elapsed_ms(started_at))
+          process(
+            parser: parser,
+            request_env: request_env,
+            request_url: request_url,
+            request_body: request_body,
+            response_env: response_env,
+            latency_ms: elapsed_ms(started_at),
+            streaming: streaming,
+            stream_buffer: stream_buffer
+          )
         end
       end
 
       private
 
-      def process(request_env, request_url, request_body, response_env, latency_ms)
-        parser = Parsers::Registry.find_for(request_url)
+      def process(parser:, request_env:, request_url:, request_body:, response_env:,
+                  latency_ms:, streaming:, stream_buffer:)
         return unless parser
 
-        parsed = parse_response(parser, request_url, request_body, response_env)
+        parsed =
+          if streaming
+            parse_stream(parser, request_url, request_body, response_env, stream_buffer)
+          else
+            parse_response(parser, request_url, request_body, response_env)
+          end
         return unless parsed
 
         Tracker.record(
@@ -42,6 +59,8 @@ module LlmCostTracker
           input_tokens: parsed.input_tokens,
           output_tokens: parsed.output_tokens,
           latency_ms: latency_ms,
+          stream: parsed.stream,
+          usage_source: parsed.usage_source,
           metadata: resolved_tags(request_env).merge(parsed.metadata)
         )
       rescue LlmCostTracker::Error
@@ -54,7 +73,9 @@ module LlmCostTracker
         response_body = read_body(response_env.body)
         unless response_body
           Logging.warn(
-            "Unable to read response body for #{request_url}; streaming/SSE responses require manual tracking."
+            "Unable to read response body for #{request_url}; " \
+            "streaming responses are captured automatically for OpenAI/Anthropic/Gemini " \
+            "or via LlmCostTracker.track_stream for custom clients."
           )
           return nil
         end
@@ -62,10 +83,37 @@ module LlmCostTracker
         parser.parse(request_url, request_body, response_env.status, response_body)
       end
 
-      def enforce_budget_before_request(request_url)
-        return unless Parsers::Registry.find_for(request_url)
+      def parse_stream(parser, request_url, request_body, response_env, stream_buffer)
+        body = stream_buffer&.string
+        body = read_body(response_env.body) if body.nil? || body.empty?
+
+        if body.nil? || body.empty?
+          Logging.warn(
+            "Unable to capture streaming response for #{request_url}; " \
+            "fall back to LlmCostTracker.track_stream for manual capture."
+          )
+          return nil
+        end
+
+        events = Parsers::SSE.parse(body)
+        parser.parse_stream(request_url, request_body, response_env.status, events)
+      end
+
+      def install_stream_tap(request_env)
+        return nil unless request_env.respond_to?(:request) && request_env.request
 
-        Tracker.enforce_budget!
+        original = request_env.request.on_data
+        return nil unless original
+
+        buffer = StringIO.new
+        request_env.request.on_data = proc do |chunk, size, env|
+          buffer << chunk.to_s
+          original.call(chunk, size, env)
+        end
+        buffer
+      rescue StandardError => e
+        Logging.warn("Unable to install streaming tap: #{e.class}: #{e.message}")
+        nil
       end
 
       def read_body(body)

data/lib/llm_cost_tracker/parsed_usage.rb

@@ -10,11 +10,13 @@ module LlmCostTracker
     :cached_input_tokens,
     :cache_read_input_tokens,
     :cache_creation_input_tokens,
-    :reasoning_tokens
+    :reasoning_tokens,
+    :stream,
+    :usage_source
   )
 
   class ParsedUsage
-    TRACKING_KEYS = %i[provider model input_tokens output_tokens total_tokens].freeze
+    TRACKING_KEYS = %i[provider model input_tokens output_tokens total_tokens stream usage_source].freeze
 
     def self.build(**attributes)
       new(
@@ -26,7 +28,9 @@ module LlmCostTracker
         cached_input_tokens: attributes[:cached_input_tokens],
         cache_read_input_tokens: attributes[:cache_read_input_tokens],
         cache_creation_input_tokens: attributes[:cache_creation_input_tokens],
-        reasoning_tokens: attributes[:reasoning_tokens]
+        reasoning_tokens: attributes[:reasoning_tokens],
+        stream: attributes[:stream] || false,
+        usage_source: attributes[:usage_source]
       )
     end
 

data/lib/llm_cost_tracker/parsers/anthropic.rb

@@ -16,6 +16,10 @@ module LlmCostTracker
         false
       end
 
+      def provider_names
+        %w[anthropic]
+      end
+
       def parse(_request_url, request_body, response_status, response_body)
         return nil unless response_status == 200
 
@@ -33,9 +37,83 @@ module LlmCostTracker
           total_tokens: usage["input_tokens"].to_i + usage["output_tokens"].to_i +
             usage["cache_read_input_tokens"].to_i + usage["cache_creation_input_tokens"].to_i,
           cache_read_input_tokens: usage["cache_read_input_tokens"],
-          cache_creation_input_tokens: usage["cache_creation_input_tokens"]
+          cache_creation_input_tokens: usage["cache_creation_input_tokens"],
+          usage_source: :response
         )
       end
+
+      def parse_stream(_request_url, request_body, response_status, events)
+        return nil unless response_status == 200
+
+        request = safe_json_parse(request_body)
+        model = stream_model(events) || request["model"]
+        usage = stream_usage(events)
+
+        if usage
+          input = usage["input_tokens"].to_i
+          output = usage["output_tokens"].to_i
+          cache_read = usage["cache_read_input_tokens"].to_i
+          cache_creation = usage["cache_creation_input_tokens"].to_i
+
+          ParsedUsage.build(
+            provider: "anthropic",
+            model: model,
+            input_tokens: input,
+            output_tokens: output,
+            total_tokens: input + output + cache_read + cache_creation,
+            cache_read_input_tokens: usage["cache_read_input_tokens"],
+            cache_creation_input_tokens: usage["cache_creation_input_tokens"],
+            stream: true,
+            usage_source: :stream_final
+          )
+        else
+          ParsedUsage.build(
+            provider: "anthropic",
+            model: model,
+            input_tokens: 0,
+            output_tokens: 0,
+            total_tokens: 0,
+            stream: true,
+            usage_source: :unknown
+          )
+        end
+      end
+
+      private
+
+      def stream_usage(events)
+        start_usage = nil
+        latest_delta = nil
+
+        events.each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          case data["type"]
+          when "message_start"
+            start_usage = data.dig("message", "usage")
+          when "message_delta"
+            latest_delta = data["usage"] if data["usage"].is_a?(Hash)
+          end
+        end
+
+        return nil unless start_usage || latest_delta
+
+        (start_usage || {}).merge(latest_delta || {}) do |_key, start_val, delta_val|
+          delta_val.nil? ? start_val : delta_val
+        end
+      end
+
+      def stream_model(events)
+        events.each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          model = data.dig("message", "model")
+          return model if model && !model.empty?
+        end
+        nil
+      end
     end
   end
 end

data/lib/llm_cost_tracker/parsers/base.rb

@@ -5,19 +5,31 @@ require "json"
 module LlmCostTracker
   module Parsers
     class Base
-      # Parse a provider response into a {LlmCostTracker::ParsedUsage}, or return
-      # nil when the response is not trackable (non-200, missing usage, etc).
-      #
-      # @return [LlmCostTracker::ParsedUsage, nil]
       def parse(request_url, request_body, response_status, response_body)
         raise NotImplementedError
       end
 
-      # Returns true if this parser can handle the given URL.
+      def provider_names
+        []
+      end
+
       def match?(url)
         raise NotImplementedError
       end
 
+      def streaming_request?(_request_url, request_body)
+        return false if request_body.nil?
+
+        body = request_body.to_s
+        return false if body.empty?
+
+        body.include?('"stream":true') || body.include?('"stream": true') || body.include?("stream: true")
+      end
+
+      def parse_stream(_request_url, _request_body, _response_status, _events)
+        nil
+      end
+
       private
 
       def safe_json_parse(body)

data/lib/llm_cost_tracker/parsers/gemini.rb

@@ -9,6 +9,7 @@ module LlmCostTracker
     class Gemini < Base
       HOSTS = %w[generativelanguage.googleapis.com].freeze
       TRACKED_PATH_PATTERN = %r{/models/[^/:]+:(?:generateContent|streamGenerateContent)\z}
+      STREAM_PATH_PATTERN  = /:streamGenerateContent\z/
 
       def match?(url)
         uri = URI.parse(url.to_s)
@@ -17,6 +18,16 @@ module LlmCostTracker
         false
       end
 
+      def provider_names
+        %w[gemini]
+      end
+
+      def streaming_request?(request_url, request_body)
+        return true if streaming_url?(request_url)
+
+        super
+      end
+
       def parse(request_url, _request_body, response_status, response_body)
         return nil unless response_status == 200
 
@@ -24,31 +35,73 @@ module LlmCostTracker
         usage    = response["usageMetadata"]
         return nil unless usage
 
-        # Extract model from URL: /v1beta/models/gemini-2.5-flash:generateContent
+        build_parsed_usage(request_url, usage, usage_source: :response)
+      end
+
+      def parse_stream(request_url, _request_body, response_status, events)
+        return nil unless response_status == 200
+
+        usage = merged_stream_usage(events)
         model = extract_model_from_url(request_url)
 
+        if usage
+          build_parsed_usage(request_url, usage, stream: true, usage_source: :stream_final)
+        else
+          ParsedUsage.build(
+            provider: "gemini",
+            model: model,
+            input_tokens: 0,
+            output_tokens: 0,
+            total_tokens: 0,
+            stream: true,
+            usage_source: :unknown
+          )
+        end
+      end
+
+      private
+
+      def build_parsed_usage(request_url, usage, usage_source:, stream: false)
         ParsedUsage.build(
           provider: "gemini",
-          model: model,
+          model: extract_model_from_url(request_url),
           input_tokens: usage["promptTokenCount"].to_i,
           output_tokens: output_tokens(usage),
           total_tokens: usage["totalTokenCount"].to_i,
-          cached_input_tokens: usage["cachedContentTokenCount"]
+          cached_input_tokens: usage["cachedContentTokenCount"],
+          stream: stream,
+          usage_source: usage_source
         )
       end
 
-      private
+      def merged_stream_usage(events)
+        latest = nil
+        events.each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          meta = data["usageMetadata"]
+          latest = meta if meta.is_a?(Hash)
+        end
+        latest
+      end
 
       def output_tokens(usage)
         usage["candidatesTokenCount"].to_i + usage["thoughtsTokenCount"].to_i
       end
 
+      def streaming_url?(request_url)
+        URI.parse(request_url.to_s).path.match?(STREAM_PATH_PATTERN)
+      rescue URI::InvalidURIError
+        false
+      end
+
       def extract_model_from_url(url)
         uri = URI.parse(url.to_s)
         match = uri.path.match(%r{/models/([^/:]+)})
-        match ? match[1] : "unknown"
+        match && match[1]
       rescue URI::InvalidURIError
-        "unknown"
+        nil
       end
     end
   end

data/lib/llm_cost_tracker/parsers/openai.rb

@@ -20,10 +20,18 @@ module LlmCostTracker
         false
       end
 
+      def provider_names
+        %w[openai]
+      end
+
       def parse(request_url, request_body, response_status, response_body)
         parse_openai_usage(request_url, request_body, response_status, response_body)
       end
 
+      def parse_stream(request_url, request_body, response_status, events)
+        parse_openai_stream_usage(request_url, request_body, response_status, events)
+      end
+
       private
 
       def provider_for(_request_url)