llm_cost_tracker 0.7.0 → 0.7.2

data/lib/llm_cost_tracker/middleware/faraday.rb

@@ -2,10 +2,10 @@
 
 require "faraday"
 require "json"
+require "uri"
 
 require_relative "../logging"
-require_relative "../request_url"
-require_relative "../stream_capture"
+require_relative "../capture/stream"
 
 module LlmCostTracker
   module Middleware
@@ -20,51 +20,59 @@ module LlmCostTracker
 
         request_url  = request_env.url.to_s
         request_body = read_body(request_env.body) || ""
-        parser       = Parsers::Registry.find_for(request_url)
+        parser       = Parsers.find_for(request_url)
         streaming    = parser&.streaming_request?(request_url, request_body)
         stream_buffer = install_stream_tap(request_env) if streaming
 
         Tracker.enforce_budget! if parser
-        started_at = monotonic_time
+        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: elapsed_ms(started_at),
+            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(
-          provider: parsed.provider,
-          model: parsed.model,
-          input_tokens: parsed.input_tokens,
-          output_tokens: parsed.output_tokens,
+          capture: parsed,
           latency_ms: latency_ms,
-          stream: parsed.stream,
-          usage_source: parsed.usage_source,
-          provider_response_id: parsed.provider_response_id,
-          metadata: resolved_tags(request_env).merge(parsed.metadata)
+          metadata: metadata,
+          context_tags: context_tags
         )
       rescue LlmCostTracker::Error
         raise
@@ -72,49 +80,72 @@ 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(
-            "Unable to read response body for #{RequestUrl.label(request_url)}; " \
+            "Unable to read response body for #{request_url_label(request_url)}; " \
             "known streaming responses are captured automatically, or via LlmCostTracker.track_stream " \
             "for custom clients."
           )
           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
-        body = read_body(response_env.body) if body.nil? || body.empty?
+        body = read_body(response_env.body) if body.blank?
 
-        if body.nil? || body.empty?
+        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)
-        return nil unless request_env.respond_to?(:request) && request_env.request
+        request = request_env.try(:request)
+        return nil unless request
 
-        original = request_env.request.on_data
+        original = request.on_data
         return nil unless original
 
         state = { buffer: StringIO.new, bytes: 0, overflowed: false }
-        request_env.request.on_data = proc do |chunk, size, env|
+        request.on_data = proc do |chunk, size, env|
           chunk = chunk.to_s
           unless state[:overflowed]
-            if state[:bytes] + chunk.bytesize <= StreamCapture::LIMIT_BYTES
+            if state[:bytes] + chunk.bytesize <= Capture::Stream::LIMIT_BYTES
               state[:buffer] << chunk
               state[:bytes] += chunk.bytesize
             else
@@ -136,38 +167,49 @@ module LlmCostTracker
         when nil then ""
         when Hash, Array then body.to_json
         else
-          body.respond_to?(:to_str) ? body.to_str : nil
+          body.try(:to_str)
         end
       end
 
       def resolved_tags(request_env)
-        tags = @tags.respond_to?(:call) ? call_tags(request_env) : @tags
+        tags =
+          if @tags.respond_to?(:call)
+            @tags.arity.zero? ? @tags.call : @tags.call(request_env)
+          else
+            @tags
+          end
         return {} if tags.nil?
 
         tags.to_h
       end
 
-      def call_tags(request_env)
-        @tags.arity.zero? ? @tags.call : @tags.call(request_env)
-      end
-
-      def monotonic_time
-        Process.clock_gettime(Process::CLOCK_MONOTONIC)
-      end
-
-      def elapsed_ms(started_at)
-        ((monotonic_time - started_at) * 1000).round
+      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 #{RequestUrl.label(request_url)}; " \
+          return "Unable to capture streaming response for #{request_url_label(request_url)}; " \
                  "recording usage_source=unknown. Use LlmCostTracker.track_stream for manual capture."
         end
 
-        "Streaming response for #{RequestUrl.label(request_url)} exceeded #{StreamCapture::LIMIT_BYTES} bytes; " \
+        "Streaming response for #{request_url_label(request_url)} exceeded #{Capture::Stream::LIMIT_BYTES} bytes; " \
           "recording usage_source=unknown. Use LlmCostTracker.track_stream for manual capture."
       end
+
+      def request_url_label(value)
+        uri = URI.parse(value.to_s)
+        uri.query = nil
+        uri.fragment = nil
+        uri.try(:user=, nil)
+        uri.try(:password=, nil)
+        uri.to_s
+      rescue URI::InvalidURIError
+        value.to_s.split("?", 2).first
+      end
     end
   end
 end

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)
@@ -24,36 +24,38 @@ module LlmCostTracker
 
         request = safe_json_parse(request_body)
         cache_read = usage["cache_read_input_tokens"].to_i
-        cache_write = usage["cache_creation_input_tokens"].to_i
 
-        ParsedUsage.build(
+        UsageCapture.build(
           provider: "anthropic",
           provider_response_id: response["id"],
+          pricing_mode: pricing_mode(request: request, response: response, usage: usage),
           model: response["model"] || request["model"],
-          input_tokens: usage["input_tokens"].to_i,
-          output_tokens: usage["output_tokens"].to_i,
-          total_tokens: usage["input_tokens"].to_i + usage["output_tokens"].to_i + cache_read + cache_write,
-          cache_read_input_tokens: usage["cache_read_input_tokens"],
-          cache_write_input_tokens: usage["cache_creation_input_tokens"],
+          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)
-        model = stream_model(events) || request["model"]
+        model = find_event_value(events) { |data| data.dig("message", "model") } || request["model"]
         usage = stream_usage(events)
-        response_id = stream_response_id(events)
+        response_id = find_event_value(events) { |data| data.dig("message", "id") || data["id"] }
 
         if usage
-          build_stream_result(model, usage, response_id)
+          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
+            provider_response_id: response_id,
+            pricing_mode: pricing_mode(request: request, response: nil, usage: usage)
           )
         end
       end
@@ -75,33 +77,64 @@ module LlmCostTracker
         end
       end
 
-      def stream_model(events)
-        find_event_value(events) { |data| data.dig("message", "model") }
-      end
+      def build_stream_result(model:, usage:, response_id:, pricing_mode:)
+        cache_read = usage["cache_read_input_tokens"].to_i
 
-      def stream_response_id(events)
-        find_event_value(events) { |data| data.dig("message", "id") || data["id"] }
+        UsageCapture.build(
+          provider: "anthropic",
+          provider_response_id: response_id,
+          pricing_mode: pricing_mode,
+          model: model,
+          token_usage: token_usage(usage: usage, cache_read: cache_read),
+          stream: true,
+          usage_source: :stream_final
+        )
       end
 
-      def build_stream_result(model, usage, response_id)
+      def token_usage(usage:, cache_read:)
         input = usage["input_tokens"].to_i
         output = usage["output_tokens"].to_i
-        cache_read = usage["cache_read_input_tokens"].to_i
-        cache_write = usage["cache_creation_input_tokens"].to_i
+        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
+        else
+          cache_write = usage["cache_creation_input_tokens"].to_i
+          cache_write_1h = 0
+        end
 
-        ParsedUsage.build(
-          provider: "anthropic",
-          provider_response_id: response_id,
-          model: model,
+        TokenUsage.build(
           input_tokens: input,
           output_tokens: output,
-          total_tokens: input + output + cache_read + cache_write,
+          total_tokens: input + output + cache_read + cache_write + cache_write_1h,
           cache_read_input_tokens: usage["cache_read_input_tokens"],
-          cache_write_input_tokens: usage["cache_creation_input_tokens"],
-          stream: true,
-          usage_source: :stream_final
+          cache_write_input_tokens: cache_write,
+          cache_write_1h_input_tokens: cache_write_1h
         )
       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"]
+
+        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
 end

data/lib/llm_cost_tracker/parsers/base.rb

@@ -1,12 +1,13 @@
 # frozen_string_literal: true
 
+require "active_support/core_ext/object/blank"
 require "json"
 require "uri"
 
 module LlmCostTracker
   module Parsers
     class Base
-      def parse(request_url, request_body, response_status, response_body)
+      def parse(**)
         raise NotImplementedError
       end
 
@@ -19,23 +20,22 @@ module LlmCostTracker
       end
 
       def streaming_request?(_request_url, request_body)
-        return false if request_body.nil?
+        return false if request_body.blank?
 
         body = request_body.to_s
-        return false if body.empty?
 
         request = safe_json_parse(body)
         request.is_a?(Hash) && request["stream"] == true
       end
 
-      def parse_stream(_request_url, _request_body, _response_status, _events)
+      def parse_stream(**)
         nil
       end
 
       private
 
       def safe_json_parse(body)
-        return {} if body.nil? || body.empty?
+        return {} if body.blank?
 
         JSON.parse(body)
       rescue JSON::ParserError
@@ -49,7 +49,7 @@ module LlmCostTracker
 
       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? || host_matches?(uri, hosts)
+          host_match = hosts.nil? || hosts.include?(uri.host.to_s.downcase)
           path_match = path_matches?(
             uri,
             exact_paths: exact_paths,
@@ -69,10 +69,6 @@ module LlmCostTracker
         nil
       end
 
-      def host_matches?(uri, hosts)
-        hosts.include?(uri.host.to_s.downcase)
-      end
-
       def path_matches?(uri, exact_paths: nil, path_includes: nil, path_suffixes: nil, path_pattern: nil)
         path = uri.path.to_s
         matches = true
@@ -98,28 +94,23 @@ module LlmCostTracker
       def find_event_value(events, reverse: false)
         each_event_data(events, reverse:) do |data|
           value = yield(data)
-          return value if event_value_present?(value)
+          return value if value.present?
         end
 
         nil
       end
 
-      def build_unknown_stream_usage(provider:, model:, provider_response_id:)
-        ParsedUsage.build(
+      def build_unknown_stream_usage(provider:, model:, provider_response_id:, pricing_mode: nil)
+        UsageCapture.build(
           provider: provider,
           provider_response_id: provider_response_id,
-          model: model || ParsedUsage::UNKNOWN_MODEL,
-          input_tokens: 0,
-          output_tokens: 0,
-          total_tokens: 0,
+          pricing_mode: pricing_mode,
+          model: model || UsageCapture::UNKNOWN_MODEL,
+          token_usage: TokenUsage.build(input_tokens: 0, output_tokens: 0, total_tokens: 0),
           stream: true,
           usage_source: :unknown
         )
       end
-
-      def event_value_present?(value)
-        !value.nil? && (!value.respond_to?(:empty?) || !value.empty?)
-      end
     end
   end
 end

data/lib/llm_cost_tracker/parsers/gemini.rb

@@ -18,63 +18,74 @@ module LlmCostTracker
       end
 
       def streaming_request?(request_url, request_body)
-        return true if streaming_url?(request_url)
+        return true if match_uri?(request_url, path_pattern: STREAM_PATH_PATTERN)
 
         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
 
-        build_parsed_usage(
-          request_url,
-          usage,
+        request = safe_json_parse(request_body)
+        build_usage_capture(
+          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_parsed_usage(
-            request_url,
-            usage,
+          build_usage_capture(
+            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_parsed_usage(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
 
-        ParsedUsage.build(
+        UsageCapture.build(
           provider: "gemini",
           model: extract_model_from_url(request_url),
-          input_tokens: [usage["promptTokenCount"].to_i - cache_read, 0].max,
-          output_tokens: output_tokens(usage),
-          total_tokens: total_tokens(usage, cache_read),
-          cache_read_input_tokens: usage["cachedContentTokenCount"],
-          hidden_output_tokens: usage["thoughtsTokenCount"],
+          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"],
+            hidden_output_tokens: usage["thoughtsTokenCount"]
+          ),
           stream: stream,
           usage_source: usage_source,
           provider_response_id: provider_response_id
@@ -92,21 +103,17 @@ module LlmCostTracker
         usage["candidatesTokenCount"].to_i + usage["thoughtsTokenCount"].to_i
       end
 
-      def total_tokens(usage, cache_read)
+      def total_tokens(usage:, cache_read:, tool_use_prompt:)
         total = usage["totalTokenCount"]
         return total.to_i unless total.nil?
 
-        [usage["promptTokenCount"].to_i - cache_read, 0].max + cache_read + output_tokens(usage)
+        [usage["promptTokenCount"].to_i - cache_read, 0].max + cache_read + tool_use_prompt + output_tokens(usage)
       end
 
       def stream_response_id(events)
         find_event_value(events) { |data| data["responseId"] }
       end
 
-      def streaming_url?(request_url)
-        match_uri?(request_url, path_pattern: STREAM_PATH_PATTERN)
-      end
-
       def extract_model_from_url(url)
         uri = parsed_uri(url)
         return nil unless uri
@@ -114,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