llm_cost_tracker 0.2.0 → 0.3.0

data/lib/llm_cost_tracker/price_sync/validator.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module PriceSync
+    class Validator
+      Result = Data.define(:accepted, :rejected, :flagged)
+      Issue = Data.define(:model, :reason, :old_price, :new_price)
+
+      MAX_INPUT_PER_MILLION = 100.0
+      MAX_OUTPUT_PER_MILLION = 500.0
+      MAX_RELATIVE_CHANGE = 3.0
+
+      def validate_batch(merged_prices, existing_registry:)
+        merged_prices.each_with_object(Result.new(accepted: {}, rejected: [], flagged: [])) do |(model, price), result|
+          old_price = normalize_entry(existing_registry[model])
+          status, reason = validate(new_price: price, old_price: old_price)
+
+          case status
+          when :rejected
+            result.rejected << Issue.new(model: model, reason: reason, old_price: old_price, new_price: price)
+          when :flagged
+            result.flagged << Issue.new(model: model, reason: reason, old_price: old_price, new_price: price)
+            result.accepted[model] = price
+          else
+            result.accepted[model] = price
+          end
+        end
+      end
+
+      private
+
+      def validate(new_price:, old_price:)
+        overrides = Array(normalize_entry(old_price)["_validator_override"])
+
+        return [:rejected, "input > $#{MAX_INPUT_PER_MILLION}/1M"] if new_price.input > MAX_INPUT_PER_MILLION
+        return [:rejected, "output > $#{MAX_OUTPUT_PER_MILLION}/1M"] if new_price.output > MAX_OUTPUT_PER_MILLION
+        return [:ok, nil] if overrides.include?("skip_relative_change")
+
+        if old_price.any? && changed_too_much?(old_price, new_price)
+          return [:flagged, "price changed >#{MAX_RELATIVE_CHANGE}x"]
+        end
+
+        [:ok, nil]
+      end
+
+      def changed_too_much?(old_price, new_price)
+        %i[input output].any? do |field|
+          old_value = old_price[field.to_s].to_f
+          next false if old_value.zero?
+
+          new_value = new_price.public_send(field).to_f
+          next false if new_value.zero?
+
+          ratio = [new_value / old_value, old_value / new_value].max
+          ratio > MAX_RELATIVE_CHANGE
+        end
+      end
+
+      def normalize_entry(entry)
+        (entry || {}).each_with_object({}) do |(key, value), normalized|
+          normalized[key.to_s] = value
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/price_sync.rb

@@ -0,0 +1,310 @@
+# frozen_string_literal: true
+
+require "date"
+require "fileutils"
+require "json"
+require "yaml"
+
+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/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
+  # rubocop:disable Metrics/ModuleLength, Metrics/ClassLength
+  module PriceSync
+    DEFAULT_OUTPUT_PATH = PriceRegistry::DEFAULT_PRICES_PATH
+    YAML_EXTENSIONS = %w[.yml .yaml].freeze
+
+    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 = build_refresh_plan(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?
+        write_registry(plan.path, 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 = build_refresh_plan(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 build_refresh_plan(path:, seed_path:, fetcher:, today:)
+        path = path.to_s
+        registry = load_registry(path, seed_path: seed_path)
+        current_models = registry.fetch("models", {})
+        source_results, failed_sources = fetch_all(current_models, fetcher)
+        merged, discrepancies = Merger.new.merge(source_results)
+        validated = Validator.new.validate_batch(merged, existing_registry: current_models)
+        updated_models = apply_changes(current_models, validated.accepted, today)
+        refresh_succeeded = source_results.any? { |_source, result| result.prices.any? }
+
+        RefreshPlan.new(
+          path: path,
+          registry: registry,
+          updated_registry: registry.merge(
+            "metadata" => updated_metadata(
+              registry["metadata"],
+              today,
+              refresh_succeeded: refresh_succeeded,
+              source_results: source_results
+            ),
+            "models" => updated_models
+          ),
+          accepted: validated.accepted,
+          changes: price_changes(current_models, updated_models),
+          orphaned_models: compute_orphaned(current_models, merged.keys),
+          failed_sources: failed_sources,
+          discrepancies: discrepancies,
+          rejected: validated.rejected,
+          flagged: validated.flagged,
+          sources_used: source_usage(source_results),
+          source_results: source_results
+        )
+      end
+
+      def fetch_all(current_models, fetcher)
+        results = {}
+        failures = {}
+
+        sources.each do |source|
+          results[source.name.to_sym] = source.fetch(current_models: current_models, fetcher: fetcher)
+        rescue Error => e
+          failures[source.name.to_sym] = e.message
+        end
+
+        [results, failures]
+      end
+
+      def apply_changes(current_models, accepted, today)
+        merged = seed_models(current_models)
+
+        accepted.each do |model, price|
+          next if manual_model?(merged[model])
+
+          merged[model] = registry_entry_for(merged[model], price, today)
+        end
+
+        merged.sort.to_h
+      end
+
+      def compute_orphaned(current_models, merged_models)
+        seed_models(current_models).keys.reject do |model|
+          manual_model?(current_models[model]) || merged_models.include?(model)
+        end.sort
+      end
+
+      def load_registry(path, seed_path:)
+        source_path = File.exist?(path) ? path : seed_path.to_s
+        normalize_registry(load_registry_file(source_path))
+      rescue Errno::ENOENT, JSON::ParserError, Psych::Exception, ArgumentError, TypeError, NoMethodError => e
+        raise Error, "Unable to load pricing registry #{source_path.inspect}: #{e.message}"
+      end
+
+      def load_registry_file(path)
+        contents = File.read(path)
+        return YAML.safe_load(contents, aliases: false) || {} if yaml_file?(path)
+
+        JSON.parse(contents)
+      end
+
+      def normalize_registry(registry)
+        {
+          "metadata" => normalize_hash(registry.fetch("metadata", {})),
+          "models" => normalize_models(registry.fetch("models", {}))
+        }
+      end
+
+      def normalize_models(models)
+        (models || {}).each_with_object({}) do |(model, entry), normalized|
+          normalized[model.to_s] = normalize_hash(entry)
+        end
+      end
+
+      def normalize_hash(hash)
+        (hash || {}).each_with_object({}) do |(key, value), normalized|
+          normalized[key.to_s] = value
+        end
+      end
+
+      def seed_models(current_models)
+        normalize_models(current_models).transform_values do |entry|
+          next entry if entry.key?("_source")
+
+          entry.merge("_source" => "seed")
+        end
+      end
+
+      def manual_model?(entry)
+        normalize_hash(entry)["_source"] == "manual"
+      end
+
+      def registry_entry_for(existing_entry, price, today)
+        normalize_hash(existing_entry)
+          .except(*PriceRegistry::PRICE_KEYS)
+          .merge(price.to_registry_entry(today: today))
+      end
+
+      def updated_metadata(existing, today, refresh_succeeded:, source_results:)
+        metadata = normalize_hash(existing)
+        metadata["currency"] ||= "USD"
+        metadata["unit"] ||= "1M tokens"
+        return metadata unless refresh_succeeded
+
+        metadata["updated_at"] = today.iso8601
+        metadata["source_urls"] = source_urls(source_results)
+        metadata
+      end
+
+      def source_usage(source_results)
+        source_results.transform_values do |result|
+          SourceUsage.new(prices_count: result.prices.size, source_version: result.source_version)
+        end
+      end
+
+      def price_changes(current_models, updated_models)
+        current_models = normalize_models(current_models)
+        updated_models = normalize_models(updated_models)
+
+        (current_models.keys | updated_models.keys).sort.each_with_object({}) do |model, changes|
+          fields = price_field_changes(current_models[model], updated_models[model])
+          changes[model] = fields if fields.any?
+        end
+      end
+
+      def price_field_changes(current_entry, updated_entry)
+        current_price = comparable_price(current_entry)
+        updated_price = comparable_price(updated_entry)
+
+        (current_price.keys | updated_price.keys).sort.each_with_object({}) do |field, changes|
+          from = current_price[field]
+          to = updated_price[field]
+          next if from == to
+
+          changes[field] = { "from" => from, "to" => to }
+        end
+      end
+
+      def comparable_price(entry)
+        normalize_hash(entry).slice(*PriceRegistry::PRICE_KEYS)
+      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
+
+      def source_urls(source_results)
+        names = source_results.keys.map(&:to_sym)
+        sources.select { |source| names.include?(source.name.to_sym) }.map(&:url)
+      end
+
+      def write_registry(path, registry)
+        FileUtils.mkdir_p(File.dirname(path))
+        payload = yaml_file?(path) ? YAML.dump(registry) : "#{JSON.pretty_generate(registry)}\n"
+        File.write(path, payload)
+      end
+
+      def yaml_file?(path)
+        YAML_EXTENSIONS.include?(File.extname(path).downcase)
+      end
+    end
+  end
+  # rubocop:enable Metrics/ModuleLength, Metrics/ClassLength
+end

data/lib/llm_cost_tracker/storage/active_record_store.rb

@@ -19,7 +19,9 @@ module LlmCostTracker
             tags:          tags_for_storage(tags),
             tracked_at:    event.tracked_at
           }
-          attributes[:latency_ms] = event.latency_ms if model_class.latency_column?
+          attributes[:latency_ms]   = event.latency_ms   if model_class.latency_column?
+          attributes[:stream]       = event.stream       if model_class.stream_column?
+          attributes[:usage_source] = event.usage_source if model_class.usage_source_column?
 
           model_class.create!(attributes)
         end

data/lib/llm_cost_tracker/stream_collector.rb

@@ -0,0 +1,158 @@
+# frozen_string_literal: true
+
+require "monitor"
+
+require_relative "value_helpers"
+
+module LlmCostTracker
+  class StreamCollector
+    attr_reader :provider
+
+    def initialize(provider:, model:, latency_ms: nil, metadata: {})
+      @provider = provider.to_s
+      @model = model
+      @latency_ms = latency_ms
+      @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 }
+    end
+
+    def metadata
+      @monitor.synchronize { ValueHelpers.deep_dup(@metadata) }
+    end
+
+    def model=(value)
+      @monitor.synchronize do
+        ensure_open!
+        @model = 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,
+          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,
+        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 } : {}
+    end
+
+    def elapsed_ms
+      ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - @started_at) * 1000).round
+    end
+  end
+end

data/lib/llm_cost_tracker/tags_column.rb

@@ -24,5 +24,13 @@ 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
   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:,
+                 metadata: {}, latency_ms: nil, stream: false, usage_source: nil)
         usage = EventMetadata.usage_data(input_tokens, output_tokens, metadata)
 
         cost_data = Pricing.cost_for(
@@ -43,13 +37,13 @@ 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),
           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 +71,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 +115,13 @@ 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
     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.0"
 end