llm_cost_tracker 0.7.3 → 0.9.0

data/lib/llm_cost_tracker/masking.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Masking
+    SENSITIVE_KEYS = %i[
+      provider_api_key_id provider_workspace_id provider_organization_id provider_project_id
+    ].to_set.freeze
+    MASK_TAIL_LENGTH = 4
+
+    module_function
+
+    def mask_value(key, value)
+      string = value.to_s
+      return string unless SENSITIVE_KEYS.include?(key.to_sym)
+      return string if string.length <= MASK_TAIL_LENGTH
+
+      "***#{string[-MASK_TAIL_LENGTH, MASK_TAIL_LENGTH]}"
+    end
+
+    def format_attribution(attribution, separator: ", ")
+      return "" if attribution.nil? || attribution.empty?
+
+      attribution.map { |key, value| "#{key}=#{mask_value(key, value)}" }.join(separator)
+    end
+
+    def mask_hash(hash)
+      return hash unless hash.is_a?(Hash)
+
+      hash.each_with_object({}) do |(key, value), masked|
+        masked[key] = case value
+                      when Hash then mask_hash(value)
+                      when Array then value.map { |entry| entry.is_a?(Hash) ? mask_hash(entry) : entry }
+                      else
+                        mask_value(key, value)
+                      end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -2,10 +2,12 @@
 
 require "faraday"
 require "json"
+require "stringio"
 require "uri"
 
 require_relative "../logging"
 require_relative "../capture/stream"
+require_relative "../timing"
 
 module LlmCostTracker
   module Middleware
@@ -19,31 +21,85 @@ module LlmCostTracker
         return @app.call(request_env) unless LlmCostTracker.configuration.enabled
 
         request_url  = request_env.url.to_s
-        request_body = read_body(request_env.body) || ""
+        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
         stream_buffer = install_stream_tap(request_env) if streaming
 
         Tracker.enforce_budget! if parser
         context_tags, metadata = tag_snapshot(request_env) if parser
-        started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        started_at = LlmCostTracker::Timing.now_monotonic
 
+        invoke_app_with_capture(
+          request_env: request_env, parser: parser, request_url: request_url,
+          request_body: request_body, streaming: streaming, stream_buffer: stream_buffer,
+          context_tags: context_tags, metadata: metadata, started_at: started_at
+        )
+      end
+
+      private
+
+      def invoke_app_with_capture(request_env:, parser:, request_url:, request_body:, streaming:,
+                                  stream_buffer:, context_tags:, metadata:, started_at:)
+        response_received = false
         @app.call(request_env).on_complete do |response_env|
+          response_received = true
           process(
-            parser: parser,
-            request_url: request_url,
-            request_body: request_body,
-            response_env: response_env,
-            latency_ms: ((Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at) * 1000).round,
-            streaming: streaming,
-            stream_buffer: stream_buffer,
-            context_tags: context_tags,
-            metadata: metadata
+            parser: parser, request_url: request_url, request_body: request_body,
+            response_env: response_env, latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at),
+            streaming: streaming, stream_buffer: stream_buffer,
+            context_tags: context_tags, metadata: metadata
+          )
+        end
+      rescue StandardError => e
+        if streaming && parser && !response_received
+          process_interrupted_stream(
+            parser: parser, request_url: request_url, request_body: request_body,
+            latency_ms: LlmCostTracker::Timing.elapsed_ms(started_at),
+            context_tags: context_tags, metadata: metadata, error: e
           )
         end
+        raise
       end
 
-      private
+      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)
+
+        body = JSON.parse(body_string)
+        return body_string if body["stream_options"].is_a?(Hash) && body["stream_options"].key?("include_usage")
+
+        body["stream_options"] = (body["stream_options"] || {}).merge("include_usage" => true)
+        new_body = body.to_json
+        request_env.body = new_body
+        new_body
+      end
+
+      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(
+          provider: parser.provider_for(request_url),
+          model: request["model"] || UsageCapture::UNKNOWN_MODEL,
+          token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
+          stream: true,
+          usage_source: :unknown
+        )
+        merged_metadata = (metadata || {}).merge(
+          stream_interrupted: true,
+          stream_interrupted_error: "#{error.class}: #{error.message}"
+        )
+        Tracker.record(
+          capture: capture,
+          latency_ms: latency_ms,
+          metadata: merged_metadata,
+          context_tags: context_tags
+        )
+      rescue StandardError => e
+        Logging.warn("Error recording interrupted stream: #{e.class}: #{e.message}")
+      end
 
       def process(parser:, request_url:, request_body:, response_env:,
                   latency_ms:, streaming:, stream_buffer:, context_tags:, metadata:)
@@ -101,21 +157,14 @@ module LlmCostTracker
       end
 
       def parse_stream(parser:, request_url:, request_body:, response_env:, stream_buffer:)
-        if stream_buffer&.dig(:overflowed)
-          Logging.warn(capture_warning(request_url, stream_buffer))
-          return parser.parse_stream(
-            request_url: request_url,
-            request_body: request_body,
-            response_status: response_env.status,
-            response_headers: response_env.response_headers
-          )
-        end
+        overflowed = stream_buffer&.dig(:overflowed) == true
+        Logging.warn(capture_warning(request_url, stream_buffer)) if overflowed
 
         body = stream_buffer&.dig(:buffer)&.string
         body = read_body(response_env.body) if body.blank?
 
         if body.blank?
-          Logging.warn(capture_warning(request_url, stream_buffer))
+          Logging.warn(capture_warning(request_url, stream_buffer)) unless overflowed
           return parser.parse_stream(
             request_url: request_url,
             request_body: request_body,
@@ -124,7 +173,7 @@ module LlmCostTracker
           )
         end
 
-        events = Parsers::SSE.parse(body)
+        events = overflowed ? [] : Parsers::SSE.parse(body)
         parser.parse_stream(
           request_url: request_url,
           request_body: request_body,
@@ -134,8 +183,19 @@ module LlmCostTracker
         )
       end
 
+      def forward_on_data_chunk(callable, chunk, size, env)
+        arity = callable.arity
+        return callable.call(chunk, size, env) if arity.negative?
+
+        case arity
+        when 0, 1 then callable.call(chunk)
+        when 2 then callable.call(chunk, size)
+        else callable.call(chunk, size, env)
+        end
+      end
+
       def install_stream_tap(request_env)
-        request = request_env.try(:request)
+        request = request_env.request
         return nil unless request
 
         original = request.on_data
@@ -144,16 +204,16 @@ module LlmCostTracker
         state = { buffer: StringIO.new, bytes: 0, overflowed: false }
         request.on_data = proc do |chunk, size, env|
           chunk = chunk.to_s
-          unless state[:overflowed]
-            if state[:bytes] + chunk.bytesize <= Capture::Stream::LIMIT_BYTES
-              state[:buffer] << chunk
-              state[:bytes] += chunk.bytesize
-            else
-              state[:overflowed] = true
-              state[:buffer] = nil
-            end
+          remaining = Capture::Stream::LIMIT_BYTES - state[:bytes]
+          if chunk.bytesize <= remaining
+            state[:buffer] << chunk
+            state[:bytes] += chunk.bytesize
+          else
+            state[:buffer] << chunk.byteslice(0, remaining) if remaining.positive?
+            state[:bytes] += [remaining, 0].max
+            state[:overflowed] = true
           end
-          original.call(chunk, size, env)
+          forward_on_data_chunk(original, chunk, size, env)
         end
         state
       rescue StandardError => e
@@ -204,8 +264,8 @@ module LlmCostTracker
         uri = URI.parse(value.to_s)
         uri.query = nil
         uri.fragment = nil
-        uri.try(:user=, nil)
-        uri.try(:password=, nil)
+        uri.user = nil
+        uri.password = nil
         uri.to_s
       rescue URI::InvalidURIError
         value.to_s.split("?", 2).first

data/lib/llm_cost_tracker/parsers/anthropic.rb

@@ -31,7 +31,8 @@ module LlmCostTracker
           pricing_mode: pricing_mode(request: request, response: response, usage: usage),
           model: response["model"] || request["model"],
           token_usage: token_usage(usage: usage, cache_read: cache_read),
-          usage_source: :response
+          usage_source: :response,
+          service_line_items: service_line_items(usage)
         )
       end
 
@@ -60,20 +61,28 @@ module LlmCostTracker
         end
       end
 
+      def provider_for(_request_url)
+        "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)
-        start_usage = find_event_value(events, reverse: true) do |data|
-          data.dig("message", "usage") if data["type"] == "message_start"
-        end
         latest_delta = find_event_value(events, reverse: true) do |data|
           data["usage"] if data["type"] == "message_delta" && data["usage"].is_a?(Hash)
         end
+        return nil unless latest_delta
 
-        return nil unless start_usage || latest_delta
+        start_usage = find_event_value(events, reverse: true) do |data|
+          data.dig("message", "usage") if data["type"] == "message_start"
+        end
 
-        (start_usage || {}).merge(latest_delta || {}) do |_key, start_val, delta_val|
-          delta_val.nil? ? start_val : delta_val
+        (start_usage || {}).merge(latest_delta) do |_key, start_val, delta_val|
+          delta_val || start_val
         end
       end
 
@@ -87,7 +96,44 @@ module LlmCostTracker
           model: model,
           token_usage: token_usage(usage: usage, cache_read: cache_read),
           stream: true,
-          usage_source: :stream_final
+          usage_source: :stream_final,
+          service_line_items: service_line_items(usage)
+        )
+      end
+
+      def service_line_items(usage)
+        server_tool_use = usage["server_tool_use"]
+        return [] unless server_tool_use.is_a?(Hash)
+
+        [
+          service_line_item(
+            component_key: :web_search_request,
+            quantity: server_tool_use["web_search_requests"],
+            provider_field: "usage.server_tool_use.web_search_requests"
+          ),
+          service_line_item(
+            component_key: :web_fetch_request,
+            quantity: server_tool_use["web_fetch_requests"],
+            provider_field: "usage.server_tool_use.web_fetch_requests"
+          ),
+          service_line_item(
+            component_key: :code_execution_request,
+            quantity: server_tool_use["code_execution_requests"],
+            provider_field: "usage.server_tool_use.code_execution_requests"
+          )
+        ].compact
+      end
+
+      def service_line_item(component_key:, quantity:, provider_field:)
+        quantity = quantity.to_i
+        return if quantity.zero?
+
+        Billing::LineItem.build(
+          component_key: component_key,
+          quantity: quantity,
+          cost_status: Billing::CostStatus::UNKNOWN,
+          pricing_basis: :provider_usage,
+          provider_field: provider_field
         )
       end
 
@@ -97,32 +143,46 @@ module LlmCostTracker
         cache_creation = usage["cache_creation"]
         if cache_creation.is_a?(Hash)
           cache_write = cache_creation["ephemeral_5m_input_tokens"].to_i
-          cache_write_1h = cache_creation["ephemeral_1h_input_tokens"].to_i
+          cache_write_extended = cache_creation["ephemeral_1h_input_tokens"].to_i
         else
+          warn_unexpected_cache_creation(cache_creation, usage)
           cache_write = usage["cache_creation_input_tokens"].to_i
-          cache_write_1h = 0
+          cache_write_extended = 0
         end
+        hidden_output = (
+          usage["thinking_tokens"] || usage["thinking_output_tokens"] ||
+            usage.dig("output_tokens_details", "reasoning_tokens")
+        ).to_i
 
         TokenUsage.build(
           input_tokens: input,
           output_tokens: output,
-          total_tokens: input + output + cache_read + cache_write + cache_write_1h,
-          cache_read_input_tokens: usage["cache_read_input_tokens"],
+          total_tokens: input + output + cache_read + cache_write + cache_write_extended,
+          cache_read_input_tokens: cache_read,
           cache_write_input_tokens: cache_write,
-          cache_write_1h_input_tokens: cache_write_1h
+          cache_write_extended_input_tokens: cache_write_extended,
+          hidden_output_tokens: hidden_output
         )
       end
 
+      def warn_unexpected_cache_creation(cache_creation, usage)
+        return if cache_creation.nil? || usage.key?("cache_creation_input_tokens")
+
+        Logging.warn("Anthropic usage.cache_creation has unexpected shape: #{cache_creation.class}")
+      end
+
       def pricing_mode(request:, response:, 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)
 
         modes << Pricing.normalize_mode(speed)
         modes << Pricing.normalize_mode(service_tier)
-        modes << "data_residency" if inference_geo(request: request, response: response, usage: usage) == "us"
+        geo = inference_geo(request: request, response: response, usage: usage).downcase
+        modes << "data_residency" if DATA_RESIDENCY_GEOS.include?(geo)
 
         modes = modes.compact.uniq
         modes.empty? ? nil : modes.join("_")

data/lib/llm_cost_tracker/parsers/base.rb

@@ -32,7 +32,9 @@ module LlmCostTracker
         nil
       end
 
-      private
+      def auto_enable_stream_usage?(_request_url)
+        false
+      end
 
       def safe_json_parse(body)
         return {} if body.blank?
@@ -42,6 +44,8 @@ module LlmCostTracker
         {}
       end
 
+      private
+
       def uri_matches?(url)
         uri = parsed_uri(url)
         uri ? yield(uri) : false
@@ -59,7 +63,10 @@ module LlmCostTracker
           )
           extra_match = block_given? ? yield(uri) : true
 
-          host_match && path_match && extra_match ? true : false
+          next false unless host_match && path_match
+          next false unless extra_match
+
+          true
         end
       end
 
@@ -100,7 +107,8 @@ module LlmCostTracker
         nil
       end
 
-      def build_unknown_stream_usage(provider:, model:, provider_response_id:, pricing_mode: nil)
+      def build_unknown_stream_usage(provider:, model:, provider_response_id:, pricing_mode: nil,
+                                     service_line_items: nil)
         UsageCapture.build(
           provider: provider,
           provider_response_id: provider_response_id,
@@ -108,7 +116,8 @@ module LlmCostTracker
           model: model || UsageCapture::UNKNOWN_MODEL,
           token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
-          usage_source: :unknown
+          usage_source: :unknown,
+          service_line_items: service_line_items
         )
       end
     end

data/lib/llm_cost_tracker/parsers/gemini.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require_relative "../billing/line_item"
 require_relative "base"
 
 module LlmCostTracker
@@ -8,6 +9,7 @@ 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)
@@ -31,12 +33,14 @@ module LlmCostTracker
         return nil unless usage
 
         request = safe_json_parse(request_body)
+        model = extract_model_from_url(request_url)
         build_usage_capture(
           request_url: request_url,
           usage: usage,
           usage_source: :response,
           provider_response_id: response["responseId"],
-          pricing_mode: pricing_mode(request: request, response_headers: response_headers)
+          pricing_mode: pricing_mode(request: request, response_headers: response_headers),
+          service_line_items: grounding_line_items_for_response(response, model: model)
         )
       end
 
@@ -48,6 +52,7 @@ module LlmCostTracker
         model = extract_model_from_url(request_url)
         response_id = stream_response_id(events)
         mode = pricing_mode(request: request, response_headers: response_headers)
+        service_line_items = grounding_line_items_for_stream(events, model: model)
 
         if usage
           build_usage_capture(
@@ -56,39 +61,51 @@ module LlmCostTracker
             stream: true,
             usage_source: :stream_final,
             provider_response_id: response_id,
-            pricing_mode: mode
+            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
+            pricing_mode: mode,
+            service_line_items: service_line_items
           )
         end
       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)
+                              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(
           provider: "gemini",
           model: extract_model_from_url(request_url),
           pricing_mode: pricing_mode,
           token_usage: TokenUsage.build(
-            input_tokens: [usage["promptTokenCount"].to_i - cache_read, 0].max + tool_use_prompt,
-            output_tokens: output_tokens(usage),
-            total_tokens: total_tokens(usage: usage, cache_read: cache_read, tool_use_prompt: tool_use_prompt),
-            cache_read_input_tokens: usage["cachedContentTokenCount"],
+            input_tokens: regular_input_tokens(usage: usage, cache_read: cache_read, audio_input: audio_input) +
+                          tool_use_prompt,
+            output_tokens: regular_output_tokens(usage: usage, audio_output: audio_output),
+            total_tokens: usage["totalTokenCount"],
+            cache_read_input_tokens: cache_read,
+            audio_input_tokens: audio_input,
+            audio_output_tokens: audio_output,
             hidden_output_tokens: usage["thoughtsTokenCount"]
           ),
           stream: stream,
           usage_source: usage_source,
-          provider_response_id: provider_response_id
+          provider_response_id: provider_response_id,
+          service_line_items: service_line_items
         )
       end
 
@@ -100,14 +117,41 @@ module LlmCostTracker
       end
 
       def output_tokens(usage)
-        usage["candidatesTokenCount"].to_i
+        (usage["candidatesTokenCount"] || usage["responseTokenCount"]).to_i + usage["thoughtsTokenCount"].to_i
+      end
+
+      def regular_input_tokens(usage:, cache_read:, audio_input:)
+        [usage["promptTokenCount"].to_i - cache_read - audio_input, 0].max
       end
 
-      def total_tokens(usage:, cache_read:, tool_use_prompt:)
-        total = usage["totalTokenCount"]
-        return total.to_i unless total.nil?
+      def regular_output_tokens(usage:, audio_output:)
+        [output_tokens(usage) - audio_output, 0].max
+      end
 
-        [usage["promptTokenCount"].to_i - cache_read, 0].max + cache_read + tool_use_prompt + output_tokens(usage)
+      def audio_input_tokens(usage)
+        prompt_audio = modality_tokens(usage["promptTokensDetails"] || usage["prompt_tokens_details"], "AUDIO")
+        cache_audio = modality_tokens(usage["cacheTokensDetails"] || usage["cache_tokens_details"], "AUDIO")
+        [prompt_audio - cache_audio, 0].max
+      end
+
+      def audio_output_tokens(usage)
+        modality_tokens(
+          usage["candidatesTokensDetails"] ||
+            usage["candidates_tokens_details"] ||
+            usage["responseTokensDetails"] ||
+            usage["response_tokens_details"],
+          "AUDIO"
+        )
+      end
+
+      def modality_tokens(details, modality)
+        Array(details).sum do |detail|
+          next 0 unless detail.is_a?(Hash)
+
+          next 0 unless detail["modality"] == modality
+
+          (detail["tokenCount"] || detail["token_count"]).to_i
+        end
       end
 
       def stream_response_id(events)
@@ -133,12 +177,58 @@ module LlmCostTracker
           request.dig("config", "service_tier") ||
           request.dig("config", "serviceTier")
         )
-        request_mode == "flex" ? request_mode : nil
+        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_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"])
+          count if count.positive?
+        end
+        grounding_line_items(quantity || 0, model: model)
+      end
+
+      def grounding_request_count(candidates)
+        Array(candidates).sum do |candidate|
+          next 0 unless candidate.is_a?(Hash)
+
+          metadata = candidate["groundingMetadata"] || candidate["grounding_metadata"] || {}
+          queries = metadata["webSearchQueries"] || metadata["web_search_queries"] || []
+          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)
+        [
+          Billing::LineItem.build(
+            component_key: :grounding_request,
+            quantity: billed_quantity,
+            cost_status: Billing::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:)
+        per_query_billing?(model) ? query_count : 1
+      end
+
+      def per_query_billing?(model)
+        model.to_s.match?(PER_QUERY_GROUNDING_MODEL_PATTERN)
+      end
     end
   end
 end

data/lib/llm_cost_tracker/parsers/openai.rb

@@ -21,7 +21,19 @@ module LlmCostTracker
         gb.api.openai.com
         ae.api.openai.com
       ].freeze
-      TRACKED_PATHS = %w[/v1/chat/completions /v1/completions /v1/embeddings /v1/responses].freeze
+      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
 
       def match?(url)
         match_uri?(url, hosts: HOSTS, exact_paths: TRACKED_PATHS)
@@ -49,7 +61,9 @@ module LlmCostTracker
         )
       end
 
-      private
+      def auto_enable_stream_usage?(request_url)
+        openai_chat_completions_url?(request_url)
+      end
 
       def provider_for(_request_url)
         "openai"