llm_cost_tracker 0.14.2 → 0.15.0

data/lib/llm_cost_tracker/integrations/ruby_llm/v1/reply.rb

@@ -0,0 +1,155 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Integrations
+    module RubyLlm
+      module V1
+        class Reply
+          KEPT_BODY = :@llm_cost_tracker_body
+
+          attr_reader :slug, :model
+
+          def self.body(response)
+            raw = response.try(:raw)
+            body = raw.respond_to?(:body) ? raw.body : raw
+            body = (raw || response).instance_variable_get(KEPT_BODY) unless body.is_a?(Hash)
+            body.is_a?(Hash) ? body : {}
+          end
+
+          def self.response_id(response)
+            body = body(response)
+            body["id"] || body["responseId"]
+          end
+
+          def initialize(provider, response, requested_model)
+            @provider = provider
+            @response = response
+            @slug = provider.slug.to_s
+            @model = (response.try(:model_id) || response.try(:model))&.to_s || requested_model
+            @body = Reply.body(response)
+          end
+
+          def usage = @usage ||= (@body["usage"] || {}).deep_symbolize_keys
+
+          def event(stream:, output_tokens: nil, audio_input_tokens: 0, service_line_items: [], usage_source: nil)
+            counts = token_counts
+            output_tokens = counts[:output] if output_tokens.nil?
+            return if counts[:input].nil? && output_tokens.nil? && service_line_items.empty? && usage_source.nil?
+
+            Event.build(
+              provider: @slug,
+              model: @model,
+              pricing_mode: pricing_mode,
+              token_usage: token_usage(counts, output_tokens, audio_input_tokens),
+              service_line_items: service_line_items + gemini_line_items,
+              stream: stream,
+              usage_source: usage_source || Usage::Source::SDK_RESPONSE,
+              provider_response_id: Reply.response_id(@response)
+            )
+          end
+
+          def service_line_items
+            case @slug
+            when "anthropic"
+              server_tool_use = @body.dig("usage", "server_tool_use")&.symbolize_keys
+              Providers::Anthropic::UsageExtractor.service_line_items(server_tool_use: server_tool_use)
+            when "openai" then Providers::Openai::ServiceCharges.service_line_items_for(@body, model: @model)
+            when "gemini" then Providers::Gemini::Parser.new.service_line_items_for(@body, model: @model)
+            when "openrouter", "xai", "perplexity" then Providers::Openai::ServiceCharges.billed_line_items(usage)
+            else []
+            end
+          end
+
+          def token_counts
+            tokens = @response.try(:tokens)
+            return { input: @response.try(:input_tokens), output: @response.try(:output_tokens) } unless tokens
+
+            usage = @body["usage"] || {}
+            { input: input_tokens(tokens, usage), output: output_tokens(tokens, usage), cache_read: tokens.cache_read,
+              cache_write: tokens.cache_write, thinking: tokens.thinking }
+          end
+
+          private
+
+          def token_usage(counts, output_tokens, audio_input_tokens)
+            gemini = gemini_usage
+            return Providers::Gemini::UsageExtractor.token_usage(gemini) if gemini
+
+            five_minute, one_hour = cache_writes(counts[:cache_write])
+            Usage::TokenUsage.build(
+              input_tokens: counts[:input].to_i - audio_input_tokens,
+              audio_input_tokens: audio_input_tokens,
+              output_tokens: output_tokens.to_i,
+              cache_read_input_tokens: counts[:cache_read].to_i,
+              cache_write_input_tokens: five_minute,
+              cache_write_extended_input_tokens: one_hour,
+              hidden_output_tokens: counts[:thinking].to_i
+            )
+          end
+
+          def gemini_line_items
+            gemini = gemini_usage
+            gemini ? Providers::Gemini::UsageExtractor.line_items(gemini) : []
+          end
+
+          def pricing_mode
+            case @slug
+            when "anthropic", "bedrock"
+              Providers::Anthropic::UsageExtractor.pricing_mode(
+                request: { model: @model, service_tier: @body["serviceTier"].try(:[], "type") },
+                usage: @body["usage"]&.deep_symbolize_keys
+              )
+            when "gemini", "vertexai"
+              Providers::Gemini::Parser.new.pricing_mode(
+                request: {}, usage: @body["usageMetadata"], response_headers: nil, host: host, model: @model
+              )
+            when "openai", "xai", "mistral"
+              service_tier = @body["service_tier"] || @body.dig("usage", "service_tier")
+              Providers::Openai::ResponseParser.combined_pricing_mode(
+                provider: @slug, host: host, model: @model, service_tier: service_tier
+              )
+            else @body["service_tier"]
+            end
+          end
+
+          def host = URI(@provider.api_base).host
+
+          def gemini_usage
+            gemini = @body["usageMetadata"] if @slug == "gemini"
+            gemini if gemini.is_a?(Hash)
+          end
+
+          def input_tokens(tokens, usage)
+            input = usage["inputTokens"] || tokens.input
+            @slug == "anthropic" ? [input, usage["input_tokens"]].compact.max : input
+          end
+
+          def output_tokens(tokens, usage)
+            output = tokens.output
+            thinking = tokens.thinking.to_i
+            input = (usage["input_tokens"] || usage["prompt_tokens"]).to_i
+            output && usage["total_tokens"] == input + output + thinking ? output + thinking : output
+          end
+
+          def cache_writes(total)
+            cache = cache_creation
+            return [total.to_i, 0] unless cache.is_a?(Hash)
+
+            five_minute = cache["ephemeral_5m_input_tokens"].to_i
+            one_hour = cache["ephemeral_1h_input_tokens"].to_i
+            [five_minute + [total.to_i - five_minute - one_hour, 0].max, one_hour]
+          end
+
+          def cache_creation
+            usage = @body["usage"] || {}
+            case @slug
+            when "anthropic" then usage["cache_creation"]
+            when "bedrock"
+              Array(usage["cacheDetails"]).to_h { |d| ["ephemeral_#{d['ttl']}_input_tokens", d["inputTokens"]] }
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/integrations/ruby_llm/v1.rb

@@ -0,0 +1,154 @@
+# frozen_string_literal: true
+
+require_relative "../base"
+require_relative "v1/patches"
+require_relative "v1/reply"
+
+module LlmCostTracker
+  module Integrations
+    module RubyLlm
+      module V1
+        extend Base
+
+        minimum_version "1.15.0"
+        maximum_version "2.0.0"
+
+        class << self
+          def integration_name = :ruby_llm
+
+          def patch_targets
+            [
+              patch_target("RubyLLM::Provider", with: ProviderPatch),
+              patch_target("RubyLLM::Providers::Gemini::Transcription", with: GeminiTranscriptionPatch, optional: true),
+              patch_target("RubyLLM::Streaming", with: StreamPatch, optional: true),
+              *%w[RubyLLM::Providers::OpenAI RubyLLM::Providers::Gemini].map do |name|
+                patch_target(name, with: ResponseBodyPatch, optional: true, skip_when_methods_missing: true)
+              end
+            ]
+          end
+
+          def blocking_seam(resource, record_method, **extras)
+            {
+              provider: resource.slug.to_s,
+              record: lambda do |response, request, latency_ms|
+                public_send(record_method, resource, response, request: request, latency_ms: latency_ms, **extras)
+              end
+            }
+          end
+
+          def request_params(args, kwargs)
+            input = args.first
+            if input.is_a?(Array)
+              input = input.map { |msg| msg.try(:content).then { |content| content.try(:text) || content } || msg }
+            end
+            kwargs.merge(input: input, model: model_id_from_request(kwargs[:model])).with_indifferent_access
+          end
+
+          def provider_response_id_for(response) = Reply.response_id(response)
+
+          def keep_usage(result, body)
+            usage = body.slice("usage", "usageMetadata") if body.is_a?(Hash)
+            result.instance_variable_set(Reply::KEPT_BODY, usage) if usage
+            result
+          end
+
+          def record_completion(provider, response, request:, latency_ms:, has_block:)
+            reply = reply_for(provider, response, request)
+            stream = has_block || request[:stream] == true
+            record_usage(reply, latency_ms, stream: stream, service_line_items: reply.service_line_items)
+          end
+
+          def record_embedding(provider, response, request:, latency_ms:)
+            return unless active?
+
+            record_usage(reply_for(provider, response, request), latency_ms, stream: false, output_tokens: 0)
+          end
+
+          def record_transcription(provider, response, request:, latency_ms:)
+            reply = reply_for(provider, response, request)
+            counts = reply.token_counts
+            no_tokens = counts[:input].to_i.zero? && counts[:output].to_i.zero?
+            duration = billed_duration(reply.usage, response, no_tokens)
+            line_items = Providers::Openai::ServiceCharges.transcription_line_items(duration)
+            record_usage(
+              reply,
+              latency_ms,
+              stream: false,
+              audio_input_tokens: audio_input_tokens(reply, counts),
+              service_line_items: line_items,
+              usage_source: (Usage::Source::UNKNOWN if no_tokens && line_items.empty?)
+            )
+          end
+
+          def record_image(provider, response, request:, latency_ms:)
+            reply = reply_for(provider, response, request)
+            usage = response.try(:usage)
+            usage = (usage.is_a?(Hash) ? usage : {}).with_indifferent_access
+            record_passthrough(
+              provider: reply.slug,
+              model: reply.model,
+              response: response,
+              latency_ms: latency_ms,
+              usage_source: usage.empty? ? Usage::Source::UNKNOWN : Usage::Source::SDK_RESPONSE,
+              **image_tokens(usage, reply.model)
+            )
+          end
+
+          def record_moderation(provider, response, request:, latency_ms:)
+            reply = reply_for(provider, response, request)
+            record_passthrough(
+              provider: reply.slug, model: reply.model, response:, latency_ms:, input_tokens: 0, output_tokens: 0
+            )
+          end
+
+          private
+
+          def reply_for(provider, response, request)
+            Reply.new(provider, response, model_id_from_request(request[:model]))
+          end
+
+          def model_id_from_request(value)
+            return value.to_s if value.is_a?(String) || value.is_a?(Symbol)
+
+            (value.try(:id) || value.try(:model_id) || value.try(:model))&.to_s
+          end
+
+          def record_usage(reply, latency_ms, **options)
+            return unless active?
+
+            record_safely do
+              event = reply.event(**options)
+              LlmCostTracker::Tracker.record(event: event, latency_ms: latency_ms) if event
+            end
+          end
+
+          def billed_duration(usage, response, no_tokens)
+            return usage if usage[:type].to_s == "duration" || usage[:prompt_audio_seconds]
+
+            { type: "duration", seconds: response.duration&.ceil } if no_tokens
+          end
+
+          def audio_input_tokens(reply, counts)
+            tokens = Providers::Openai::UsageExtractor.audio_input_tokens(reply.usage)
+            match = Pricing::Matcher.lookup(provider: reply.slug, model: reply.model) if tokens.zero?
+            match&.prices&.key?("audio_input") ? counts[:input].to_i : tokens
+          end
+
+          def image_tokens(usage, model)
+            extractor = Providers::Openai::UsageExtractor
+            image_input = extractor.image_input_tokens(usage)
+            image_output, text_output = extractor.split_output(
+              output_tokens: usage[:output_tokens].to_i,
+              image_output_details: extractor.image_output_tokens(usage),
+              text_output_details: extractor.text_output_tokens(usage),
+              audio_output: 0,
+              default_to_image: model.to_s.match?(/\A(gpt-image-|gemini-.*-image)/)
+            )
+            { input_tokens: [usage[:input_tokens].to_i - image_input, 0].max, image_input_tokens: image_input,
+              output_tokens: text_output, image_output_tokens: image_output }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/integrations/ruby_llm/v2/attempt/token_counts.rb

@@ -0,0 +1,159 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Integrations
+    module RubyLlm
+      module V2
+        class Attempt
+          class TokenCounts
+            REFUSED = { input_tokens: 0, output_tokens: 0 }.freeze
+
+            def initialize(usage, payload, raw:, request:)
+              @usage = usage
+              @tokens = usage[:tokens]
+              @payload = payload
+              @raw = raw
+              @request = request
+            end
+
+            def event(result:, response_id: nil)
+              return if @usage[:status] != :succeeded && @tokens.to_h == REFUSED
+
+              model = @payload[:response_model] || @usage[:model]
+              line_items = line_items(model, result)
+              source = usage_source(known? || line_items.any?, result)
+              return unless source
+
+              Event.build(
+                provider: @usage[:provider],
+                model: model,
+                token_usage: token_usage(model),
+                pricing_mode: pricing_mode(model),
+                stream: @payload[:streaming],
+                usage_source: source,
+                provider_response_id: result.try(:id) || response_id,
+                service_line_items: line_items
+              )
+            end
+
+            def converse_event(counts)
+              five_minute, one_hour = cache_writes(counts["cacheWriteInputTokens"].to_i, counts["cacheDetails"])
+              Event.build(
+                provider: @usage[:provider],
+                model: @usage[:model],
+                token_usage: Usage::TokenUsage.build(
+                  input_tokens: counts["inputTokens"],
+                  output_tokens: counts["outputTokens"],
+                  cache_read_input_tokens: counts["cacheReadInputTokens"],
+                  cache_write_input_tokens: five_minute,
+                  cache_write_extended_input_tokens: one_hour,
+                  hidden_output_tokens: @tokens.thinking
+                ),
+                pricing_mode: pricing_mode(@usage[:model]),
+                stream: @payload[:streaming],
+                usage_source: Usage::Source::SDK_RESPONSE
+              )
+            end
+
+            def host
+              base = @raw ? @raw.env.url : RubyLLM.config.try("#{@usage[:provider]}_api_base")
+              URI(base.to_s).host if base
+            end
+
+            private
+
+            def known? = @tokens.to_h.any? || !@tokens.reported_cost.nil?
+
+            def usage_source(known, result)
+              return Usage::Source::SDK_RESPONSE if known
+              return Usage::Source::UNKNOWN unless @usage[:status] == :succeeded
+              return if result.nil? || @usage[:operation] == :chat
+
+              @usage[:operation] == :moderation ? Usage::Source::SDK_RESPONSE : Usage::Source::UNKNOWN
+            end
+
+            def token_usage(model)
+              input = @tokens.input.to_i
+              output = @tokens.output.to_i
+              audio = @usage[:operation] == :transcription && audio_priced?(model) ? input : 0
+              image = @usage[:operation] == :image ? output : 0
+              five_minute, one_hour = cache_writes(@tokens.cache_write.to_i, nil)
+              Usage::TokenUsage.build(
+                input_tokens: input - audio,
+                audio_input_tokens: audio,
+                image_input_tokens: @usage[:image_tokens],
+                output_tokens: output - image,
+                image_output_tokens: image,
+                cache_read_input_tokens: @tokens.cache_read.to_i,
+                cache_write_input_tokens: five_minute,
+                cache_write_extended_input_tokens: one_hour,
+                hidden_output_tokens: @tokens.thinking.to_i
+              )
+            end
+
+            def audio_priced?(model)
+              Pricing::Matcher.lookup(provider: @usage[:provider], model: model)&.prices&.key?("audio_input")
+            end
+
+            def cache_writes(total, details)
+              one_hour = if details.is_a?(Array)
+                           details.grep(Hash).sum { |detail| detail["ttl"] == "1h" ? detail["inputTokens"].to_i : 0 }
+                         else
+                           @payload[:caching].try(:[], :ttl).to_s == "1h" ? total : 0
+                         end
+              [[total - one_hour, 0].max, one_hour]
+            end
+
+            def line_items(model, result)
+              return result_line_items(model, result) unless known? && @usage[:operation] != :speech
+
+              service_line_items(model, result)
+            end
+
+            def service_line_items(model, result)
+              cost = @tokens.reported_cost
+              return Providers::Openai::ServiceCharges.billed_line_items(cost: cost) if cost
+
+              calls = Array(result.try(:server_tool_calls)).map(&:raw).grep(Hash)
+              grounding = { "candidates" => calls.map { |call| { "groundingMetadata" => call } } }
+              server_tool_use = @tokens.server_tool_use&.symbolize_keys
+              Providers::Anthropic::UsageExtractor.service_line_items(server_tool_use: server_tool_use) +
+                Providers::Openai::ServiceCharges.line_items_from_output(calls, request: @request, model: model) +
+                Providers::Gemini::Parser.new.service_line_items_for(grounding, model: model)
+            end
+
+            def result_line_items(model, result)
+              return [] unless result
+
+              charges = Providers::Openai::ServiceCharges
+              case @usage[:operation]
+              when :speech then charges.speech_line_items("input" => @payload[:input], "model" => model)
+              when :transcription
+                seconds = result.try(:duration)
+                seconds ? charges.transcription_line_items(type: "duration", seconds: seconds.to_f.ceil) : []
+              when :ocr then charges.ocr_line_items(result.try(:raw).to_h)
+              when :rerank then charges.rerank_line_items(result.try(:raw).to_h)
+              else []
+              end
+            end
+
+            def pricing_mode(model)
+              case @usage[:provider]
+              when "anthropic", "bedrock"
+                Providers::Anthropic::UsageExtractor.pricing_mode(request: @request.merge(model: model), usage: nil)
+              when "gemini"
+                Providers::Gemini::Parser.new.pricing_mode(
+                  request: @request, usage: nil, response_headers: @raw&.headers
+                )
+              else
+                Providers::Openai::ResponseParser.combined_pricing_mode(
+                  provider: @usage[:provider], host: host, model: model, service_tier: @request[:service_tier]
+                )
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/llm_cost_tracker/integrations/ruby_llm/v2/attempt.rb

@@ -0,0 +1,175 @@
+# frozen_string_literal: true
+
+require "json"
+require_relative "attempt/token_counts"
+
+module LlmCostTracker
+  module Integrations
+    module RubyLlm
+      module V2
+        class Attempt
+          attr_reader :usage, :latency_ms
+
+          class << self
+            def stream_window
+              Capture::EventWindow.new(notable: method(:notable_event?),
+                                       trim: stream_parsers.last.method(:trim_stream_event))
+            end
+
+            def batch_event(usage, result, base)
+              base = nil unless URI(base.to_s).host.to_s.match?(/\A(?:us|eu)\.|aiplatform\./i)
+              raw = Faraday::Response.new(status: 200, response_body: result.try(:raw), url: URI(base.to_s))
+              parsed = new(usage, { response: result }, final: true, raw: raw).event
+              return unless parsed
+
+              source = parsed.usage_source == Usage::Source::UNKNOWN ? parsed.usage_source : Usage::Source::SDK_BATCH_RESULT
+              parsed.with(pricing_mode: Pricing::Mode.merge("batch", parsed.pricing_mode), usage_source: source)
+            end
+
+            def request_params(raw, payload)
+              body = raw&.env&.request_body
+              params = body.is_a?(String) ? JSON.parse(body) : body
+              ((params.presence if params.is_a?(Hash)) || payload[:provider_options].to_h).with_indifferent_access
+            rescue JSON::ParserError
+              payload[:provider_options].to_h.with_indifferent_access
+            end
+
+            def stream_parsers
+              @stream_parsers ||= [Providers::Anthropic::Parser.new, Providers::Gemini::Parser.new,
+                                   Providers::Openai::Parser.new]
+            end
+
+            private
+
+            def notable_event?(data) = stream_parsers.any? { |parser| parser.retain_stream_event?(data) }
+          end
+
+          def initialize(usage, payload, final:, latency_ms: nil, events: nil, response: nil, raw: nil, provider: nil)
+            @usage = usage
+            @payload = payload
+            @final = final
+            @latency_ms = latency_ms
+            @events = events
+            @response = response
+            @raw = raw
+            @provider = provider
+          end
+
+          def event
+            base = Faraday::Response.new(url: URI(@provider.api_base.to_s)) if @provider
+            @response.is_a?(Hash) ? transcript_event(base) : response_event(base)
+          end
+
+          private
+
+          def transcript_event(base)
+            counts = TokenCounts.new(@usage, @payload, raw: base, request: Attempt.request_params(base, @payload))
+            if @final && openai_usage?(@response["usage"])
+              event = openai_event(@usage, @response, Attempt.request_params(base, @payload), counts.host)
+            end
+            (event || counts.event(result: (@payload[:result] if @final)))&.with(stream: true)
+          end
+
+          def response_event(base)
+            own = faraday_response(@raw, *(shared_responses if @final))
+            raw = faraday_response(@raw, *shared_responses) || base
+            request = Attempt.request_params(raw, @payload)
+            usage = billed_units(own)
+            converse_event(usage, own) || stream_event(usage, raw, request) || parsed_event(usage, own, request) ||
+              TokenCounts.new(usage, @payload, raw: raw, request: request)
+                         .event(result: (result if @final), response_id: response_id(own))
+          end
+
+          def result = @payload[:response] || @payload[:result]
+
+          def shared_responses = [@response, result.try(:raw)]
+
+          def faraday_response(*candidates) = candidates.find { |candidate| candidate.is_a?(Faraday::Response) }
+
+          def converse_event(usage, own)
+            data = bodies(own).reverse.find { |body| body["usage"].try(:key?, "inputTokens") }
+            return unless data
+
+            request = { service_tier: data["serviceTier"].try(:[], "type") }
+            TokenCounts.new(usage, @payload, raw: nil, request: request).converse_event(data["usage"])
+          end
+
+          def stream_event(usage, raw, request)
+            return unless @events
+
+            context = raw_context(raw, request).merge(response_status: 200, events: @events)
+            Attempt.stream_parsers.lazy.filter_map { |parser| parser.parse_stream(**context) }
+                   .find { |parsed| parsed.usage_source != Usage::Source::UNKNOWN }
+                   &.with(provider: usage[:provider], usage_source: Usage::Source::SDK_RESPONSE)
+          end
+
+          def parsed_event(usage, own, request)
+            body = own&.body
+            return unless body.is_a?(Hash)
+
+            response = raw_context(own, request).merge(response_status: own.status, response_body: body)
+            event = body_event(usage, own, request, response)
+            event&.with(provider: usage[:provider], usage_source: Usage::Source::SDK_RESPONSE)
+          end
+
+          def body_event(usage, own, request, response)
+            body = response[:response_body]
+            if body["type"] == "message" then Providers::Anthropic::Parser.new.parse(**response)
+            elsif body.key?("usageMetadata") || body["object"] == "interaction" then gemini_event(response)
+            elsif openai_usage?(body["usage"]) then openai_event(usage, body, request, own.env.url.host)
+            end
+          end
+
+          def raw_context(raw, request)
+            { request_url: raw.env.url.to_s, request_body: request, response_headers: raw.headers }
+          end
+
+          def gemini_event(response)
+            event = Providers::Gemini::Parser.new.parse(**response)
+            body = response[:response_body]
+            details = body.dig("usageMetadata", body.key?("embedding") ? "promptTokensDetails" : "promptTokenDetails")
+            video = Providers::Gemini::UsageExtractor.modality_tokens(details, "VIDEO")
+            return event unless event && video.positive?
+
+            event.with(line_items: [*event.line_items,
+                                    Charges::LineItem.build(dimension_key: "video_input", quantity: video)])
+          end
+
+          def openai_event(usage, body, request, host)
+            request[:model] ||= usage[:model]
+            Providers::Openai::ResponseParser.event_from_response(
+              response: body,
+              request: request,
+              provider: usage[:provider],
+              host: host,
+              usage_source: Usage::Source::SDK_RESPONSE
+            )
+          end
+
+          def openai_usage?(usage)
+            usage.is_a?(Hash) &&
+              (usage.key?("input_tokens") || usage.key?("prompt_tokens") || usage.key?("cost_in_usd_ticks") ||
+               usage["type"] == "duration")
+          end
+
+          def bodies(own) = [own&.body, *@events&.map { |event| event[:data] }].grep(Hash)
+
+          def response_id(own) = bodies(own).filter_map { |body| body["id"] || body["responseId"] }.first
+
+          def billed_units(own)
+            units = bodies(own).filter_map do |data|
+              (data["type"] == "message-end" ? data.dig("delta", "usage") : data["usage"] || data["meta"])
+                .try(:[], "billed_units")
+            end.last
+            return @usage unless units
+
+            tokens = @usage[:tokens]
+            @usage.merge(tokens: RubyLLM::Tokens.new(input: units["input_tokens"] || tokens.input,
+                                                     output: units["output_tokens"] || tokens.output),
+                         image_tokens: units["image_tokens"])
+          end
+        end
+      end
+    end
+  end
+end