llm_cost_tracker 0.6.1 → 0.7.1

data/lib/llm_cost_tracker/capture/stream_collector.rb

@@ -0,0 +1,182 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/object/blank"
+require "active_support/core_ext/object/deep_dup"
+
+require_relative "stream"
+
+module LlmCostTracker
+  module Capture
+    class StreamCollector
+      attr_reader :provider
+
+      def initialize(provider:, model:, latency_ms: nil, provider_response_id: nil, pricing_mode: nil, metadata: {})
+        @provider = provider.to_s
+        @model = model
+        @latency_ms = latency_ms
+        @provider_response_id = provider_response_id
+        @pricing_mode = pricing_mode
+        @metadata = (metadata || {}).deep_dup
+        @events = []
+        @captured_bytes = 0
+        @overflowed = false
+        @explicit_usage = nil
+        @started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        @finished = false
+        @mutex = Mutex.new
+      end
+
+      def model
+        @mutex.synchronize { @model }
+      end
+
+      def metadata
+        @mutex.synchronize { @metadata.deep_dup }
+      end
+
+      def provider_response_id
+        @mutex.synchronize { @provider_response_id }
+      end
+
+      def model=(value)
+        @mutex.synchronize do
+          ensure_open!
+          @model = value
+        end
+      end
+
+      def provider_response_id=(value)
+        @mutex.synchronize do
+          ensure_open!
+          @provider_response_id = value
+        end
+      end
+
+      def event(data, type: nil)
+        @mutex.synchronize do
+          ensure_open!
+          capture_event(data, type: type) unless data.nil?
+        end
+        self
+      end
+
+      def usage(input_tokens:, output_tokens:, **extra)
+        @mutex.synchronize do
+          ensure_open!
+          @provider_response_id = extra.delete(:provider_response_id) || @provider_response_id
+          @explicit_usage = TokenUsage.from_hash(extra.merge(
+                                                   input_tokens: input_tokens.to_i,
+                                                   output_tokens: output_tokens.to_i
+                                                 ))
+        end
+        self
+      end
+
+      def finish!(errored: false)
+        snapshot = @mutex.synchronize do
+          return if @finished
+
+          @finished = true
+          {
+            events: @events.dup,
+            overflowed: @overflowed,
+            explicit_usage: @explicit_usage,
+            model: @model,
+            latency_ms: @latency_ms,
+            provider_response_id: @provider_response_id,
+            pricing_mode: @pricing_mode,
+            metadata: @metadata.deep_dup
+          }
+        end
+
+        capture = build_usage_capture(snapshot)
+        provider_response_id = capture.provider_response_id || snapshot[:provider_response_id]
+        capture = capture.with(provider_response_id: provider_response_id)
+
+        Tracker.record(
+          capture: capture,
+          latency_ms: snapshot[:latency_ms] ||
+            ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - @started_at) * 1000).round,
+          pricing_mode: snapshot[:pricing_mode],
+          metadata: (errored ? { stream_errored: true } : {}).merge(snapshot[:metadata])
+        )
+      end
+
+      private
+
+      def ensure_open!
+        return unless @finished
+
+        raise FrozenError, "can't modify finished LlmCostTracker::Capture::StreamCollector"
+      end
+
+      def build_usage_capture(snapshot)
+        return build_from_explicit_usage(snapshot) if snapshot[:explicit_usage]
+        return build_unknown_usage(snapshot) if snapshot[:overflowed]
+
+        capture = Parsers.find_for_provider(@provider)&.parse_stream(nil, nil, 200, snapshot[:events])
+        if capture
+          model = present_model(capture.model) || present_model(snapshot[:model]) || UsageCapture::UNKNOWN_MODEL
+          return capture.with(provider: @provider, model: model)
+        end
+
+        build_unknown_usage(snapshot)
+      end
+
+      def present_model(value)
+        return nil if value.nil?
+
+        string = value.to_s.presence
+        return nil if string.nil? || string == "unknown"
+
+        string
+      end
+
+      def build_from_explicit_usage(snapshot)
+        UsageCapture.build(
+          provider: @provider,
+          model: snapshot[:model] || UsageCapture::UNKNOWN_MODEL,
+          token_usage: snapshot[:explicit_usage],
+          stream: true,
+          usage_source: :manual
+        )
+      end
+
+      def build_unknown_usage(snapshot)
+        UsageCapture.build(
+          provider: @provider,
+          model: snapshot[:model] || UsageCapture::UNKNOWN_MODEL,
+          token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
+          stream: true,
+          usage_source: :unknown
+        )
+      end
+
+      def capture_event(data, type:)
+        size = type.to_s.bytesize + estimated_bytes(data) + 32
+        if @captured_bytes + size <= Capture::Stream::LIMIT_BYTES
+          @events << { event: type, data: data.deep_dup }
+          @captured_bytes += size
+        else
+          @overflowed = true
+          @events.clear
+        end
+      end
+
+      def estimated_bytes(value)
+        case value
+        when Hash
+          value.sum { |key, nested| estimated_bytes(key) + estimated_bytes(nested) + 4 }
+        when Array
+          value.sum { |nested| estimated_bytes(nested) + 2 }
+        when String
+          value.bytesize + 2
+        when Numeric, true, false, nil
+          value.to_s.bytesize
+        else
+          value.to_s.bytesize + 2
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/{integrations → capture}/stream_tracker.rb

@@ -1,17 +1,13 @@
 # frozen_string_literal: true
 
-require "monitor"
+require "active_support/core_ext/object/deep_dup"
+require "active_support/core_ext/object/try"
 
 require_relative "../logging"
-require_relative "../stream_collector"
-require_relative "../value_helpers"
-require_relative "object_reader"
 
 module LlmCostTracker
-  module Integrations
+  module Capture
     class StreamTracker
-      def self.wrap(stream, collector:, active:, finish: nil) = new(stream, collector, active, finish).wrap
-
       def initialize(stream, collector, active, finish)
         @stream = stream
         @collector = collector
@@ -19,13 +15,22 @@ module LlmCostTracker
         @finish = finish || proc { |errored:| @collector.finish!(errored: errored) }
         @finished = false
         @capture_failed = false
-        @monitor = Monitor.new
+        @mutex = Mutex.new
       end
 
       def wrap
         return @stream unless @stream
 
-        iterator_wrapped = @stream.instance_variable_defined?(:@iterator) && wrap_iterator?
+        iterator_wrapped = false
+        if @stream.instance_variable_defined?(:@iterator)
+          iterator = @stream.instance_variable_get(:@iterator)
+          if iterator.respond_to?(:each)
+            @stream.instance_variable_set(:@iterator, Enumerator.new do |yielder|
+              each_from(iterator) { |event| yielder << event }
+            end)
+            iterator_wrapped = true
+          end
+        end
         wrap_each if !iterator_wrapped && @stream.respond_to?(:each)
 
         @stream
@@ -36,14 +41,6 @@ module LlmCostTracker
 
       private
 
-      def wrap_iterator?
-        iterator = @stream.instance_variable_get(:@iterator)
-        return false unless iterator.respond_to?(:each)
-
-        @stream.instance_variable_set(:@iterator, tracked_iterator(iterator))
-        true
-      end
-
       def wrap_each
         tracker = self
         original_each = @stream.method(:each)
@@ -54,17 +51,18 @@ module LlmCostTracker
         end
       end
 
-      def tracked_iterator(iterator)
-        Enumerator.new do |yielder|
-          each_from(iterator) { |event| yielder << event }
-        end
-      end
-
       def each_from(iterable)
         errored = false
-        iterate(iterable) do |event|
-          capture(event)
-          yield event
+        if iterable.respond_to?(:each)
+          iterable.each do |event|
+            capture(event)
+            yield event
+          end
+        else
+          iterable.call do |event|
+            capture(event)
+            yield event
+          end
         end
       rescue StandardError
         errored = true
@@ -73,41 +71,17 @@ module LlmCostTracker
         finish!(errored: errored)
       end
 
-      def iterate(iterable, &)
-        if iterable.respond_to?(:each)
-          iterable.each(&)
-        else
-          iterable.call(&)
-        end
-      end
-
       def capture(event)
-        payload = normalize(event_payload(event))
-        @collector.event(payload, type: event_type(event, payload))
-      rescue StandardError => e
-        warn_capture_failure(e)
-      end
-
-      def event_payload(event)
-        if event.respond_to?(:deep_to_h)
-          event.deep_to_h
-        elsif event.respond_to?(:to_h)
-          event.to_h
-        else
-          event_attributes(event)
-        end
-      end
-
-      def event_attributes(event)
-        %i[type id model usage response message].each_with_object({}) do |key, attributes|
-          value = ObjectReader.read(event, key)
+        raw_payload = event.try(:deep_to_h) || event.try(:to_h)
+        raw_payload ||= %i[type id model usage response message].each_with_object({}) do |key, attributes|
+          value = event.try(key)
           attributes[key] = value unless value.nil?
         end
-      end
-
-      def event_type(event, payload)
-        value = ObjectReader.first(event, :type) || payload["type"]
-        value&.to_s
+        payload = normalize(raw_payload)
+        type = event.try(:type) || payload["type"]
+        @collector.event(payload, type: type&.to_s)
+      rescue StandardError => e
+        warn_capture_failure(e)
       end
 
       def normalize(value)
@@ -123,23 +97,17 @@ module LlmCostTracker
         when NilClass
           nil
         else
-          converted = object_hash(value)
-          converted ? normalize(converted) : ValueHelpers.deep_dup(value)
-        end
-      end
-
-      def object_hash(value)
-        if value.respond_to?(:deep_to_h)
-          value.deep_to_h
-        elsif value.respond_to?(:to_h)
-          value.to_h
+          converted = begin
+            value.try(:deep_to_h) || value.try(:to_h)
+          rescue StandardError
+            nil
+          end
+          converted ? normalize(converted) : value.deep_dup
         end
-      rescue StandardError
-        nil
       end
 
       def warn_capture_failure(error)
-        should_warn = @monitor.synchronize do
+        should_warn = @mutex.synchronize do
           next false if @capture_failed
 
           @capture_failed = true
@@ -151,7 +119,7 @@ module LlmCostTracker
       end
 
       def finish!(errored:)
-        should_finish = @monitor.synchronize do
+        should_finish = @mutex.synchronize do
           next false if @finished
 
           @finished = true

data/lib/llm_cost_tracker/configuration/instrumentation.rb

@@ -16,7 +16,7 @@ module LlmCostTracker
     def normalize_instrumentation_names(names)
       names.flatten.flat_map do |name|
         key = name.to_sym
-        next available_instrumentation_names if key == :all
+        next Integrations.names if key == :all
 
         validate_instrumentation_name!(key)
         key
@@ -24,14 +24,10 @@ module LlmCostTracker
     end
 
     def validate_instrumentation_name!(name)
-      return if available_instrumentation_names.include?(name)
+      return if Integrations.names.include?(name)
 
       raise Error, "Unknown integration: #{name.inspect}. " \
-                   "Use one of: #{available_instrumentation_names.join(', ')}"
-    end
-
-    def available_instrumentation_names
-      Integrations::Registry.names
+                   "Use one of: #{Integrations.names.join(', ')}"
     end
   end
 end

data/lib/llm_cost_tracker/configuration.rb

@@ -1,27 +1,21 @@
 # frozen_string_literal: true
 
 require_relative "errors"
-require_relative "tag_key"
-require_relative "value_helpers"
+require_relative "tags/key"
 require_relative "configuration/instrumentation"
-require_relative "configuration/storage_backend"
 
 module LlmCostTracker
   class Configuration
     include ConfigurationInstrumentation
-    include ConfigurationStorageBackend
 
     OPENAI_COMPATIBLE_PROVIDERS = { "openrouter.ai" => "openrouter", "api.deepseek.com" => "deepseek" }.freeze
 
     BUDGET_EXCEEDED_BEHAVIORS = %i[notify raise block_requests].freeze
-    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 daily_budget per_call_budget
-                                  log_level prices_file max_tag_count max_tag_value_bytesize].freeze
+    SHARED_SCALAR_ATTRIBUTES = %i[enabled on_budget_exceeded monthly_budget daily_budget per_call_budget log_level
+                                  prices_file max_tag_count max_tag_value_bytesize].freeze
     SHARED_ENUM_ATTRIBUTES = {
       budget_exceeded_behavior: [BUDGET_EXCEEDED_BEHAVIORS, :notify],
-      storage_error_behavior: [STORAGE_ERROR_BEHAVIORS, :warn],
       unknown_pricing_behavior: [UNKNOWN_PRICING_BEHAVIORS, :warn]
     }.freeze
     DEFAULT_REDACTED_TAG_KEYS = %w[api_key access_token authorization credential password refresh_token secret].freeze
@@ -34,23 +28,18 @@ module LlmCostTracker
       :instrumented_integrations,
       :report_tag_breakdowns,
       :redacted_tag_keys,
-      :storage_backend,
-      :storage_error_behavior,
       :unknown_pricing_behavior,
       :openai_compatible_providers
     )
 
     def initialize
       @enabled = true
-      self.storage_backend = :log
-      @custom_storage     = nil
       @default_tags       = {}
       @on_budget_exceeded = nil
       @monthly_budget     = nil
       @daily_budget       = nil
       @per_call_budget    = nil
       self.budget_exceeded_behavior = :notify
-      self.storage_error_behavior = :warn
       self.unknown_pricing_behavior = :warn
       @log_level          = :info
       @prices_file        = nil
@@ -81,7 +70,7 @@ module LlmCostTracker
 
     def report_tag_breakdowns=(value)
       ensure_shared_configuration_mutable!
-      @report_tag_breakdowns = normalize_report_tag_breakdowns(value)
+      @report_tag_breakdowns = Array(value).map { |key| Tags::Key.validate!(key, error_class: Error) }
     end
 
     def redacted_tag_keys=(value)
@@ -104,38 +93,20 @@ module LlmCostTracker
     end
 
     def finalize!
-      @default_tags = ValueHelpers.deep_freeze(@default_tags || {})
-      @pricing_overrides = ValueHelpers.deep_freeze(@pricing_overrides || {})
-      @instrumented_integrations = ValueHelpers.deep_freeze(@instrumented_integrations || [])
-      @report_tag_breakdowns = ValueHelpers.deep_freeze(Array(@report_tag_breakdowns))
-      @redacted_tag_keys = ValueHelpers.deep_freeze(Array(@redacted_tag_keys))
-      @openai_compatible_providers = ValueHelpers.deep_freeze(@openai_compatible_providers || {})
+      @default_tags = deep_freeze(@default_tags || {})
+      @pricing_overrides = deep_freeze(@pricing_overrides || {})
+      @instrumented_integrations = deep_freeze(@instrumented_integrations || [])
+      @report_tag_breakdowns = deep_freeze(Array(@report_tag_breakdowns))
+      @redacted_tag_keys = deep_freeze(Array(@redacted_tag_keys))
+      @openai_compatible_providers = deep_freeze(@openai_compatible_providers || {})
       @finalized = true
       self
     end
 
-    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(
-        :@instrumented_integrations,
-        ValueHelpers.deep_dup(@instrumented_integrations || [])
-      )
-      copy.instance_variable_set(:@report_tag_breakdowns, ValueHelpers.deep_dup(@report_tag_breakdowns || []))
-      copy.instance_variable_set(:@redacted_tag_keys, ValueHelpers.deep_dup(@redacted_tag_keys || []))
-      copy.instance_variable_set(
-        :@openai_compatible_providers,
-        ValueHelpers.deep_dup(@openai_compatible_providers || {})
-      )
-      copy.instance_variable_set(:@finalized, false)
-      copy
+    def finalized?
+      @finalized
     end
 
-    def active_record? = storage_backend == :active_record
-
     private
 
     def normalize_enum(name, value, allowed, default:)
@@ -152,14 +123,28 @@ module LlmCostTracker
       end
     end
 
-    def normalize_report_tag_breakdowns(value)
-      Array(value).map { |key| TagKey.validate!(key, error_class: Error) }
-    end
-
     def ensure_shared_configuration_mutable!
       return unless finalized?
 
       raise FrozenError, "can't modify frozen LlmCostTracker::Configuration"
     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

data/lib/llm_cost_tracker/doctor/capture_verifier.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require_relative "check"
+require_relative "../ingestion"
+
+module LlmCostTracker
+  class Doctor
+    class CaptureVerifier
+      class << self
+        def call
+          new.checks
+        end
+
+        def report(checks = call)
+          (["LLM Cost Tracker capture verification"] + checks.map do |check|
+            "[#{check.status}] #{check.name}: #{check.message}"
+          end).join("\n")
+        end
+
+        def healthy?(checks = call)
+          checks.none? { |check| check.status == :error }
+        end
+      end
+
+      def checks
+        [
+          enabled_check,
+          *integration_checks,
+          *storage_checks
+        ].compact
+      end
+
+      private
+
+      def enabled_check
+        return Check.new(:ok, "tracking", "enabled") if LlmCostTracker.configuration.enabled
+
+        Check.new(:error, "tracking", "disabled; set config.enabled = true before verifying capture")
+      end
+
+      def integration_checks
+        enabled = LlmCostTracker.configuration.instrumented_integrations
+        if enabled.empty?
+          return [
+            Check.new(:ok, "sdk integrations", "none enabled; Faraday middleware and manual capture remain available")
+          ]
+        end
+
+        LlmCostTracker::Integrations.checks.map do |check|
+          Check.new(check.status, "sdk integration #{check.name}", check.message)
+        end
+      end
+
+      def storage_checks
+        LlmCostTracker::Ingestion.verify
+      rescue LlmCostTracker::Error => e
+        [Check.new(:error, "storage", e.message)]
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/doctor/check.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  class Doctor
+    Check = Data.define(:status, :name, :message)
+  end
+end