llm_cost_tracker 0.14.1 → 0.15.0

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

@@ -9,67 +9,157 @@ module LlmCostTracker
           "web_fetch_request" => :web_fetch_requests
         }.freeze
         DATA_RESIDENCY_GEOS = %w[us].freeze
-        private_constant :SERVER_TOOL_LINE_ITEMS, :DATA_RESIDENCY_GEOS
-
-        def self.token_usage(usage)
-          input = usage[:input_tokens].to_i
-          output = usage[:output_tokens].to_i
-          cache_read = usage[:cache_read_input_tokens].to_i
-          cache_write, cache_write_extended = cache_writes(usage)
-
-          Usage::TokenUsage.build(
-            input_tokens: input,
-            output_tokens: output,
-            cache_read_input_tokens: cache_read,
-            cache_write_input_tokens: cache_write,
-            cache_write_extended_input_tokens: cache_write_extended,
-            hidden_output_tokens: usage.dig(:output_tokens_details, :thinking_tokens).to_i
-          )
-        end
+        BILLED_REFUSAL_CATEGORIES = %w[bio frontier_llm reasoning_extraction].freeze
+        BEDROCK_REGIONAL_PROFILE = /\A(?!global\.)[a-z]+\.anthropic\./
+        private_constant :SERVER_TOOL_LINE_ITEMS,
+                         :DATA_RESIDENCY_GEOS,
+                         :BILLED_REFUSAL_CATEGORIES,
+                         :BEDROCK_REGIONAL_PROFILE
 
-        def self.pricing_mode(request:, usage:)
-          speed = usage&.dig(:speed) || request&.dig(:speed)
-          service_tier = usage&.dig(:service_tier) || request&.dig(:service_tier)
-          geo = (usage&.dig(:inference_geo) || request&.dig(:inference_geo)).to_s.downcase
+        class << self
+          def token_usage(usage)
+            entries = [usage, *iterations(usage, "compaction")]
+            writes = entries.map { |entry| cache_writes(entry) }
 
-          modes = [Pricing::Mode.normalize(speed), Pricing::Mode.normalize(service_tier)]
-          modes << "data_residency" if DATA_RESIDENCY_GEOS.include?(geo)
-          Pricing::Mode.compose(modes)
-        end
+            Usage::TokenUsage.build(
+              input_tokens: entries.sum { |entry| entry[:input_tokens].to_i },
+              output_tokens: entries.sum { |entry| entry[:output_tokens].to_i },
+              cache_read_input_tokens: entries.sum { |entry| entry[:cache_read_input_tokens].to_i },
+              cache_write_input_tokens: writes.sum(&:first),
+              cache_write_extended_input_tokens: writes.sum(&:last),
+              hidden_output_tokens: usage.dig(:output_tokens_details, :thinking_tokens).to_i
+            )
+          end
+
+          def served_model(usage, content = nil)
+            served = iterations(usage, "fallback_message").last&.dig(:model)
+            served ||= fallback_blocks(content).last&.dig(:to, :model)
+            served&.to_s.presence
+          end
 
-        def self.service_line_items(usage)
-          server_tool_use = usage[:server_tool_use]
-          return [] unless server_tool_use.is_a?(Hash)
+          def iteration_line_items(usage, content: nil)
+            other_model_iterations(usage, content).map do |iteration|
+              tokens = token_usage(iteration)
+              Charges::LineItem.build(
+                dimension_key: "model_iteration",
+                quantity: 1,
+                pricing_basis: "provider_usage",
+                provider_field: "usage.iterations.#{iteration[:type]}",
+                details: tokens.to_h.select { |_key, count| count.positive? }.merge(model: iteration[:model].to_s)
+              )
+            end
+          end
 
-          SERVER_TOOL_LINE_ITEMS.filter_map do |dimension_key, count_key|
-            quantity = server_tool_use[count_key].to_i
-            next if quantity.zero?
+          def refusal_line_items(usage, stop_reason:, refusal_category:)
+            return [] unless stop_reason.to_s == "refusal" && token_usage(usage).output_tokens.zero?
+            return [] if BILLED_REFUSAL_CATEGORIES.include?(refusal_category.to_s)
 
-            Charges::LineItem.build(
-              dimension_key: dimension_key,
-              quantity: quantity,
-              cost_status: Charges::CostStatus::UNKNOWN,
+            [Charges::LineItem.build(
+              dimension_key: "billed_request",
+              quantity: 1,
+              rate_amount: 0,
+              cost: 0,
               pricing_basis: "provider_usage",
-              provider_field: "usage.server_tool_use.#{count_key}"
-            )
+              price_source: "provider_response",
+              provider_field: "stop_details.category",
+              details: { category: refusal_category&.to_s }.compact
+            )]
           end
-        end
 
-        def self.cache_writes(usage)
-          cache_creation = usage[:cache_creation]
-          if cache_creation.is_a?(Hash)
-            [cache_creation[:ephemeral_5m_input_tokens].to_i, cache_creation[:ephemeral_1h_input_tokens].to_i]
-          else
-            warn_unexpected_cache_creation(cache_creation, usage)
-            [usage[:cache_creation_input_tokens].to_i, 0]
+          def pricing_mode(request:, usage:, host: nil, model: request&.dig(:model))
+            modes = %i[speed service_tier].map { |key| Pricing::Mode.normalize(reported(key, usage, request)) }
+            modes << "data_residency" if data_residency?(usage, request) || regional_host?(host, model)
+            Pricing::Mode.compose(modes)
           end
-        end
 
-        def self.warn_unexpected_cache_creation(cache_creation, usage)
-          return if cache_creation.nil?
-          return if usage.key?(:cache_creation_input_tokens)
+          def regional_host?(host, model)
+            Openai::Hosts.vertex_non_global?(host) &&
+              Pricing::Matcher.modifier_priced?(provider: "anthropic", model: model, modifier: "data_residency")
+          end
+
+          def service_line_items(usage)
+            server_tool_use = usage[:server_tool_use]
+            return [] unless server_tool_use.is_a?(Hash)
+
+            SERVER_TOOL_LINE_ITEMS.filter_map do |dimension_key, count_key|
+              quantity = server_tool_use[count_key].to_i
+              next if quantity.zero?
+
+              Charges::LineItem.build(
+                dimension_key: dimension_key,
+                quantity: quantity,
+                cost_status: Charges::CostStatus::UNKNOWN,
+                pricing_basis: "provider_usage",
+                provider_field: "usage.server_tool_use.#{count_key}"
+              )
+            end
+          end
 
-          Logging.warn("Anthropic usage.cache_creation has unexpected shape: #{cache_creation.class}")
+          private
+
+          def other_model_iterations(usage, content)
+            iterations(usage, "advisor_message") + billed_fallback_iterations(usage, content)
+          end
+
+          def billed_fallback_iterations(usage, content)
+            served = served_model(usage, content)
+            return [] unless served
+
+            categories = fallback_blocks(content).map { |block| block.dig(:trigger, :category).to_s }
+            declined_hops(usage, served).zip(categories).flat_map do |hop, category|
+              billed_hop?(hop, category) ? hop : []
+            end
+          end
+
+          def declined_hops(usage, served)
+            declined = iterations(usage, "message").reject { |entry| ["", served].include?(entry[:model].to_s) }
+            declined.group_by { |entry| entry[:model].to_s }.values
+          end
+
+          def billed_hop?(hop, category)
+            hop.any? { |entry| entry[:output_tokens].to_i.positive? } || BILLED_REFUSAL_CATEGORIES.include?(category)
+          end
+
+          def fallback_blocks(content)
+            Array(content).select { |block| block.is_a?(Hash) && block[:type].to_s == "fallback" }
+          end
+
+          def iterations(usage, type)
+            Array(usage[:iterations]).select { |entry| entry.is_a?(Hash) && entry[:type].to_s == type }
+          end
+
+          def reported(key, usage, request)
+            usage&.dig(key) || request&.dig(key)
+          end
+
+          def data_residency?(usage, request)
+            DATA_RESIDENCY_GEOS.include?(reported(:inference_geo, usage, request).to_s.downcase) ||
+              bedrock_regional?(request&.dig(:model))
+          end
+
+          def bedrock_regional?(model)
+            model.to_s.match?(BEDROCK_REGIONAL_PROFILE) &&
+              Pricing::Matcher.modifier_priced?(provider: "bedrock", model: model.to_s, modifier: "data_residency")
+          end
+
+          def cache_writes(usage)
+            cache_creation = usage[:cache_creation]
+            if cache_creation.is_a?(Hash)
+              five_minute = cache_creation[:ephemeral_5m_input_tokens].to_i
+              one_hour = cache_creation[:ephemeral_1h_input_tokens].to_i
+              [five_minute + [usage[:cache_creation_input_tokens].to_i - five_minute - one_hour, 0].max, one_hour]
+            else
+              warn_unexpected_cache_creation(cache_creation, usage)
+              [usage[:cache_creation_input_tokens].to_i, 0]
+            end
+          end
+
+          def warn_unexpected_cache_creation(cache_creation, usage)
+            return if cache_creation.nil?
+            return if usage.key?(:cache_creation_input_tokens)
+
+            Logging.warn("Anthropic usage.cache_creation has unexpected shape: #{cache_creation.class}")
+          end
         end
       end
     end

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

@@ -4,7 +4,7 @@ module LlmCostTracker
   module Providers
     module Azure
       module Hosts
-        OPENAI_HOST_PATTERN = /\A[a-z0-9][a-z0-9-]*\.(?:openai\.azure\.com|services\.ai\.azure\.com)\z/i
+        OPENAI_HOST_PATTERN = /\A[a-z0-9][a-z0-9-]*\.(?:openai|services\.ai|cognitiveservices)\.azure\.com\z/i
         def self.openai?(host)
           host.to_s.match?(OPENAI_HOST_PATTERN)
         end

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

@@ -5,14 +5,12 @@ module LlmCostTracker
     module Azure
       class Parser < LlmCostTracker::Parsers::Base
         include Openai::ResponseParser
+        include Openai::StreamParser
 
-        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}
+        PATH_PATTERN = %r{
+          \A/openai/(?:deployments/[^/]+|v1)/(?:#{Openai::Parser::TRACKED_ENDPOINTS.join('|')}|responses/resp_[^/]+)\z
+        }x
+        DEPLOYMENT_PATH_PATTERN = %r{/openai/deployments/([^/]+)/}
 
         class << self
           def match?(url)
@@ -31,12 +29,7 @@ module LlmCostTracker
         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]
+          super || parsed_uri(request_url)&.path.to_s[DEPLOYMENT_PATH_PATTERN, 1]
         end
 
         def auto_enable_stream_usage?(request_url, request_parsed)

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

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Providers
+    module Gemini
+      class Grounding
+        CANDIDATE_FIELDS = {
+          "grounding_request" => "response.candidates.groundingMetadata.webSearchQueries",
+          "maps_grounding_request" => "response.candidates.groundingMetadata.groundingChunks.maps"
+        }.freeze
+        TOOL_KINDS = {
+          "google_search" => "grounding_request",
+          "google_maps" => "maps_grounding_request"
+        }.freeze
+        TOOL_COUNT_FIELD = "response.usage.grounding_tool_count"
+
+        class << self
+          def from_response(response)
+            usage = response["usage"]
+            return from_candidates(response["candidates"]) unless usage.is_a?(Hash)
+
+            from_tool_counts(usage["grounding_tool_count"])
+          end
+
+          def from_candidates(candidates)
+            counts = Array(candidates).each_with_object(Hash.new(0)) do |candidate, totals|
+              metadata = candidate["groundingMetadata"]
+              next unless metadata.is_a?(Hash)
+
+              kind, count = metadata_count(metadata)
+              totals[kind] += count
+            end
+            new(counts)
+          end
+
+          private
+
+          def from_tool_counts(entries)
+            counts = Array(entries).each_with_object(Hash.new(0)) do |entry, totals|
+              kind = TOOL_KINDS[entry["type"]]
+              totals[kind] += entry["count"].to_i if kind
+            end
+            new(counts, provider_field: TOOL_COUNT_FIELD)
+          end
+
+          def metadata_count(metadata)
+            queries = unique_count(metadata["webSearchQueries"])
+            if Array(metadata["groundingChunks"]).any? { |chunk| chunk.is_a?(Hash) && chunk.key?("maps") }
+              ["maps_grounding_request", [queries, 1].max]
+            else
+              ["grounding_request", queries + unique_count(metadata["imageSearchQueries"])]
+            end
+          end
+
+          def unique_count(queries)
+            Array(queries).map { |query| query.to_s.strip }.reject(&:empty?).uniq.size
+          end
+        end
+
+        def initialize(counts, provider_field: nil)
+          @counts = counts
+          @provider_field = provider_field
+        end
+
+        def any?
+          @counts.values.any?(&:positive?)
+        end
+
+        def line_items(model:)
+          per_query = ModelFamilies.per_query_grounding?(model)
+          @counts.filter_map do |kind, count|
+            next unless count.positive?
+
+            Charges::LineItem.build(
+              dimension_key: kind,
+              quantity: per_query ? count : 1,
+              cost_status: Charges::CostStatus::UNKNOWN,
+              pricing_basis: "provider_usage",
+              provider_field: @provider_field || CANDIDATE_FIELDS.fetch(kind),
+              details: { web_search_queries: count }
+            )
+          end
+        end
+      end
+    end
+  end
+end