llm_cost_tracker 0.2.0.alpha2 → 0.3.0

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/pricing.rb

@@ -1,9 +1,12 @@
 # frozen_string_literal: true
 
+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.
@@ -59,9 +62,14 @@ module LlmCostTracker
         cached = @prices_cache
         return cached[:value] if cached && cached[:key] == cache_key
 
-        value = PRICES.merge(file_prices).merge(overrides).freeze
-        @prices_cache = { key: cache_key, value: value }.freeze
-        value
+        MUTEX.synchronize do
+          cached = @prices_cache
+          return cached[:value] if cached && cached[:key] == cache_key
+
+          value = PRICES.merge(file_prices).merge(overrides).freeze
+          @prices_cache = { key: cache_key, value: value }.freeze
+          value
+        end
       end
 
       private
@@ -116,9 +124,14 @@ module LlmCostTracker
         cached = @sorted_price_keys_cache
         return cached[:keys] if cached && cached[:table].equal?(table)
 
-        keys = table.keys.sort_by { |key| -key.length }
-        @sorted_price_keys_cache = { table: table, keys: keys }.freeze
-        keys
+        MUTEX.synchronize do
+          cached = @sorted_price_keys_cache
+          return cached[:keys] if cached && cached[:table].equal?(table)
+
+          keys = table.keys.sort_by { |key| -key.length }
+          @sorted_price_keys_cache = { table: table, keys: keys }.freeze
+          keys
+        end
       end
     end
   end

data/lib/llm_cost_tracker/retention.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Retention
+    DEFAULT_BATCH_SIZE = 5_000
+
+    class << self
+      def prune(older_than:, batch_size: DEFAULT_BATCH_SIZE, now: Time.now.utc)
+        cutoff = resolve_cutoff(older_than, now)
+        require_relative "llm_api_call" unless defined?(LlmCostTracker::LlmApiCall)
+
+        deleted = 0
+        loop do
+          batch = LlmCostTracker::LlmApiCall.where(tracked_at: ...cutoff).limit(batch_size).delete_all
+          deleted += batch
+          break if batch < batch_size
+        end
+        deleted
+      end
+
+      private
+
+      def resolve_cutoff(older_than, now)
+        case older_than
+        when Time, DateTime then older_than.utc
+        when ActiveSupport::Duration then now - older_than
+        when Integer then now - (older_than * 86_400)
+        else
+          raise ArgumentError, "older_than must be a Duration, Time, or Integer days: #{older_than.inspect}"
+        end
+      end
+    end
+  end
+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/tag_query.rb

@@ -9,7 +9,8 @@ module LlmCostTracker
         normalized_tags = normalize_tags(tags)
         return model.all if normalized_tags.empty?
 
-        return json_query(model, normalized_tags) if model.tags_json_column?
+        return postgres_json_query(model, normalized_tags) if model.tags_jsonb_column?
+        return mysql_json_query(model, normalized_tags) if model.tags_mysql_json_column?
 
         text_query(model, normalized_tags)
       end
@@ -20,10 +21,14 @@ module LlmCostTracker
 
       private
 
-      def json_query(model, tags)
+      def postgres_json_query(model, tags)
         model.where("tags @> ?::jsonb", tags.to_json)
       end
 
+      def mysql_json_query(model, tags)
+        model.where("JSON_CONTAINS(tags, ?)", tags.to_json)
+      end
+
       def text_query(model, tags)
         tags.reduce(model.all) do |relation, (key, value)|
           relation.where("tags LIKE ? ESCAPE '\\'", "%#{model.sanitize_sql_like(json_tag_fragment(key, value))}%")

data/lib/llm_cost_tracker/tags_column.rb

@@ -3,14 +3,34 @@
 module LlmCostTracker
   module TagsColumn
     def tags_json_column?
+      tags_jsonb_column? || tags_mysql_json_column?
+    end
+
+    def tags_jsonb_column?
       column = columns_hash["tags"]
       return false unless column
 
-      %i[json jsonb].include?(column.type) || column.sql_type.to_s.downcase == "jsonb"
+      column.type == :jsonb || column.sql_type.to_s.downcase == "jsonb"
+    end
+
+    def tags_mysql_json_column?
+      column = columns_hash["tags"]
+      return false unless column
+      return false if tags_jsonb_column?
+
+      column.type == :json && connection.adapter_name.match?(/mysql/i)
     end
 
     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