llm_cost_tracker 0.2.0 → 0.3.1

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,97 @@ 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,
+          provider_response_id: response["responseId"]
+        )
+      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,
+            provider_response_id: stream_response_id(events)
+          )
+        else
+          ParsedUsage.build(
+            provider: "gemini",
+            provider_response_id: stream_response_id(events),
+            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, provider_response_id: nil)
         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,
+          provider_response_id: provider_response_id
         )
       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 stream_response_id(events)
+        events.each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          id = data["responseId"]
+          return id if id && !id.to_s.empty?
+        end
+        nil
+      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)

data/lib/llm_cost_tracker/parsers/openai_compatible.rb

@@ -19,10 +19,21 @@ module LlmCostTracker
         false
       end
 
+      def provider_names
+        [
+          "openai_compatible",
+          *LlmCostTracker.configuration.openai_compatible_providers.each_value.map(&:to_s)
+        ].uniq.freeze
+      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)
@@ -33,11 +44,7 @@ module LlmCostTracker
       end
 
       def provider_for_host(host)
-        configured_providers[host.to_s.downcase]&.to_s
-      end
-
-      def configured_providers
-        LlmCostTracker.configuration.openai_compatible_providers
+        LlmCostTracker.configuration.openai_compatible_providers[host.to_s.downcase]&.to_s
       end
 
       def tracked_path?(path)

data/lib/llm_cost_tracker/parsers/openai_usage.rb

@@ -16,14 +16,82 @@ module LlmCostTracker
 
         ParsedUsage.build(
           provider: provider_for(request_url),
+          provider_response_id: response["id"],
           model: response["model"] || request["model"],
           input_tokens: (usage["prompt_tokens"] || usage["input_tokens"]).to_i,
           output_tokens: (usage["completion_tokens"] || usage["output_tokens"]).to_i,
           total_tokens: usage["total_tokens"].to_i,
-          cached_input_tokens: cached_input_tokens(usage)
+          cached_input_tokens: cached_input_tokens(usage),
+          usage_source: :response
         )
       end
 
+      def parse_openai_stream_usage(request_url, request_body, response_status, events)
+        return nil unless response_status == 200
+
+        request = safe_json_parse(request_body)
+        model = detect_stream_model(events) || request["model"]
+        usage = detect_stream_usage(events)
+
+        if usage
+          ParsedUsage.build(
+            provider: provider_for(request_url),
+            provider_response_id: detect_stream_response_id(events),
+            model: model,
+            input_tokens: (usage["prompt_tokens"] || usage["input_tokens"]).to_i,
+            output_tokens: (usage["completion_tokens"] || usage["output_tokens"]).to_i,
+            total_tokens: usage["total_tokens"].to_i,
+            cached_input_tokens: cached_input_tokens(usage),
+            stream: true,
+            usage_source: :stream_final
+          )
+        else
+          ParsedUsage.build(
+            provider: provider_for(request_url),
+            provider_response_id: detect_stream_response_id(events),
+            model: model,
+            input_tokens: 0,
+            output_tokens: 0,
+            total_tokens: 0,
+            stream: true,
+            usage_source: :unknown
+          )
+        end
+      end
+
+      def detect_stream_usage(events)
+        events.reverse_each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          usage = data["usage"]
+          return usage if usage.is_a?(Hash) && !usage.empty?
+        end
+        nil
+      end
+
+      def detect_stream_model(events)
+        events.each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          model = data["model"]
+          return model if model && !model.to_s.empty?
+        end
+        nil
+      end
+
+      def detect_stream_response_id(events)
+        events.each do |event|
+          data = event[:data]
+          next unless data.is_a?(Hash)
+
+          id = data["id"] || data.dig("response", "id")
+          return id if id && !id.to_s.empty?
+        end
+        nil
+      end
+
       def cached_input_tokens(usage)
         details = usage["prompt_tokens_details"] || usage["input_tokens_details"] || {}
         details["cached_tokens"]

data/lib/llm_cost_tracker/parsers/registry.rb

@@ -1,23 +1,35 @@
 # frozen_string_literal: true
 
+require "monitor"
+
 module LlmCostTracker
   module Parsers
     class Registry
+      MUTEX = Monitor.new
+
       class << self
         def parsers
-          @parsers ||= default_parsers
+          @parsers || MUTEX.synchronize { @parsers ||= default_parsers.freeze }
         end
 
         def register(parser)
-          parsers.unshift(parser)
+          MUTEX.synchronize do
+            current = @parsers || default_parsers.freeze
+            @parsers = ([parser] + current).freeze
+          end
         end
 
         def find_for(url)
           parsers.find { |parser| parser.match?(url) }
         end
 
+        def find_for_provider(provider)
+          provider_name = provider.to_s
+          parsers.find { |parser| parser.provider_names.include?(provider_name) }
+        end
+
         def reset!
-          @parsers = nil
+          MUTEX.synchronize { @parsers = nil }
         end
 
         private

data/lib/llm_cost_tracker/parsers/sse.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+require "json"
+
+module LlmCostTracker
+  module Parsers
+    module SSE
+      DONE_MARKER = "[DONE]"
+
+      class << self
+        def parse(body)
+          return [] if body.nil? || body.empty?
+
+          return parse_json_array(body) if probably_json_array?(body)
+
+          parse_event_stream(body)
+        end
+
+        private
+
+        def parse_event_stream(body)
+          events = []
+          current_event = nil
+          data_lines = []
+
+          body.each_line do |raw|
+            line = raw.chomp
+
+            if line.empty?
+              events << finalize_event(current_event, data_lines) if data_lines.any?
+              current_event = nil
+              data_lines = []
+              next
+            end
+
+            next if line.start_with?(":")
+
+            field, _, value = line.partition(":")
+            value = value[1..] if value.start_with?(" ")
+
+            case field
+            when "event" then current_event = value
+            when "data"  then data_lines << value
+            end
+          end
+
+          events << finalize_event(current_event, data_lines) if data_lines.any?
+          events.compact
+        end
+
+        def parse_json_array(body)
+          parsed = JSON.parse(body)
+          return [] unless parsed.is_a?(Array)
+
+          parsed.map { |entry| { event: nil, data: entry } }
+        rescue JSON::ParserError
+          []
+        end
+
+        def finalize_event(event_name, data_lines)
+          payload = data_lines.join("\n")
+          return nil if payload == DONE_MARKER
+
+          { event: event_name, data: decode_data(payload) }
+        end
+
+        def decode_data(payload)
+          return payload if payload.empty?
+
+          JSON.parse(payload)
+        rescue JSON::ParserError
+          payload
+        end
+
+        def probably_json_array?(body)
+          body.lstrip.start_with?("[")
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/price_registry.rb

@@ -11,7 +11,7 @@ module LlmCostTracker
     DEFAULT_PRICES_PATH = File.expand_path("prices.json", __dir__)
     EMPTY_PRICES = {}.freeze
     PRICE_KEYS = %w[input cached_input output cache_read_input cache_creation_input].freeze
-    METADATA_KEYS = %w[_source _updated _notes].freeze
+    METADATA_KEYS = %w[_source _source_version _fetched_at _updated _notes _validator_override].freeze
     MUTEX = Monitor.new
 
     class << self
@@ -26,9 +26,7 @@ module LlmCostTracker
       end
 
       def normalize_price_table(table)
-        (table || {}).each_with_object({}) do |(model, price), normalized|
-          normalized[model.to_s] = normalize_price_entry(price)
-        end
+        normalize_price_entries(table, context: "price table")
       end
 
       def file_prices(path)
@@ -47,7 +45,7 @@ module LlmCostTracker
           @file_prices_cache = { key: cache_key, value: value }.freeze
           value
         end
-      rescue Errno::ENOENT, JSON::ParserError, Psych::Exception, ArgumentError, TypeError, NoMethodError => e
+      rescue Errno::ENOENT, JSON::ParserError, Psych::Exception, ArgumentError, TypeError => e
         raise Error, "Unable to load prices_file #{path.inspect}: #{e.message}"
       end
 
@@ -60,15 +58,23 @@ module LlmCostTracker
       end
 
       def normalize_price_entry(price)
-        (price || {}).each_with_object({}) do |(key, value), normalized|
+        price.each_with_object({}) do |(key, value), normalized|
           key = key.to_s
           normalized[key.to_sym] = Float(value) if PRICE_KEYS.include?(key)
         end
       end
 
       def normalize_file_prices(table, path:)
-        (table || {}).each_with_object({}) do |(model, price), normalized|
-          warn_unknown_keys(model, price, path)
+        normalize_price_entries(table, context: path)
+      end
+
+      def normalize_price_entries(table, context:)
+        table = {} if table.nil?
+        raise ArgumentError, "#{context} must be a hash of models" unless table.is_a?(Hash)
+
+        table.each_with_object({}) do |(model, price), normalized|
+          price = validate_price_entry(price, model: model, context: context)
+          warn_unknown_keys(model, price, context)
           normalized[model.to_s] = normalize_price_entry(price)
         end
       end
@@ -95,8 +101,17 @@ module LlmCostTracker
       end
 
       def price_file_models(registry)
+        raise ArgumentError, "prices_file must be a hash" unless registry.is_a?(Hash)
+
         registry.fetch("models", registry)
       end
+
+      def validate_price_entry(price, model:, context:)
+        return {} if price.nil?
+        return price if price.is_a?(Hash)
+
+        raise ArgumentError, "price entry for #{model.inspect} in #{context} must be a hash"
+      end
     end
   end
 end

data/lib/llm_cost_tracker/price_sync/fetcher.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+require "digest"
+require "net/http"
+require "time"
+require "uri"
+
+module LlmCostTracker
+  module PriceSync
+    class Fetcher
+      Response = Data.define(:body, :etag, :last_modified, :not_modified, :fetched_at) do
+        def source_version
+          etag || last_modified || Digest::SHA256.hexdigest(body.to_s)
+        end
+      end
+
+      USER_AGENT = "llm_cost_tracker price sync"
+      MAX_REDIRECTS = 5
+      OPEN_TIMEOUT = 5
+      READ_TIMEOUT = 10
+      WRITE_TIMEOUT = 10
+
+      def get(url, etag: nil, redirects: 0)
+        raise Error, "Too many redirects while fetching #{url}" if redirects > MAX_REDIRECTS
+
+        uri = URI.parse(url)
+        request = Net::HTTP::Get.new(uri)
+        request["User-Agent"] = USER_AGENT
+        request["If-None-Match"] = etag if etag
+
+        response = Net::HTTP.start(
+          uri.host,
+          uri.port,
+          use_ssl: uri.scheme == "https",
+          open_timeout: OPEN_TIMEOUT,
+          read_timeout: READ_TIMEOUT,
+          write_timeout: WRITE_TIMEOUT
+        ) do |http|
+          http.request(request)
+        end
+
+        case response
+        when Net::HTTPSuccess
+          build_response(response, not_modified: false)
+        when Net::HTTPNotModified
+          build_response(response, body: nil, not_modified: true)
+        when Net::HTTPRedirection
+          location = response["location"]
+          raise Error, "Redirect without location while fetching #{url}" if location.nil? || location.empty?
+
+          get(URI.join(url, location).to_s, etag: etag, redirects: redirects + 1)
+        else
+          raise Error, "Unable to fetch #{url}: HTTP #{response.code}"
+        end
+      rescue SocketError, SystemCallError, Timeout::Error => e
+        raise Error, "Unable to fetch #{url}: #{e.class}: #{e.message}"
+      end
+
+      private
+
+      def build_response(response, not_modified:, body: response.body)
+        Response.new(
+          body: body,
+          etag: response["etag"],
+          last_modified: response["last-modified"],
+          not_modified: not_modified,
+          fetched_at: Time.now.utc.iso8601
+        )
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/price_sync/merger.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module PriceSync
+    class Merger
+      Discrepancy = Data.define(:model, :field, :values)
+
+      PRIORITY_ORDER = %i[litellm openrouter].freeze
+      SUPPLEMENTAL_FIELDS = %i[cached_input cache_read_input cache_creation_input].freeze
+
+      def merge(results_by_source)
+        prices = collect_prices(results_by_source)
+        discrepancies = []
+
+        merged = prices.group_by(&:model).sort.to_h.transform_values do |candidates|
+          sorted = sort_candidates(candidates)
+          discrepancies.concat(detect_discrepancies(sorted))
+          fill_missing_fields(sorted.first, sorted.drop(1))
+        end
+
+        [merged, discrepancies]
+      end
+
+      private
+
+      def collect_prices(results_by_source)
+        results_by_source.flat_map do |source_name, result|
+          result.prices.map do |price|
+            price.with(source: source_name)
+          end
+        end
+      end
+
+      def sort_candidates(candidates)
+        candidates.sort_by do |price|
+          PRIORITY_ORDER.index(price.source.to_sym) || PRIORITY_ORDER.length
+        end
+      end
+
+      def fill_missing_fields(primary, fallbacks)
+        SUPPLEMENTAL_FIELDS.reduce(primary) do |current, field|
+          next current if current.public_send(field)
+
+          fallback = fallbacks.find { |candidate| candidate.public_send(field) }
+          fallback ? current.with(field => fallback.public_send(field)) : current
+        end
+      end
+
+      def detect_discrepancies(candidates)
+        return [] if candidates.length < 2
+
+        RawPrice::PRICE_FIELDS.filter_map do |field|
+          values = candidates.each_with_object({}) do |price, collected|
+            value = price.public_send(field)
+            collected[price.source] = value unless value.nil?
+          end
+          next if values.size < 2
+          next unless discrepant?(values.values)
+
+          Discrepancy.new(model: candidates.first.model, field: field, values: values)
+        end
+      end
+
+      def discrepant?(values)
+        min, max = values.minmax
+        return max != min if min.to_f.zero?
+
+        ((max - min).abs / min.to_f) >= 0.05
+      end
+    end
+  end
+end