llm_cost_tracker 0.10.0 → 0.12.0

data/lib/llm_cost_tracker/providers/azure/parser.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Providers
+    module Azure
+      class Parser < LlmCostTracker::Parsers::Base
+        include Openai::ResponseParser
+
+        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|
+              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
+end

data/lib/llm_cost_tracker/providers/gemini/model_families.rb

@@ -5,10 +5,7 @@ module LlmCostTracker
     module Gemini
       module ModelFamilies
         PER_QUERY_GROUNDING_MODEL_PATTERN = /\bgemini-(?:[3-9]|[1-9]\d)\b/i
-
-        module_function
-
-        def per_query_grounding?(model)
+        def self.per_query_grounding?(model)
           model.to_s.match?(PER_QUERY_GROUNDING_MODEL_PATTERN)
         end
       end

data/lib/llm_cost_tracker/providers/gemini/parser.rb

@@ -0,0 +1,177 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Providers
+    module Gemini
+      class Parser < LlmCostTracker::Parsers::Base
+        HOSTS = %w[generativelanguage.googleapis.com].freeze
+        TRACKED_PATH_PATTERN = %r{/models/[^/:]+:(?:generateContent|streamGenerateContent)\z}
+        STREAM_PATH_PATTERN  = /:streamGenerateContent\z/
+
+        class << self
+          def match?(url)
+            match_uri?(url, hosts: HOSTS, path_pattern: TRACKED_PATH_PATTERN)
+          end
+
+          def provider_names
+            %w[gemini]
+          end
+        end
+
+        def streaming_request?(request_url, request_parsed)
+          return true if match_uri?(request_url, path_pattern: STREAM_PATH_PATTERN)
+
+          super
+        end
+
+        def parse(request_url:, request_body:, response_status:, response_body:, response_headers: nil)
+          return nil unless response_status == 200
+
+          response = safe_json_parse(response_body)
+          usage    = response["usageMetadata"]
+          return nil unless usage
+
+          request = safe_json_parse(request_body)
+          model = extract_model_from_url(request_url)
+          build_event(
+            request_url: request_url,
+            usage: usage,
+            usage_source: Usage::Source::RESPONSE,
+            provider_response_id: response["responseId"],
+            pricing_mode: pricing_mode(request: request, usage: usage, response_headers: response_headers),
+            service_line_items: grounding_line_items(grounding_request_count(response["candidates"]), model: model)
+          )
+        end
+
+        def parse_stream(response_status:, request_url: nil, request_body: nil, events: [], response_headers: nil)
+          return nil unless response_status == 200
+
+          request = safe_json_parse(request_body)
+          usage = merged_stream_usage(events)
+          model = extract_model_from_url(request_url)
+          response_id = stream_response_id(events)
+          mode = pricing_mode(request: request, usage: usage, response_headers: response_headers)
+          service_line_items = grounding_line_items_for_stream(events, model: model)
+
+          if usage
+            build_event(
+              request_url: request_url,
+              usage: usage,
+              stream: true,
+              usage_source: Usage::Source::STREAM_FINAL,
+              provider_response_id: response_id,
+              pricing_mode: mode,
+              service_line_items: service_line_items
+            )
+          else
+            build_unknown_stream_usage(
+              provider: "gemini",
+              model: model,
+              provider_response_id: response_id,
+              pricing_mode: mode,
+              service_line_items: service_line_items
+            )
+          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_event(request_url:,
+                        usage:,
+                        usage_source:,
+                        stream: false,
+                        provider_response_id: nil,
+                        pricing_mode: nil,
+                        service_line_items: nil)
+          Event.build(
+            provider: "gemini",
+            model: extract_model_from_url(request_url),
+            pricing_mode: pricing_mode,
+            token_usage: UsageExtractor.token_usage(usage),
+            stream: stream,
+            usage_source: usage_source,
+            provider_response_id: provider_response_id,
+            service_line_items: service_line_items
+          )
+        end
+
+        def merged_stream_usage(events)
+          find_event_value(events, reverse: true) do |data|
+            meta = data["usageMetadata"]
+            meta if meta.is_a?(Hash)
+          end
+        end
+
+        def stream_response_id(events)
+          find_event_value(events) { |data| data["responseId"] }
+        end
+
+        def extract_model_from_url(url)
+          uri = parsed_uri(url)
+          return nil unless uri
+
+          match = uri.path.match(%r{/models/([^/:]+)})
+          match && match[1]
+        end
+
+        def pricing_mode(request:, usage:, response_headers:)
+          body_mode = Pricing::Mode.normalize(usage && usage["serviceTier"])
+          return body_mode if body_mode
+
+          header_mode = Pricing::Mode.normalize(response_header(response_headers, "x-gemini-service-tier"))
+          return header_mode if header_mode
+
+          request_mode = Pricing::Mode.normalize(request["service_tier"] || request["serviceTier"])
+          request_mode == "flex" ? request_mode : nil
+        end
+
+        def response_header(headers, name)
+          headers.to_h.find { |key, _value| key.to_s.downcase == name }&.last
+        end
+
+        def grounding_line_items_for_stream(events, model:)
+          quantity = find_event_value(events, reverse: true) do |data|
+            count = grounding_request_count(data["candidates"])
+            count if count.positive?
+          end
+          grounding_line_items(quantity || 0, model: model)
+        end
+
+        def grounding_request_count(candidates)
+          Array(candidates).sum do |candidate|
+            queries = candidate.dig("groundingMetadata", "webSearchQueries") || []
+            Array(queries).size
+          end
+        end
+
+        def grounding_line_items(query_count, model:)
+          return [] unless query_count.positive?
+
+          billed_quantity = grounding_billed_quantity(query_count, model: model)
+          [
+            Charges::LineItem.build(
+              dimension_key: "grounding_request",
+              quantity: billed_quantity,
+              cost_status: Charges::CostStatus::UNKNOWN,
+              pricing_basis: "provider_usage",
+              provider_field: "response.candidates.groundingMetadata.webSearchQueries",
+              details: { web_search_queries: query_count }
+            )
+          ]
+        end
+
+        def grounding_billed_quantity(query_count, model:)
+          ModelFamilies.per_query_grounding?(model) ? query_count : 1
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/providers/gemini/usage_extractor.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Providers
+    module Gemini
+      module UsageExtractor
+        def self.token_usage(usage)
+          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)
+          image_input = image_input_tokens(usage)
+          image_output = image_output_tokens(usage)
+
+          Usage::TokenUsage.build(
+            input_tokens: regular_input_tokens(usage: usage,
+                                               cache_read: cache_read,
+                                               audio_input: audio_input,
+                                               image_input: image_input) +
+                          tool_use_prompt,
+            output_tokens: regular_output_tokens(usage: usage,
+                                                 audio_output: audio_output,
+                                                 image_output: image_output),
+            total_tokens: usage["totalTokenCount"],
+            cache_read_input_tokens: cache_read,
+            audio_input_tokens: audio_input,
+            audio_output_tokens: audio_output,
+            image_input_tokens: image_input,
+            image_output_tokens: image_output,
+            hidden_output_tokens: usage["thoughtsTokenCount"]
+          )
+        end
+
+        def self.gross_output_tokens(usage)
+          usage["candidatesTokenCount"].to_i + usage["thoughtsTokenCount"].to_i
+        end
+
+        def self.regular_input_tokens(usage:, cache_read:, audio_input:, image_input:)
+          [usage["promptTokenCount"].to_i - cache_read - audio_input - image_input, 0].max
+        end
+
+        def self.regular_output_tokens(usage:, audio_output:, image_output:)
+          [gross_output_tokens(usage) - audio_output - image_output, 0].max
+        end
+
+        def self.audio_input_tokens(usage)
+          prompt_audio = modality_tokens(usage["promptTokensDetails"], "AUDIO")
+          cache_audio = modality_tokens(usage["cacheTokensDetails"], "AUDIO")
+          [prompt_audio - cache_audio, 0].max
+        end
+
+        def self.audio_output_tokens(usage)
+          modality_tokens(usage["candidatesTokensDetails"], "AUDIO")
+        end
+
+        def self.image_input_tokens(usage)
+          prompt_image = modality_tokens(usage["promptTokensDetails"], "IMAGE")
+          cache_image = modality_tokens(usage["cacheTokensDetails"], "IMAGE")
+          [prompt_image - cache_image, 0].max
+        end
+
+        def self.image_output_tokens(usage)
+          modality_tokens(usage["candidatesTokensDetails"], "IMAGE")
+        end
+
+        def self.modality_tokens(details, modality)
+          Array(details).sum do |detail|
+            next 0 unless detail["modality"] == modality
+
+            detail["tokenCount"].to_i
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/providers/openai/hosts.rb

@@ -20,13 +20,7 @@ module LlmCostTracker
 
         DATA_RESIDENCY_HOST_PATTERN = /\A[a-z]{2,3}\.api\.openai\.com\z/
 
-        module_function
-
-        def api?(host)
-          API_HOSTS.include?(host.to_s.downcase)
-        end
-
-        def data_residency?(host)
+        def self.data_residency?(host)
           host.to_s.downcase.match?(DATA_RESIDENCY_HOST_PATTERN)
         end
       end

data/lib/llm_cost_tracker/providers/openai/model_families.rb

@@ -19,26 +19,23 @@ module LlmCostTracker
         NON_REASONING_GPT5_PATTERN = /\Agpt-5(?:\.\d+)?-chat\b/i
 
         CHAT_COMPLETIONS_SEARCH_MODEL_PATTERN = /-search-(?:preview|api)\b/i
-
-        module_function
-
-        def data_residency?(model)
+        def self.data_residency?(model)
           model.to_s.match?(DATA_RESIDENCY_MODEL_PATTERN)
         end
 
-        def image_output?(model)
+        def self.image_output?(model)
           model.to_s.match?(IMAGE_OUTPUT_MODEL_PATTERN)
         end
 
-        def character_billed_tts?(model)
+        def self.character_billed_tts?(model)
           model.to_s.match?(CHARACTER_BILLED_TTS_MODEL_PATTERN)
         end
 
-        def chat_completions_search?(model)
+        def self.chat_completions_search?(model)
           model.to_s.match?(CHAT_COMPLETIONS_SEARCH_MODEL_PATTERN)
         end
 
-        def reasoning?(model)
+        def self.reasoning?(model)
           name = model.to_s
           return false if name.empty?
           return false if NON_REASONING_GPT5_PATTERN.match?(name)

data/lib/llm_cost_tracker/providers/openai/parser.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Providers
+    module Openai
+      class Parser < LlmCostTracker::Parsers::Base
+        include ResponseParser
+
+        TRACKED_PATHS = %w[
+          /v1/chat/completions
+          /v1/completions
+          /v1/embeddings
+          /v1/responses
+          /v1/images/generations
+          /v1/images/edits
+          /v1/images/variations
+          /v1/audio/transcriptions
+          /v1/audio/translations
+          /v1/audio/speech
+          /v1/moderations
+        ].freeze
+
+        class << self
+          def match?(url)
+            match_uri?(url, hosts: Hosts::API_HOSTS, exact_paths: TRACKED_PATHS)
+          end
+
+          def provider_names
+            %w[openai]
+          end
+        end
+
+        def provider_for(_request_url)
+          "openai"
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/providers/openai/response_parser.rb

@@ -0,0 +1,152 @@
+# frozen_string_literal: true
+
+require "active_support/core_ext/hash/keys"
+
+require_relative "hosts"
+require_relative "model_families"
+require_relative "service_charges"
+require_relative "usage_extractor"
+
+module LlmCostTracker
+  module Providers
+    module Openai
+      module ResponseParser
+        include LlmCostTracker::Providers::Openai::ServiceCharges
+
+        class << self
+          def combined_pricing_mode(host:, model:, service_tier:)
+            modes = [Pricing::Mode.normalize(service_tier)]
+            modes << "data_residency" if Hosts.data_residency?(host) && ModelFamilies.data_residency?(model)
+            Pricing::Mode.compose(modes)
+          end
+
+          def event_from_response(response:, request:, provider:, host:, usage_source:, pricing_mode: nil)
+            usage = response["usage"]&.deep_symbolize_keys
+            return nil unless usage
+
+            model = response["model"] || request["model"]
+            service_line_items =
+              ServiceCharges.service_line_items_for(response, request: request, model: response["model"]) +
+              ServiceCharges.transcription_line_items(usage)
+            Event.build(
+              provider: provider,
+              provider_response_id: response["id"],
+              pricing_mode: pricing_mode || combined_pricing_mode(
+                host: host, model: model, service_tier: response["service_tier"] || request["service_tier"]
+              ),
+              model: model,
+              token_usage: UsageExtractor.token_usage(usage, model: model),
+              usage_source: usage_source,
+              service_line_items: service_line_items
+            )
+          end
+        end
+
+        def parse(request_url:, request_body:, response_status:, response_body:, **)
+          return nil unless response_status == 200
+
+          ResponseParser.event_from_response(
+            response: safe_json_parse(response_body),
+            request: safe_json_parse(request_body),
+            provider: provider_for(request_url),
+            host: parsed_uri(request_url)&.host,
+            usage_source: Usage::Source::RESPONSE
+          )
+        end
+
+        def parse_stream(response_status:, request_url: nil, request_body: nil, events: [], **)
+          return nil unless response_status == 200
+
+          request = safe_json_parse(request_body)
+          usage = detect_stream_usage(events)
+          context = stream_capture_context(events: events, request: request, request_url: request_url)
+
+          return build_known_stream_usage(usage: usage, **context) if usage
+
+          warn_missing_stream_usage(request_url: request_url, request: request)
+          build_unknown_stream_usage(**context)
+        end
+
+        def auto_enable_stream_usage?(request_url)
+          openai_chat_completions_url?(request_url)
+        end
+
+        private
+
+        def stream_capture_context(events:, request:, request_url:)
+          model = find_event_value(events) do |data|
+            data["model"] || data.dig("response", "model") || data.dig("chunk", "model")
+          end || request["model"]
+          {
+            provider: provider_for(request_url),
+            model: model,
+            provider_response_id: find_event_value(events) do |data|
+              data["id"] || data.dig("response", "id") || data.dig("chunk", "id")
+            end,
+            pricing_mode: pricing_mode(
+              request_url: request_url,
+              model: model,
+              service_tier: stream_pricing_mode(events) || request["service_tier"]
+            ),
+            service_line_items: openai_stream_service_line_items(events, request: request, model: model)
+          }
+        end
+
+        def build_known_stream_usage(usage:,
+                                     provider:,
+                                     model:,
+                                     provider_response_id:,
+                                     pricing_mode:,
+                                     service_line_items:)
+          Event.build(
+            provider: provider,
+            provider_response_id: provider_response_id,
+            pricing_mode: pricing_mode,
+            model: model,
+            token_usage: UsageExtractor.token_usage(usage, model: model),
+            stream: true,
+            usage_source: Usage::Source::STREAM_FINAL,
+            service_line_items: service_line_items
+          )
+        end
+
+        def warn_missing_stream_usage(request_url:, request:)
+          return unless request["stream"]
+          return unless openai_chat_completions_url?(request_url)
+          return if request.dig("stream_options", "include_usage")
+
+          Logging.warn(
+            "OpenAI-compatible chat-completions stream finished without a final usage chunk. " \
+            "Set `stream_options: { include_usage: true }` in your request body so the gem can " \
+            "record token counts. This call was stored with usage_source=#{Usage::Source::UNKNOWN}."
+          )
+        end
+
+        def openai_chat_completions_url?(request_url)
+          uri = parsed_uri(request_url)
+          uri && uri.path.to_s.end_with?("/chat/completions")
+        end
+
+        def detect_stream_usage(events)
+          usage = find_event_value(events, reverse: true) do |data|
+            candidate = data["usage"] || data.dig("response", "usage") || data.dig("chunk", "usage")
+            candidate if candidate.is_a?(Hash)
+          end
+          usage&.deep_symbolize_keys
+        end
+
+        def stream_pricing_mode(events)
+          find_event_value(events, reverse: true) do |data|
+            data["service_tier"] || data.dig("response", "service_tier") || data.dig("chunk", "service_tier")
+          end
+        end
+
+        def pricing_mode(request_url:, model:, service_tier:)
+          ResponseParser.combined_pricing_mode(host: parsed_uri(request_url)&.host,
+                                               model: model,
+                                               service_tier: service_tier)
+        end
+      end
+    end
+  end
+end