llm_cost_tracker 0.14.2 → 0.15.0

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

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Integrations
+    module RubyLlm
+      module V2
+        class Frame
+          STACK = :llm_cost_tracker_ruby_llm_frames
+
+          attr_reader :payload, :provider, :workflow
+
+          class << self
+            def current = stack.last
+
+            def open(payload, compaction: false, workflow: nil)
+              new(payload, compaction, workflow).tap { |frame| stack << frame }
+            end
+
+            def close(payload)
+              index = stack.rindex { |frame| frame.payload.equal?(payload) }
+              index ? stack.slice!(index..) : []
+            end
+
+            def discard(frame) = stack.delete_if { |open| open.equal?(frame) }
+
+            private
+
+            def stack = Thread.current[STACK] ||= []
+          end
+
+          def initialize(payload, compaction, workflow)
+            @payload = payload
+            @compaction = compaction
+            @workflow = workflow
+            @reported = []
+          end
+
+          def observe(seam, value, provider)
+            @provider = provider
+            case seam
+            when :parse_completion_body then @raw = value unless value.is_a?(Hash)
+            when :build_chunk then (@window ||= Attempt.stream_window).push(value)
+            else @response = value
+            end
+          end
+
+          def request_started
+            @request_started_at = Timing.now_monotonic
+            @latency_ms = nil
+          end
+
+          def request_finished(payload)
+            @latency_ms = Timing.elapsed_ms(@request_started_at) if @request_started_at
+            @workflow = payload.slice(:workflow_id, :workflow_name, :workflow_step_name)
+          end
+
+          def add_attempt(usage)
+            raw = @raw unless usage[:status] == :succeeded && usage[:tokens].to_h.empty?
+            @reported << [usage, { latency_ms: @latency_ms, events: @window&.events, raw: raw }]
+            @window = nil
+            @raw = nil if raw
+          end
+
+          def attempts
+            final = @reported.rindex { |usage, _| usage[:status] == :succeeded } || (@reported.size - 1 if @response)
+            payload = @compaction ? @payload.except(:provider_options) : @payload
+            @reported.each_with_index.map do |(usage, reported), index|
+              Attempt.new(usage, payload, final: index == final, response: @response, provider: @provider, **reported)
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Integrations
+    module RubyLlm
+      module V2
+        module Recorder
+          RECORDED = ObjectSpace::WeakKeyMap.new
+
+          class << self
+            def record(attempt)
+              event = attempt.event
+              return unless event
+
+              LlmCostTracker::Tracker.record(event: event, latency_ms: attempt.latency_ms, **tags_for(attempt.usage))
+            end
+
+            def record_all(attempts) = attempts.filter_map { |attempt| caller_error { record(attempt) } }
+
+            def record_batch(batch, results, frame)
+              errors = Array(results).each_with_index.filter_map do |result, index|
+                next if result.nil? || RECORDED.key?(result)
+
+                caller_error do
+                  record_batch_result(batch, result, index, frame)
+                  RECORDED[result] = true
+                end
+              end
+              raise errors.first if errors.any?
+            end
+
+            def tags_for(payload)
+              context = LlmCostTracker::Tags::Context.tags
+              run_tagged = context.any? { |key, value| key.to_s == "run_id" && !value.to_s.empty? }
+              run_id = payload[:workflow_id] unless run_tagged
+              { context_tags: context,
+                metadata: { run_id: run_id }.compact.merge(payload.slice(:workflow_name, :workflow_step_name).compact) }
+            end
+
+            private
+
+            def caller_error(&)
+              V2.record_safely(&)
+              nil
+            rescue *CALLER_ERRORS => e
+              e
+            end
+
+            def record_batch_result(batch, result, index, frame)
+              usage = batch_usage(batch, result, index)
+              base = frame.provider&.api_base || RubyLLM.config.try("#{batch.provider}_api_base")
+              event = Attempt.batch_event(usage, result, base)
+              return unless event
+
+              id = event.provider_response_id || "#{batch.id}/#{index}"
+              return if Call.already_recorded?(provider: event.provider, provider_response_id: id)
+
+              V2.record_once(event.with(provider_response_id: id), **tags_for(frame.workflow.to_h))
+            end
+
+            def batch_usage(batch, result, index)
+              { operation: result.is_a?(RubyLLM::Embedding) ? :embedding : :chat, provider: batch.provider,
+                model: result.model || batch.chats.to_a[index]&.model&.id, status: :succeeded, tokens: result.tokens }
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+module LlmCostTracker
+  module Integrations
+    module RubyLlm
+      module V2
+        SEAMS = {
+          build_chunk: %w[Protocols::Anthropic Protocols::ChatCompletions Protocols::Responses Protocols::Gemini
+                          Protocols::Interactions Protocols::Converse Providers::OpenRouter::ChatCompletions
+                          Protocols::Cohere],
+          parse_completion_body: %w[Protocols::Anthropic Protocols::ChatCompletions Protocols::Responses
+                                    Protocols::Gemini Protocols::Interactions Protocols::Converse
+                                    Protocols::Mistral::Conversations Providers::Mistral::ChatCompletions],
+          parse_embedding_response: %w[Protocols::ChatCompletions Protocols::Gemini Providers::VertexAI::EmbedContent
+                                       Providers::Perplexity::Embeddings Protocols::Cohere],
+          parse_transcription_response: %w[Protocols::ChatCompletions Protocols::Gemini],
+          parse_speech_response: %w[Protocols::Gemini],
+          stream_transcription: %w[Protocols::ChatCompletions],
+          transcribe: %w[Protocol Protocols::Deepgram Protocols::Gemini::LiveTranscription],
+          parse_image_responses: %w[Protocols::ChatCompletions Protocols::Gemini Providers::XAI::Images],
+          parse_cache_response: %w[Protocols::Gemini],
+          messages: %w[Batch],
+          results: %w[Batch],
+          current_workflow: %w[Support::Instrumentation]
+        }.freeze
+        StreamTranscriptionBridge = Module.new do
+          def stream_transcription(*, **, &block)
+            super do |chunk|
+              V2.observe(:stream_transcription, chunk.raw, self)
+              block.call(chunk)
+            end
+          end
+        end
+        TranscribeBridge = Module.new do
+          def transcribe(*, **, &block)
+            V2.observe(:transcribe, {}, self) if block
+            super
+          end
+        end
+        BatchBridge = Module.new do
+          %i[messages results].each { |name| define_method(name) { V2.collect(self) { super() } } }
+        end
+        BRIDGES = (SEAMS.keys - %i[stream_transcription transcribe messages results current_workflow]).to_h do |seam|
+          bridge = Module.new do
+            define_method(seam) do |value, *args, **options, &block|
+              V2.observe(seam, options.fetch(:raw, value), self)
+              super(value, *args, **options, &block)
+            end
+          end
+          [seam, const_set("#{seam.to_s.camelize}Bridge", bridge)]
+        end.merge(
+          stream_transcription: StreamTranscriptionBridge,
+          transcribe: TranscribeBridge,
+          messages: BatchBridge,
+          results: BatchBridge
+        ).freeze
+
+        module Seams
+          class << self
+            def bridge
+              BRIDGES.each do |seam, bridge|
+                SEAMS[seam].filter_map { |target| owner(target, seam) }.each do |owner|
+                  owner.prepend(bridge) unless owner == bridge
+                end
+              end
+            end
+
+            def missing
+              SEAMS.flat_map do |seam, targets|
+                targets.reject { |target| owner(target, seam) }.map { |target| "RubyLLM::#{target}##{seam}" }
+              end
+            end
+
+            private
+
+            def owner(target, seam)
+              "RubyLLM::#{target}".safe_constantize&.instance_method(seam)&.owner
+            rescue NameError
+              nil
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,161 @@
+# frozen_string_literal: true
+
+require "active_support/notifications"
+require_relative "../base"
+require_relative "v2/attempt"
+require_relative "v2/frame"
+require_relative "v2/recorder"
+require_relative "v2/seams"
+
+module LlmCostTracker
+  module Integrations
+    module RubyLlm
+      module V2
+        extend Base
+
+        minimum_version "2.0.0"
+        maximum_version "3.0.0"
+
+        OPERATIONS = %i[chat embedding image transcription moderation speech ocr rerank judgment].freeze
+        JOBS = %w[batch.ruby_llm video_job.ruby_llm research_job.ruby_llm].freeze
+        EVENTS = [*OPERATIONS, :compaction, :request, :usage].map { |name| "#{name}.ruby_llm" }.concat(JOBS).freeze
+        CACHE_CREATED = :llm_cost_tracker_ruby_llm_cache_created
+
+        module TagRelay
+          module_function
+
+          def instrument(*) = yield
+
+          def capture_context = LlmCostTracker::Tags::Context.scoped
+
+          def with_context(tags, &) = LlmCostTracker::Tags::Context.with(tags, &)
+        end
+
+        class << self
+          def integration_name = :ruby_llm
+
+          def install
+            validate_contract!
+            Logging.warn(untested_version_message) if untested_version?
+            @subscriptions ||= EVENTS.map { |name| ActiveSupport::Notifications.subscribe(name, self) }
+            RubyLLM.config.instrumenter ||= ActiveSupport::Notifications
+            Seams.bridge
+            instrumentation = "RubyLLM::Support::Instrumentation".safe_constantize
+            instrumentation.subscribe(TagRelay) if instrumentation.respond_to?(:capture_context)
+          end
+
+          def status
+            name = integration_name.to_s
+            warning = install_warning
+            return Check.new(:warn, name, warning) if warning
+
+            missing = Seams.missing
+            return Check.new(:ok, name, "#{name} integration installed") if missing.empty?
+
+            message = "#{name} integration installed, but these RubyLLM methods are missing, so what they carry " \
+                      "is not read: #{missing.join(', ')}"
+            Check.new(:warn, name, message)
+          end
+
+          def start(name, _id, payload)
+            case name
+            when "usage.ruby_llm" then nil
+            when "request.ruby_llm" then start_request(payload)
+            else start_operation(name, payload)
+            end
+          end
+
+          def finish(name, _id, payload)
+            case name
+            when "usage.ruby_llm" then add_attempt(payload) unless %i[video research].include?(payload[:operation])
+            when "request.ruby_llm" then Frame.current&.request_finished(payload)
+            else finish_operation(payload)
+            end
+          end
+
+          def observe(seam, value, protocol = nil)
+            record_safely do
+              next record_cache_storage(value) if seam == :parse_cache_response
+
+              Frame.current&.observe(seam, value, protocol&.provider)
+            end
+          end
+
+          def collect(batch)
+            frame = Frame.open(batch, workflow: current_workflow)
+            results = yield
+            Recorder.record_batch(batch, results, frame) if active?
+            results
+          ensure
+            Frame.discard(frame)
+          end
+
+          private
+
+          def install_warning
+            problems = version_problems + target_problems
+            return "#{integration_name} integration cannot be installed: #{problems.join('; ')}" if problems.any?
+            return untested_version_message if untested_version?
+            return "#{integration_name} integration is enabled but not installed" unless @subscriptions
+
+            instrumenter = RubyLLM.config.instrumenter
+            return if instrumenter == ActiveSupport::Notifications
+
+            "RubyLLM.config.instrumenter is #{instrumenter.inspect}, not ActiveSupport::Notifications, " \
+              "so RubyLLM calls are not recorded"
+          end
+
+          def target_problems = defined?(RubyLLM.config) ? [] : ["RubyLLM is not loaded"]
+
+          def current_workflow
+            instrumentation = "RubyLLM::Support::Instrumentation".safe_constantize
+            instrumentation.current_workflow if instrumentation.respond_to?(:current_workflow)
+          end
+
+          def start_operation(name, payload)
+            return unless active?
+
+            enforce_budget!(
+              request: budget_request(payload),
+              provider: payload[:provider].to_s,
+              tags: (LlmCostTracker::Tracker.build_tags(**Recorder.tags_for(payload)) if payload[:workflow_id])
+            )
+            Frame.open(payload, compaction: name.start_with?("compaction")) if JOBS.exclude?(name)
+          end
+
+          def budget_request(payload)
+            input = payload[:input_messages]&.map { |message| message.try(:content).then { |c| c.try(:text) || c } }
+            { model: payload[:model], input: input || payload.values_at(:input, :prompt, :query) }
+          end
+
+          def finish_operation(payload)
+            errors = Recorder.record_all(Frame.close(payload).flat_map(&:attempts))
+            raise errors.first if errors.any? && !payload[:exception]
+          end
+
+          def start_request(payload)
+            cache = payload[:method] == :post && payload[:url].to_s.end_with?("cachedContents")
+            Thread.current[CACHE_CREATED] = (payload[:provider].to_s if cache)
+            enforce_budget!(request: {}, provider: payload[:provider].to_s) if cache
+            Frame.current&.request_started
+          end
+
+          def add_attempt(usage)
+            if OPERATIONS.include?(usage[:operation])
+              Frame.current&.add_attempt(usage)
+            elsif active?
+              record_safely { Recorder.record(Attempt.new(usage, {}, final: false)) }
+            end
+          end
+
+          def record_cache_storage(data)
+            event = Providers::Gemini::Parser.new.cache_storage_event(data) if Thread.current[CACHE_CREATED] && active?
+            LlmCostTracker::Tracker.record(event: event.with(provider: Thread.current[CACHE_CREATED])) if event
+          rescue LlmCostTracker::BudgetExceededError
+            nil
+          end
+        end
+      end
+    end
+  end
+end