llm_cost_tracker 0.9.0 → 0.10.0

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -18,16 +18,25 @@ module LlmCostTracker
       end
 
       def call(request_env)
-        return @app.call(request_env) unless LlmCostTracker.configuration.enabled
+        return @app.call(request_env) unless enabled?
 
         request_url  = request_env.url.to_s
         request_body = read_body(request_env.body)
         parser       = Parsers.find_for(request_url)
-        streaming    = parser&.streaming_request?(request_url, request_body)
-        request_body = inject_stream_usage_flag(request_env, parser, request_url) if streaming
+        request_parsed = parser ? safe_json_parse(request_body) : nil
+        streaming = parser&.streaming_request?(request_url, request_parsed)
+        if streaming
+          request_body = inject_stream_usage_flag(request_env, parser, request_url, request_parsed) || request_body
+        end
         stream_buffer = install_stream_tap(request_env) if streaming
 
-        Tracker.enforce_budget! if parser
+        if parser
+          Tracker.enforce_budget!(
+            provider: parser.provider_for(request_url),
+            model: parser.model_for(request_url, request_parsed),
+            request: request_parsed
+          )
+        end
         context_tags, metadata = tag_snapshot(request_env) if parser
         started_at = LlmCostTracker::Timing.now_monotonic
 
@@ -40,6 +49,26 @@ module LlmCostTracker
 
       private
 
+      def enabled?
+        return @enabled if defined?(@enabled)
+
+        @enabled = LlmCostTracker.configuration.enabled
+      end
+
+      def safe_json_parse(body)
+        return {} if body.nil? || body.empty?
+
+        JSON.parse(body)
+      rescue JSON::ParserError
+        {}
+      end
+
+      def auto_enable_stream_usage?
+        return @auto_enable_stream_usage if defined?(@auto_enable_stream_usage)
+
+        @auto_enable_stream_usage = LlmCostTracker.configuration.auto_enable_stream_usage
+      end
+
       def invoke_app_with_capture(request_env:, parser:, request_url:, request_body:, streaming:,
                                   stream_buffer:, context_tags:, metadata:, started_at:)
         response_received = false
@@ -63,16 +92,15 @@ module LlmCostTracker
         raise
       end
 
-      def inject_stream_usage_flag(request_env, parser, request_url)
-        body_string = read_body(request_env.body)
-        return body_string unless LlmCostTracker.configuration.auto_enable_stream_usage
-        return body_string unless parser&.auto_enable_stream_usage?(request_url)
+      def inject_stream_usage_flag(request_env, parser, request_url, request_parsed)
+        return nil unless auto_enable_stream_usage?
+        return nil unless parser&.auto_enable_stream_usage?(request_url)
 
-        body = JSON.parse(body_string)
-        return body_string if body["stream_options"].is_a?(Hash) && body["stream_options"].key?("include_usage")
+        stream_options = request_parsed["stream_options"]
+        return nil if stream_options.is_a?(Hash) && stream_options.key?("include_usage")
 
-        body["stream_options"] = (body["stream_options"] || {}).merge("include_usage" => true)
-        new_body = body.to_json
+        request_parsed["stream_options"] = (stream_options || {}).merge("include_usage" => true)
+        new_body = request_parsed.to_json
         request_env.body = new_body
         new_body
       end
@@ -80,9 +108,9 @@ module LlmCostTracker
       def process_interrupted_stream(parser:, request_url:, request_body:, latency_ms:,
                                      context_tags:, metadata:, error:)
         request = parser.safe_json_parse(request_body)
-        capture = UsageCapture.build(
+        event = Event.build(
           provider: parser.provider_for(request_url),
-          model: request["model"] || UsageCapture::UNKNOWN_MODEL,
+          model: request["model"] || Event::UNKNOWN_MODEL,
           token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
           usage_source: :unknown
@@ -92,7 +120,7 @@ module LlmCostTracker
           stream_interrupted_error: "#{error.class}: #{error.message}"
         )
         Tracker.record(
-          capture: capture,
+          event: event,
           latency_ms: latency_ms,
           metadata: merged_metadata,
           context_tags: context_tags
@@ -125,7 +153,7 @@ module LlmCostTracker
         return unless parsed
 
         Tracker.record(
-          capture: parsed,
+          event: parsed,
           latency_ms: latency_ms,
           metadata: metadata,
           context_tags: context_tags

data/lib/llm_cost_tracker/parsers/anthropic.rb

@@ -1,18 +1,21 @@
 # frozen_string_literal: true
 
 require_relative "base"
+require_relative "../providers/anthropic/tier_classification"
 
 module LlmCostTracker
   module Parsers
     class Anthropic < Base
       HOSTS = %w[api.anthropic.com].freeze
 
-      def match?(url)
-        match_uri?(url, hosts: HOSTS, path_includes: "/v1/messages")
-      end
+      class << self
+        def match?(url)
+          match_uri?(url, hosts: HOSTS, path_includes: "/v1/messages")
+        end
 
-      def provider_names
-        %w[anthropic]
+        def provider_names
+          %w[anthropic]
+        end
       end
 
       def parse(request_body:, response_status:, response_body:, **)
@@ -25,10 +28,10 @@ module LlmCostTracker
         request = safe_json_parse(request_body)
         cache_read = usage["cache_read_input_tokens"].to_i
 
-        UsageCapture.build(
+        Event.build(
           provider: "anthropic",
           provider_response_id: response["id"],
-          pricing_mode: pricing_mode(request: request, response: response, usage: usage),
+          pricing_mode: pricing_mode(request: request, usage: usage),
           model: response["model"] || request["model"],
           token_usage: token_usage(usage: usage, cache_read: cache_read),
           usage_source: :response,
@@ -49,14 +52,14 @@ module LlmCostTracker
             model: model,
             usage: usage,
             response_id: response_id,
-            pricing_mode: pricing_mode(request: request, response: nil, usage: usage)
+            pricing_mode: pricing_mode(request: request, usage: usage)
           )
         else
           build_unknown_stream_usage(
             provider: "anthropic",
             model: model,
             provider_response_id: response_id,
-            pricing_mode: pricing_mode(request: request, response: nil, usage: usage)
+            pricing_mode: pricing_mode(request: request, usage: usage)
           )
         end
       end
@@ -65,10 +68,6 @@ module LlmCostTracker
         "anthropic"
       end
 
-      DATA_RESIDENCY_GEOS = %w[us].freeze
-      STANDARD_EQUIVALENT_SERVICE_TIERS = %w[standard standard_only priority].freeze
-      private_constant :DATA_RESIDENCY_GEOS, :STANDARD_EQUIVALENT_SERVICE_TIERS
-
       private
 
       def stream_usage(events)
@@ -89,7 +88,7 @@ module LlmCostTracker
       def build_stream_result(model:, usage:, response_id:, pricing_mode:)
         cache_read = usage["cache_read_input_tokens"].to_i
 
-        UsageCapture.build(
+        Event.build(
           provider: "anthropic",
           provider_response_id: response_id,
           pricing_mode: pricing_mode,
@@ -171,29 +170,23 @@ module LlmCostTracker
         Logging.warn("Anthropic usage.cache_creation has unexpected shape: #{cache_creation.class}")
       end
 
-      def pricing_mode(request:, response:, usage:)
+      def pricing_mode(request:, usage:)
         modes = []
-        speed = usage&.fetch("speed", nil) || response&.fetch("speed", nil) || request["speed"]
-        service_tier = usage&.fetch("service_tier", nil) ||
-                       response&.fetch("service_tier", nil) ||
-                       request["service_tier"]
-        service_tier = nil if STANDARD_EQUIVALENT_SERVICE_TIERS.include?(service_tier.to_s)
+        speed = usage&.fetch("speed", nil) || request["speed"]
+        service_tier = usage&.fetch("service_tier", nil) || request["service_tier"]
+        service_tier = nil if Providers::Anthropic::TierClassification.standard_equivalent_tier?(service_tier)
 
         modes << Pricing.normalize_mode(speed)
         modes << Pricing.normalize_mode(service_tier)
-        geo = inference_geo(request: request, response: response, usage: usage).downcase
-        modes << "data_residency" if DATA_RESIDENCY_GEOS.include?(geo)
+        geo = inference_geo(request: request, usage: usage).downcase
+        modes << "data_residency" if Providers::Anthropic::TierClassification.data_residency_geo?(geo)
 
         modes = modes.compact.uniq
         modes.empty? ? nil : modes.join("_")
       end
 
-      def inference_geo(request:, response:, usage:)
-        (
-          usage&.fetch("inference_geo", nil) ||
-          response&.fetch("inference_geo", nil) ||
-          request["inference_geo"]
-        ).to_s
+      def inference_geo(request:, usage:)
+        (usage&.fetch("inference_geo", nil) || request["inference_geo"]).to_s
       end
     end
   end

data/lib/llm_cost_tracker/parsers/azure.rb

@@ -0,0 +1,46 @@
+# frozen_string_literal: true
+
+require_relative "base"
+require_relative "openai_usage"
+require_relative "../providers/azure/hosts"
+
+module LlmCostTracker
+  module Parsers
+    class Azure < Base
+      include OpenaiUsage
+
+      TRACKED_ENDPOINTS = %w[
+        chat/completions completions embeddings moderations responses
+        audio/transcriptions audio/translations audio/speech
+        images/generations images/edits images/variations
+      ].freeze
+
+      PATH_PATTERN = %r{\A/openai/(?:deployments/[^/]+|v1)/(?:#{TRACKED_ENDPOINTS.join('|')})\z}
+
+      class << self
+        def match?(url)
+          uri_matches?(url) do |uri|
+            LlmCostTracker::Providers::Azure::Hosts.openai?(uri.host) && uri.path.to_s.match?(PATH_PATTERN)
+          end
+        end
+
+        def provider_names
+          %w[azure_openai]
+        end
+      end
+
+      def provider_for(_request_url)
+        "azure_openai"
+      end
+
+      def model_for(request_url, request_parsed)
+        body_model = super
+        return body_model if body_model
+
+        uri = parsed_uri(request_url)
+        match = uri&.path&.match(%r{/openai/deployments/([^/]+)/})
+        match && match[1]
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/parsers/base.rb

@@ -6,51 +6,7 @@ require "uri"
 
 module LlmCostTracker
   module Parsers
-    class Base
-      def parse(**)
-        raise NotImplementedError
-      end
-
-      def provider_names
-        []
-      end
-
-      def match?(url)
-        raise NotImplementedError
-      end
-
-      def streaming_request?(_request_url, request_body)
-        return false if request_body.blank?
-
-        body = request_body.to_s
-
-        request = safe_json_parse(body)
-        request.is_a?(Hash) && request["stream"] == true
-      end
-
-      def parse_stream(**)
-        nil
-      end
-
-      def auto_enable_stream_usage?(_request_url)
-        false
-      end
-
-      def safe_json_parse(body)
-        return {} if body.blank?
-
-        JSON.parse(body)
-      rescue JSON::ParserError
-        {}
-      end
-
-      private
-
-      def uri_matches?(url)
-        uri = parsed_uri(url)
-        uri ? yield(uri) : false
-      end
-
+    module UrlMatchers
       def match_uri?(url, hosts: nil, exact_paths: nil, path_includes: nil, path_suffixes: nil, path_pattern: nil)
         uri_matches?(url) do |uri|
           host_match = hosts.nil? || hosts.include?(uri.host.to_s.downcase)
@@ -70,6 +26,11 @@ module LlmCostTracker
         end
       end
 
+      def uri_matches?(url)
+        uri = parsed_uri(url)
+        uri ? yield(uri) : false
+      end
+
       def parsed_uri(url)
         URI.parse(url.to_s)
       rescue URI::InvalidURIError
@@ -88,6 +49,51 @@ module LlmCostTracker
 
         matches
       end
+    end
+
+    class Base
+      extend UrlMatchers
+      include UrlMatchers
+
+      class << self
+        def match?(_url)
+          raise NotImplementedError
+        end
+
+        def provider_names
+          []
+        end
+      end
+
+      def parse(**)
+        raise NotImplementedError
+      end
+
+      def streaming_request?(_request_url, request_parsed)
+        request_parsed.is_a?(Hash) && request_parsed["stream"] == true
+      end
+
+      def model_for(_request_url, request_parsed)
+        request_parsed["model"] if request_parsed.is_a?(Hash)
+      end
+
+      def parse_stream(**)
+        nil
+      end
+
+      def auto_enable_stream_usage?(_request_url)
+        false
+      end
+
+      def safe_json_parse(body)
+        return {} if body.blank?
+
+        JSON.parse(body)
+      rescue JSON::ParserError
+        {}
+      end
+
+      private
 
       def each_event_data(events, reverse: false)
         enumerator = reverse ? events.reverse_each : events.each
@@ -109,11 +115,11 @@ module LlmCostTracker
 
       def build_unknown_stream_usage(provider:, model:, provider_response_id:, pricing_mode: nil,
                                      service_line_items: nil)
-        UsageCapture.build(
+        Event.build(
           provider: provider,
           provider_response_id: provider_response_id,
           pricing_mode: pricing_mode,
-          model: model || UsageCapture::UNKNOWN_MODEL,
+          model: model || Event::UNKNOWN_MODEL,
           token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
           usage_source: :unknown,

data/lib/llm_cost_tracker/parsers/gemini.rb

@@ -2,6 +2,7 @@
 
 require_relative "../billing/line_item"
 require_relative "base"
+require_relative "../providers/gemini/model_families"
 
 module LlmCostTracker
   module Parsers
@@ -9,17 +10,18 @@ module LlmCostTracker
       HOSTS = %w[generativelanguage.googleapis.com].freeze
       TRACKED_PATH_PATTERN = %r{/models/[^/:]+:(?:generateContent|streamGenerateContent)\z}
       STREAM_PATH_PATTERN  = /:streamGenerateContent\z/
-      PER_QUERY_GROUNDING_MODEL_PATTERN = /\bgemini-(?:[3-9]|[1-9]\d)\b/i
 
-      def match?(url)
-        match_uri?(url, hosts: HOSTS, path_pattern: TRACKED_PATH_PATTERN)
-      end
+      class << self
+        def match?(url)
+          match_uri?(url, hosts: HOSTS, path_pattern: TRACKED_PATH_PATTERN)
+        end
 
-      def provider_names
-        %w[gemini]
+        def provider_names
+          %w[gemini]
+        end
       end
 
-      def streaming_request?(request_url, request_body)
+      def streaming_request?(request_url, request_parsed)
         return true if match_uri?(request_url, path_pattern: STREAM_PATH_PATTERN)
 
         super
@@ -34,13 +36,13 @@ module LlmCostTracker
 
         request = safe_json_parse(request_body)
         model = extract_model_from_url(request_url)
-        build_usage_capture(
+        build_event(
           request_url: request_url,
           usage: usage,
           usage_source: :response,
           provider_response_id: response["responseId"],
           pricing_mode: pricing_mode(request: request, response_headers: response_headers),
-          service_line_items: grounding_line_items_for_response(response, model: model)
+          service_line_items: grounding_line_items(grounding_request_count(response["candidates"]), model: model)
         )
       end
 
@@ -55,7 +57,7 @@ module LlmCostTracker
         service_line_items = grounding_line_items_for_stream(events, model: model)
 
         if usage
-          build_usage_capture(
+          build_event(
             request_url: request_url,
             usage: usage,
             stream: true,
@@ -75,20 +77,24 @@ module LlmCostTracker
         end
       end
 
+      def model_for(request_url, _request_parsed)
+        extract_model_from_url(request_url)
+      end
+
       def provider_for(_request_url)
         "gemini"
       end
 
       private
 
-      def build_usage_capture(request_url:, usage:, usage_source:, stream: false, provider_response_id: nil,
-                              pricing_mode: nil, service_line_items: nil)
+      def build_event(request_url:, usage:, usage_source:, stream: false, provider_response_id: nil,
+                      pricing_mode: nil, service_line_items: nil)
         cache_read = usage["cachedContentTokenCount"].to_i
         tool_use_prompt = usage["toolUsePromptTokenCount"].to_i
         audio_input = audio_input_tokens(usage)
         audio_output = audio_output_tokens(usage)
 
-        UsageCapture.build(
+        Event.build(
           provider: "gemini",
           model: extract_model_from_url(request_url),
           pricing_mode: pricing_mode,
@@ -184,10 +190,6 @@ module LlmCostTracker
         headers.to_h.find { |key, _value| key.to_s.downcase == name }&.last
       end
 
-      def grounding_line_items_for_response(response, model:)
-        grounding_line_items(grounding_request_count(response["candidates"]), model: model)
-      end
-
       def grounding_line_items_for_stream(events, model:)
         quantity = find_event_value(events, reverse: true) do |data|
           count = grounding_request_count(data["candidates"])
@@ -223,11 +225,7 @@ module LlmCostTracker
       end
 
       def grounding_billed_quantity(query_count, model:)
-        per_query_billing?(model) ? query_count : 1
-      end
-
-      def per_query_billing?(model)
-        model.to_s.match?(PER_QUERY_GROUNDING_MODEL_PATTERN)
+        LlmCostTracker::Providers::Gemini::ModelFamilies.per_query_grounding?(model) ? query_count : 1
       end
     end
   end

data/lib/llm_cost_tracker/parsers/openai.rb

@@ -2,25 +2,13 @@
 
 require_relative "base"
 require_relative "openai_usage"
+require_relative "../providers/openai/hosts"
 
 module LlmCostTracker
   module Parsers
     class Openai < Base
       include OpenaiUsage
 
-      HOSTS = %w[
-        api.openai.com
-        us.api.openai.com
-        eu.api.openai.com
-        au.api.openai.com
-        ca.api.openai.com
-        jp.api.openai.com
-        in.api.openai.com
-        sg.api.openai.com
-        kr.api.openai.com
-        gb.api.openai.com
-        ae.api.openai.com
-      ].freeze
       TRACKED_PATHS = %w[
         /v1/chat/completions
         /v1/completions
@@ -35,34 +23,14 @@ module LlmCostTracker
         /v1/moderations
       ].freeze
 
-      def match?(url)
-        match_uri?(url, hosts: HOSTS, exact_paths: TRACKED_PATHS)
-      end
-
-      def provider_names
-        %w[openai]
-      end
-
-      def parse(request_url:, request_body:, response_status:, response_body:, **)
-        parse_openai_usage(
-          request_url: request_url,
-          request_body: request_body,
-          response_status: response_status,
-          response_body: response_body
-        )
-      end
-
-      def parse_stream(response_status:, request_url: nil, request_body: nil, events: [], **)
-        parse_openai_stream_usage(
-          request_url: request_url,
-          request_body: request_body,
-          response_status: response_status,
-          events: events
-        )
-      end
+      class << self
+        def match?(url)
+          match_uri?(url, hosts: Providers::Openai::Hosts::API_HOSTS, exact_paths: TRACKED_PATHS)
+        end
 
-      def auto_enable_stream_usage?(request_url)
-        openai_chat_completions_url?(request_url)
+        def provider_names
+          %w[openai]
+        end
       end
 
       def provider_for(_request_url)

data/lib/llm_cost_tracker/parsers/openai_compatible.rb

@@ -10,58 +10,41 @@ module LlmCostTracker
 
       TRACKED_PATH_SUFFIXES = %w[/chat/completions /completions /embeddings /responses].freeze
 
-      def match?(url)
-        match_uri?(url, path_suffixes: TRACKED_PATH_SUFFIXES) { |uri| provider_for_uri(uri) }
-      end
+      class << self
+        def match?(url)
+          match_uri?(url, path_suffixes: TRACKED_PATH_SUFFIXES) { |uri| provider_for_uri(uri) }
+        end
 
-      def provider_names
-        providers = LlmCostTracker.configuration.openai_compatible_providers
-        cached = @provider_names
-        return cached if cached && @provider_names_providers.equal?(providers)
+        def provider_names
+          providers = LlmCostTracker.configuration.openai_compatible_providers
+          cached = @provider_names
+          return cached if cached && @provider_names_providers.equal?(providers)
 
-        names = [
-          "openai_compatible",
-          *providers.each_value.map { |provider| provider.to_s.downcase }
-        ].uniq.freeze
-        return names unless providers.frozen?
+          names = [
+            "openai_compatible",
+            *providers.each_value.map { |provider| provider.to_s.downcase }
+          ].uniq.freeze
+          return names unless providers.frozen?
 
-        @provider_names_providers = providers
-        @provider_names = names
-      end
+          @provider_names_providers = providers
+          @provider_names = names
+        end
 
-      def parse(request_url:, request_body:, response_status:, response_body:, **)
-        parse_openai_usage(
-          request_url: request_url,
-          request_body: request_body,
-          response_status: response_status,
-          response_body: response_body
-        )
-      end
+        def provider_for(request_url)
+          provider_for_uri(parsed_uri(request_url)) || "openai_compatible"
+        end
 
-      def parse_stream(response_status:, request_url: nil, request_body: nil, events: [], **)
-        parse_openai_stream_usage(
-          request_url: request_url,
-          request_body: request_body,
-          response_status: response_status,
-          events: events
-        )
-      end
+        private
 
-      def auto_enable_stream_usage?(request_url)
-        openai_chat_completions_url?(request_url)
-      end
+        def provider_for_uri(uri)
+          return nil unless uri
 
-      def provider_for(request_url)
-        uri = parsed_uri(request_url)
-        provider_for_uri(uri) || "openai_compatible"
+          LlmCostTracker.configuration.openai_compatible_providers[uri.host.to_s.downcase]&.to_s
+        end
       end
 
-      private
-
-      def provider_for_uri(uri)
-        return nil unless uri
-
-        LlmCostTracker.configuration.openai_compatible_providers[uri.host.to_s.downcase]&.to_s
+      def provider_for(request_url)
+        self.class.provider_for(request_url)
       end
     end
   end