llm_cost_tracker 0.7.1 → 0.7.3

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -25,40 +25,54 @@ module LlmCostTracker
         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)
 
         @app.call(request_env).on_complete do |response_env|
           process(
             parser: parser,
-            request_env: request_env,
             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
+            stream_buffer: stream_buffer,
+            context_tags: context_tags,
+            metadata: metadata
           )
         end
       end
 
       private
 
-      def process(parser:, request_env:, request_url:, request_body:, response_env:,
-                  latency_ms:, streaming:, stream_buffer:)
+      def process(parser:, request_url:, request_body:, response_env:,
+                  latency_ms:, streaming:, stream_buffer:, context_tags:, metadata:)
         return unless parser
 
         parsed =
           if streaming
-            parse_stream(parser, request_url, request_body, response_env, stream_buffer)
+            parse_stream(
+              parser: parser,
+              request_url: request_url,
+              request_body: request_body,
+              response_env: response_env,
+              stream_buffer: stream_buffer
+            )
           else
-            parse_response(parser, request_url, request_body, response_env)
+            parse_response(
+              parser: parser,
+              request_url: request_url,
+              request_body: request_body,
+              response_env: response_env
+            )
           end
         return unless parsed
 
         Tracker.record(
           capture: parsed,
           latency_ms: latency_ms,
-          metadata: resolved_tags(request_env)
+          metadata: metadata,
+          context_tags: context_tags
         )
       rescue LlmCostTracker::Error
         raise
@@ -66,7 +80,7 @@ module LlmCostTracker
         Logging.warn("Error processing response: #{e.class}: #{e.message}")
       end
 
-      def parse_response(parser, request_url, request_body, response_env)
+      def parse_response(parser:, request_url:, request_body:, response_env:)
         response_body = read_body(response_env.body)
         unless response_body
           Logging.warn(
@@ -77,13 +91,24 @@ module LlmCostTracker
           return nil
         end
 
-        parser.parse(request_url, request_body, response_env.status, response_body)
+        parser.parse(
+          request_url: request_url,
+          request_body: request_body,
+          response_status: response_env.status,
+          response_body: response_body,
+          response_headers: response_env.response_headers
+        )
       end
 
-      def parse_stream(parser, request_url, request_body, response_env, stream_buffer)
+      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_body, response_env.status, [])
+          return parser.parse_stream(
+            request_url: request_url,
+            request_body: request_body,
+            response_status: response_env.status,
+            response_headers: response_env.response_headers
+          )
         end
 
         body = stream_buffer&.dig(:buffer)&.string
@@ -91,11 +116,22 @@ module LlmCostTracker
 
         if body.blank?
           Logging.warn(capture_warning(request_url, stream_buffer))
-          return parser.parse_stream(request_url, request_body, response_env.status, [])
+          return parser.parse_stream(
+            request_url: request_url,
+            request_body: request_body,
+            response_status: response_env.status,
+            response_headers: response_env.response_headers
+          )
         end
 
         events = Parsers::SSE.parse(body)
-        parser.parse_stream(request_url, request_body, response_env.status, events)
+        parser.parse_stream(
+          request_url: request_url,
+          request_body: request_body,
+          response_status: response_env.status,
+          events: events,
+          response_headers: response_env.response_headers
+        )
       end
 
       def install_stream_tap(request_env)
@@ -147,6 +183,13 @@ module LlmCostTracker
         tags.to_h
       end
 
+      def tag_snapshot(request_env)
+        [LlmCostTracker::Tags::Context.tags, resolved_tags(request_env)]
+      rescue StandardError => e
+        Logging.warn("Error resolving request tags: #{e.class}: #{e.message}")
+        [{}, {}]
+      end
+
       def capture_warning(request_url, stream_buffer)
         unless stream_buffer&.dig(:overflowed)
           return "Unable to capture streaming response for #{request_url_label(request_url)}; " \

data/lib/llm_cost_tracker/parsers/anthropic.rb

@@ -15,7 +15,7 @@ module LlmCostTracker
         %w[anthropic]
       end
 
-      def parse(_request_url, request_body, response_status, response_body)
+      def parse(request_body:, response_status:, response_body:, **)
         return nil unless response_status == 200
 
         response = safe_json_parse(response_body)
@@ -28,14 +28,14 @@ module LlmCostTracker
         UsageCapture.build(
           provider: "anthropic",
           provider_response_id: response["id"],
-          pricing_mode: pricing_mode(request, response, usage),
+          pricing_mode: pricing_mode(request: request, response: response, usage: usage),
           model: response["model"] || request["model"],
-          token_usage: token_usage(usage, cache_read),
+          token_usage: token_usage(usage: usage, cache_read: cache_read),
           usage_source: :response
         )
       end
 
-      def parse_stream(_request_url, request_body, response_status, events)
+      def parse_stream(response_status:, request_body: nil, events: [], **)
         return nil unless response_status == 200
 
         request = safe_json_parse(request_body)
@@ -44,13 +44,18 @@ module LlmCostTracker
         response_id = find_event_value(events) { |data| data.dig("message", "id") || data["id"] }
 
         if usage
-          build_stream_result(model, usage, response_id, pricing_mode(request, nil, usage))
+          build_stream_result(
+            model: model,
+            usage: usage,
+            response_id: response_id,
+            pricing_mode: pricing_mode(request: request, response: nil, usage: usage)
+          )
         else
           build_unknown_stream_usage(
             provider: "anthropic",
             model: model,
             provider_response_id: response_id,
-            pricing_mode: pricing_mode(request, nil, usage)
+            pricing_mode: pricing_mode(request: request, response: nil, usage: usage)
           )
         end
       end
@@ -72,7 +77,7 @@ module LlmCostTracker
         end
       end
 
-      def build_stream_result(model, usage, response_id, pricing_mode)
+      def build_stream_result(model:, usage:, response_id:, pricing_mode:)
         cache_read = usage["cache_read_input_tokens"].to_i
 
         UsageCapture.build(
@@ -80,13 +85,13 @@ module LlmCostTracker
           provider_response_id: response_id,
           pricing_mode: pricing_mode,
           model: model,
-          token_usage: token_usage(usage, cache_read),
+          token_usage: token_usage(usage: usage, cache_read: cache_read),
           stream: true,
           usage_source: :stream_final
         )
       end
 
-      def token_usage(usage, cache_read)
+      def token_usage(usage:, cache_read:)
         input = usage["input_tokens"].to_i
         output = usage["output_tokens"].to_i
         cache_creation = usage["cache_creation"]
@@ -108,10 +113,27 @@ module LlmCostTracker
         )
       end
 
-      def pricing_mode(request, response, usage)
-        usage&.fetch("service_tier", nil) ||
-          response&.fetch("service_tier", nil) ||
-          request["service_tier"]
+      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"]
+
+        modes << Pricing.normalize_mode(speed)
+        modes << Pricing.normalize_mode(service_tier)
+        modes << "data_residency" if inference_geo(request: request, response: response, usage: usage) == "us"
+
+        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
       end
     end
   end

data/lib/llm_cost_tracker/parsers/base.rb

@@ -7,7 +7,7 @@ require "uri"
 module LlmCostTracker
   module Parsers
     class Base
-      def parse(request_url, request_body, response_status, response_body)
+      def parse(**)
         raise NotImplementedError
       end
 
@@ -28,7 +28,7 @@ module LlmCostTracker
         request.is_a?(Hash) && request["stream"] == true
       end
 
-      def parse_stream(_request_url, _request_body, _response_status, _events)
+      def parse_stream(**)
         nil
       end
 

data/lib/llm_cost_tracker/parsers/gemini.rb

@@ -23,58 +23,66 @@ module LlmCostTracker
         super
       end
 
-      def parse(request_url, _request_body, response_status, response_body)
+      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)
         build_usage_capture(
-          request_url,
-          usage,
+          request_url: request_url,
+          usage: usage,
           usage_source: :response,
-          provider_response_id: response["responseId"]
+          provider_response_id: response["responseId"],
+          pricing_mode: pricing_mode(request: request, response_headers: response_headers)
         )
       end
 
-      def parse_stream(request_url, _request_body, response_status, events)
+      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, response_headers: response_headers)
 
         if usage
           build_usage_capture(
-            request_url,
-            usage,
+            request_url: request_url,
+            usage: usage,
             stream: true,
             usage_source: :stream_final,
-            provider_response_id: response_id
+            provider_response_id: response_id,
+            pricing_mode: mode
           )
         else
           build_unknown_stream_usage(
             provider: "gemini",
             model: model,
-            provider_response_id: response_id
+            provider_response_id: response_id,
+            pricing_mode: mode
           )
         end
       end
 
       private
 
-      def build_usage_capture(request_url, usage, usage_source:, stream: false, provider_response_id: nil)
+      def build_usage_capture(request_url:, usage:, usage_source:, stream: false, provider_response_id: nil,
+                              pricing_mode: nil)
         cache_read = usage["cachedContentTokenCount"].to_i
         tool_use_prompt = usage["toolUsePromptTokenCount"].to_i
 
         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, cache_read, tool_use_prompt),
+            total_tokens: total_tokens(usage: usage, cache_read: cache_read, tool_use_prompt: tool_use_prompt),
             cache_read_input_tokens: usage["cachedContentTokenCount"],
             hidden_output_tokens: usage["thoughtsTokenCount"]
           ),
@@ -92,10 +100,10 @@ module LlmCostTracker
       end
 
       def output_tokens(usage)
-        usage["candidatesTokenCount"].to_i + usage["thoughtsTokenCount"].to_i
+        usage["candidatesTokenCount"].to_i
       end
 
-      def total_tokens(usage, cache_read, tool_use_prompt)
+      def total_tokens(usage:, cache_read:, tool_use_prompt:)
         total = usage["totalTokenCount"]
         return total.to_i unless total.nil?
 
@@ -113,6 +121,24 @@ module LlmCostTracker
         match = uri.path.match(%r{/models/([^/:]+)})
         match && match[1]
       end
+
+      def pricing_mode(request:, response_headers:)
+        response_tier = response_header(response_headers, "x-gemini-service-tier")
+        response_mode = Pricing.normalize_mode(response_tier)
+        return response_mode if response_mode
+
+        request_mode = Pricing.normalize_mode(
+          request["service_tier"] ||
+          request["serviceTier"] ||
+          request.dig("config", "service_tier") ||
+          request.dig("config", "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
     end
   end
 end

data/lib/llm_cost_tracker/parsers/openai.rb

@@ -8,7 +8,19 @@ module LlmCostTracker
     class Openai < Base
       include OpenaiUsage
 
-      HOSTS = %w[api.openai.com].freeze
+      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 /v1/embeddings /v1/responses].freeze
 
       def match?(url)
@@ -19,12 +31,22 @@ module LlmCostTracker
         %w[openai]
       end
 
-      def parse(request_url, request_body, response_status, response_body)
-        parse_openai_usage(request_url, request_body, response_status, response_body)
+      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(request_url, request_body, response_status, events)
-        parse_openai_stream_usage(request_url, request_body, response_status, events)
+      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

data/lib/llm_cost_tracker/parsers/openai_compatible.rb

@@ -21,12 +21,22 @@ module LlmCostTracker
         ].uniq.freeze
       end
 
-      def parse(request_url, request_body, response_status, response_body)
-        parse_openai_usage(request_url, request_body, response_status, response_body)
+      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(request_url, request_body, response_status, events)
-        parse_openai_stream_usage(request_url, request_body, response_status, events)
+      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

data/lib/llm_cost_tracker/parsers/openai_usage.rb

@@ -5,7 +5,7 @@ module LlmCostTracker
     module OpenaiUsage
       private
 
-      def parse_openai_usage(request_url, request_body, response_status, response_body)
+      def parse_openai_usage(request_url:, request_body:, response_status:, response_body:)
         return nil unless response_status == 200
 
         response = safe_json_parse(response_body)
@@ -15,17 +15,23 @@ module LlmCostTracker
         request = safe_json_parse(request_body)
         cache_read = cache_read_input_tokens(usage)
 
+        model = response["model"] || request["model"]
+
         UsageCapture.build(
           provider: provider_for(request_url),
           provider_response_id: response["id"],
-          pricing_mode: response["service_tier"] || request["service_tier"],
-          model: response["model"] || request["model"],
-          token_usage: token_usage(usage, cache_read),
+          pricing_mode: pricing_mode(
+            request_url: request_url,
+            model: model,
+            service_tier: response["service_tier"] || request["service_tier"]
+          ),
+          model: model,
+          token_usage: token_usage(usage: usage, cache_read: cache_read),
           usage_source: :response
         )
       end
 
-      def parse_openai_stream_usage(request_url, request_body, response_status, events)
+      def parse_openai_stream_usage(response_status:, request_url: nil, request_body: nil, events: [])
         return nil unless response_status == 200
 
         request = safe_json_parse(request_body)
@@ -33,7 +39,11 @@ module LlmCostTracker
           find_event_value(events) { |data| data["model"] || data.dig("response", "model") } || request["model"]
         usage = detect_stream_usage(events)
         response_id = find_event_value(events) { |data| data["id"] || data.dig("response", "id") }
-        pricing_mode = stream_pricing_mode(events) || request["service_tier"]
+        pricing_mode = pricing_mode(
+          request_url: request_url,
+          model: model,
+          service_tier: stream_pricing_mode(events) || request["service_tier"]
+        )
 
         if usage
           cache_read = cache_read_input_tokens(usage)
@@ -42,7 +52,7 @@ module LlmCostTracker
             provider_response_id: response_id,
             pricing_mode: pricing_mode,
             model: model,
-            token_usage: token_usage(usage, cache_read),
+            token_usage: token_usage(usage: usage, cache_read: cache_read),
             stream: true,
             usage_source: :stream_final
           )
@@ -69,17 +79,35 @@ module LlmCostTracker
         end
       end
 
-      def token_usage(usage, cache_read)
+      def pricing_mode(request_url:, model:, service_tier:)
+        modes = [Pricing.normalize_mode(service_tier)]
+        modes << "data_residency" if openai_regional_processing?(request_url: request_url, model: model)
+        modes = modes.compact.uniq
+        modes.empty? ? nil : modes.join("_")
+      end
+
+      def openai_regional_processing?(request_url:, model:)
+        uri = parsed_uri(request_url)
+        return false unless %w[us.api.openai.com eu.api.openai.com].include?(uri&.host.to_s.downcase)
+
+        openai_data_residency_model?(model)
+      end
+
+      def openai_data_residency_model?(model)
+        model.to_s.match?(/\Agpt-5\.(?:4|5)(?:-(?:mini|nano|pro))?(?:-\d{4}-\d{2}-\d{2})?\z/)
+      end
+
+      def token_usage(usage:, cache_read:)
         TokenUsage.build(
-          input_tokens: regular_input_tokens(usage, cache_read),
+          input_tokens: regular_input_tokens(usage: usage, cache_read: cache_read),
           output_tokens: (usage["completion_tokens"] || usage["output_tokens"]).to_i,
-          total_tokens: total_tokens(usage, cache_read),
+          total_tokens: total_tokens(usage: usage, cache_read: cache_read),
           cache_read_input_tokens: cache_read,
           hidden_output_tokens: hidden_output_tokens(usage)
         )
       end
 
-      def regular_input_tokens(usage, cache_read)
+      def regular_input_tokens(usage:, cache_read:)
         [(usage["prompt_tokens"] || usage["input_tokens"]).to_i - cache_read.to_i, 0].max
       end
 
@@ -93,11 +121,11 @@ module LlmCostTracker
         details["reasoning_tokens"]
       end
 
-      def total_tokens(usage, cache_read)
+      def total_tokens(usage:, cache_read:)
         total = usage["total_tokens"]
         return total.to_i unless total.nil?
 
-        regular_input_tokens(usage, cache_read) +
+        regular_input_tokens(usage: usage, cache_read: cache_read) +
           cache_read.to_i +
           (usage["completion_tokens"] || usage["output_tokens"]).to_i
       end