feat(captain): add cache-aware llm usage instrumentation

enterprise/app/helpers/captain/chat_generation_recorder.rb

@@ -1,6 +1,7 @@
 module Captain::ChatGenerationRecorder
   extend ActiveSupport::Concern
   include Integrations::LlmInstrumentationConstants
+  include Integrations::LlmUsageDetailsBuilder
 
   private
 
@@ -30,12 +31,16 @@ module Captain::ChatGenerationRecorder
   end
 
   def generation_attributes(chat, message)
+    usage_details = generation_usage_details(message)
+
     {
       ATTR_GEN_AI_PROVIDER => determine_provider(model),
       ATTR_GEN_AI_REQUEST_MODEL => model,
       ATTR_GEN_AI_REQUEST_TEMPERATURE => temperature,
-      ATTR_GEN_AI_USAGE_INPUT_TOKENS => message.input_tokens,
+      ATTR_GEN_AI_USAGE_INPUT_TOKENS => usage_details[:input],
       ATTR_GEN_AI_USAGE_OUTPUT_TOKENS => message.respond_to?(:output_tokens) ? message.output_tokens : nil,
+      ATTR_GEN_AI_USAGE_TOTAL_TOKENS => usage_details[:total],
+      ATTR_LANGFUSE_OBSERVATION_USAGE_DETAILS => usage_details.to_json.presence,
       ATTR_LANGFUSE_OBSERVATION_INPUT => format_input_messages(chat),
       ATTR_LANGFUSE_OBSERVATION_OUTPUT => message.respond_to?(:content) ? message.content.to_s : nil
     }
@@ -44,4 +49,8 @@ module Captain::ChatGenerationRecorder
   def format_input_messages(chat)
     chat.messages[0...-1].map { |m| { role: m.role.to_s, content: m.content.to_s } }.to_json
   end
+
+  def generation_usage_details(message)
+    usage_details_from_message(message, provider: determine_provider(model))
+  end
 end

lib/integrations/llm_base_service.rb

@@ -146,6 +146,8 @@ class Integrations::LlmBaseService
       usage: {
         'prompt_tokens' => response.input_tokens,
         'completion_tokens' => response.output_tokens,
+        'cached_tokens' => response.respond_to?(:cached_tokens) ? response.cached_tokens : nil,
+        'cache_creation_tokens' => response.respond_to?(:cache_creation_tokens) ? response.cache_creation_tokens : nil,
         'total_tokens' => (response.input_tokens || 0) + (response.output_tokens || 0)
       },
       request_messages: messages

lib/integrations/llm_instrumentation.rb

@@ -16,7 +16,7 @@ module Integrations::LlmInstrumentation
       setup_span_attributes(span, params)
       result = yield
       executed = true
-      record_completion(span, result)
+      record_completion(span, result, params)
       result
     end
   rescue StandardError => e

lib/integrations/llm_instrumentation_completion_helpers.rb

@@ -2,6 +2,7 @@
 
 module Integrations::LlmInstrumentationCompletionHelpers
   include Integrations::LlmInstrumentationConstants
+  include Integrations::LlmUsageDetailsBuilder
 
   private
 
@@ -73,9 +74,13 @@ module Integrations::LlmInstrumentationCompletionHelpers
     usage = result[:usage] || result['usage']
     return if usage.blank?
 
-    span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, usage['prompt_tokens']) if usage['prompt_tokens']
-    span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, usage['completion_tokens']) if usage['completion_tokens']
-    span.set_attribute(ATTR_GEN_AI_USAGE_TOTAL_TOKENS, usage['total_tokens']) if usage['total_tokens']
+    usage_details = usage_details_from_hash(usage)
+    return if usage_details.blank?
+
+    span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, usage_details[:input]) if usage_details[:input]
+    span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, usage_details[:output]) if usage_details[:output]
+    span.set_attribute(ATTR_GEN_AI_USAGE_TOTAL_TOKENS, usage_details[:total]) if usage_details[:total]
+    span.set_attribute(ATTR_LANGFUSE_OBSERVATION_USAGE_DETAILS, usage_details.to_json)
   end
 
   def set_error_attributes(span, result)
@@ -85,4 +90,14 @@ module Integrations::LlmInstrumentationCompletionHelpers
     span.set_attribute(ATTR_GEN_AI_RESPONSE_ERROR, error.to_json)
     span.status = OpenTelemetry::Trace::Status.error(error.to_s.truncate(1000))
   end
+
+  def set_message_usage_metrics(span, message, provider: nil)
+    usage_details = usage_details_from_message(message, provider: provider)
+    return if usage_details.blank?
+
+    span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, usage_details[:input]) if usage_details[:input]
+    span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, usage_details[:output]) if usage_details[:output]
+    span.set_attribute(ATTR_GEN_AI_USAGE_TOTAL_TOKENS, usage_details[:total]) if usage_details[:total]
+    span.set_attribute(ATTR_LANGFUSE_OBSERVATION_USAGE_DETAILS, usage_details.to_json)
+  end
 end

lib/integrations/llm_instrumentation_constants.rb

@@ -29,4 +29,5 @@ module Integrations::LlmInstrumentationConstants
   ATTR_LANGFUSE_OBSERVATION_TYPE = 'langfuse.observation.type'
   ATTR_LANGFUSE_OBSERVATION_INPUT = 'langfuse.observation.input'
   ATTR_LANGFUSE_OBSERVATION_OUTPUT = 'langfuse.observation.output'
+  ATTR_LANGFUSE_OBSERVATION_USAGE_DETAILS = 'langfuse.observation.usage_details'
 end

lib/integrations/llm_instrumentation_helpers.rb

@@ -24,10 +24,11 @@ module Integrations::LlmInstrumentationHelpers
     set_metadata_attributes(span, params)
   end
 
-  def record_completion(span, result)
+  def record_completion(span, result, params)
     if result.respond_to?(:content)
       span.set_attribute(ATTR_GEN_AI_COMPLETION_ROLE, result.role.to_s) if result.respond_to?(:role)
       span.set_attribute(ATTR_GEN_AI_COMPLETION_CONTENT, result.content.to_s)
+      set_message_usage_metrics(span, result, provider: determine_provider(params[:model]))
     elsif result.is_a?(Hash)
       set_completion_attributes(span, result)
     end

lib/integrations/llm_instrumentation_spans.rb

@@ -4,6 +4,7 @@ require 'opentelemetry_config'
 
 module Integrations::LlmInstrumentationSpans
   include Integrations::LlmInstrumentationConstants
+  include Integrations::LlmUsageDetailsBuilder
 
   def tracer
     @tracer ||= OpentelemetryConfig.tracer
@@ -86,7 +87,29 @@ module Integrations::LlmInstrumentationSpans
   end
 
   def set_llm_turn_usage_attributes(span, message)
-    span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, message.input_tokens) if message.respond_to?(:input_tokens) && message.input_tokens
-    span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, message.output_tokens) if message.respond_to?(:output_tokens) && message.output_tokens
+    usage_details = llm_turn_usage_details(message)
+    return if usage_details.blank?
+
+    span.set_attribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, usage_details[:input]) if usage_details[:input]
+    span.set_attribute(ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, usage_details[:output]) if usage_details[:output]
+    span.set_attribute(ATTR_GEN_AI_USAGE_TOTAL_TOKENS, usage_details[:total]) if usage_details[:total]
+    span.set_attribute(ATTR_LANGFUSE_OBSERVATION_USAGE_DETAILS, usage_details.to_json)
+  end
+
+  def llm_turn_usage_details(message)
+    usage_details_from_message(message, provider: llm_turn_provider(message))
+  end
+
+  def llm_turn_provider(message)
+    model_name = message.respond_to?(:model_id) ? message.model_id : nil
+    return 'openai' if model_name.blank?
+
+    model = model_name.to_s.downcase
+
+    LlmConstants::PROVIDER_PREFIXES.each do |provider, prefixes|
+      return provider if prefixes.any? { |prefix| model.start_with?(prefix) }
+    end
+
+    'openai'
   end
 end

lib/integrations/llm_usage_details_builder.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+module Integrations::LlmUsageDetailsBuilder
+  private
+
+  def usage_details_from_hash(usage)
+    usage_hash = normalize_usage_hash(usage)
+    return {} if usage_hash.blank?
+
+    usage_details_from_hash_values(usage_hash)
+  end
+
+  def usage_details_from_hash_values(usage_hash)
+    input_tokens_source = usage_hash.key?('input_tokens') ? 'input_tokens' : 'prompt_tokens'
+    cache_read_tokens = usage_cache_read_tokens(usage_hash)
+    cache_creation_tokens = usage_cache_creation_tokens(usage_hash)
+    input_tokens = normalized_uncached_input_tokens(
+      usage_hash[input_tokens_source],
+      cache_read_tokens,
+      input_tokens_source: input_tokens_source
+    )
+    output_tokens = usage_hash['output_tokens'] || usage_hash['completion_tokens']
+    total_tokens = usage_total_tokens(
+      usage_hash['total_tokens'],
+      input_tokens,
+      output_tokens,
+      cache_read_tokens,
+      cache_creation_tokens
+    )
+
+    build_usage_details(input_tokens, output_tokens, total_tokens, cache_read_tokens, cache_creation_tokens)
+  end
+
+  def usage_details_from_message(message, provider:)
+    input_tokens = message_token(message, :input_tokens)
+    output_tokens = message_token(message, :output_tokens)
+    cache_read_tokens = message_token(message, :cached_tokens)
+    cache_creation_tokens = message_token(message, :cache_creation_tokens)
+
+    input_tokens = normalized_uncached_input_tokens(
+      input_tokens,
+      cache_read_tokens,
+      provider: provider,
+      input_tokens_source: 'input_tokens'
+    )
+    total_tokens = total_tokens_from_parts(input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens)
+
+    build_usage_details(input_tokens, output_tokens, total_tokens, cache_read_tokens, cache_creation_tokens)
+  end
+
+  def build_usage_details(input_tokens, output_tokens, total_tokens, cache_read_tokens, cache_creation_tokens)
+    compact_usage_details(
+      input: input_tokens,
+      output: output_tokens,
+      total: total_tokens,
+      cache_read_input_tokens: cache_read_tokens,
+      cache_creation_input_tokens: cache_creation_tokens
+    )
+  end
+
+  def normalize_usage_hash(usage)
+    return usage.deep_stringify_keys if usage.respond_to?(:deep_stringify_keys)
+    return usage.to_h.transform_keys(&:to_s) if usage.respond_to?(:to_h)
+
+    {}
+  end
+
+  def usage_cache_read_tokens(usage_hash)
+    usage_hash['cache_read_input_tokens'] ||
+      usage_hash['cached_tokens'] ||
+      usage_hash.dig('prompt_tokens_details', 'cached_tokens')
+  end
+
+  def usage_cache_creation_tokens(usage_hash)
+    usage_hash['cache_creation_input_tokens'] ||
+      usage_hash['cache_creation_tokens'] ||
+      usage_hash['cache_creation']&.values&.compact&.sum
+  end
+
+  def usage_total_tokens(reported_total, input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens)
+    if cache_read_tokens || cache_creation_tokens
+      total_tokens_from_parts(input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens)
+    else
+      reported_total || total_tokens_from_parts(input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens)
+    end
+  end
+
+  def message_token(message, token_name)
+    message.respond_to?(token_name) ? message.public_send(token_name) : nil
+  end
+
+  def normalized_uncached_input_tokens(input_tokens, cache_read_tokens, provider: nil, input_tokens_source: nil)
+    return input_tokens if input_tokens.nil? || cache_read_tokens.nil?
+
+    if input_tokens_source == 'prompt_tokens' || provider == 'openai'
+      [input_tokens - cache_read_tokens, 0].max
+    else
+      input_tokens
+    end
+  end
+
+  def total_tokens_from_parts(input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens)
+    values = [input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens].compact
+    values.sum if values.present?
+  end
+
+  def compact_usage_details(details)
+    details.compact
+  end
+end

spec/lib/integrations/llm_instrumentation_spec.rb

@@ -176,6 +176,7 @@ RSpec.describe Integrations::LlmInstrumentation do
           {
             usage: {
               'prompt_tokens' => 150,
+              'prompt_tokens_details' => { 'cached_tokens' => 20 },
               'completion_tokens' => 200,
               'total_tokens' => 350
             }
@@ -183,9 +184,41 @@ RSpec.describe Integrations::LlmInstrumentation do
         end
 
         expect(result[:usage]['prompt_tokens']).to eq(150)
-        expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.input_tokens', 150)
+        expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.input_tokens', 130)
         expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.output_tokens', 200)
         expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.total_tokens', 350)
+        expect(mock_span).to have_received(:set_attribute)
+          .with('langfuse.observation.usage_details', '{"input":130,"output":200,"total":350,"cache_read_input_tokens":20}')
+      end
+
+      it 'sets usage metrics for RubyLLM message responses with cached tokens' do
+        mock_span = instance_double(OpenTelemetry::Trace::Span)
+        allow(mock_span).to receive(:set_attribute)
+        allow(mock_span).to receive(:status=)
+        mock_tracer = instance_double(OpenTelemetry::Trace::Tracer)
+        allow(instance).to receive(:tracer).and_return(mock_tracer)
+        allow(mock_tracer).to receive(:in_span).and_yield(mock_span)
+
+        llm_message = instance_double(
+          RubyLLM::Message,
+          role: :assistant,
+          content: 'AI response',
+          input_tokens: 150,
+          output_tokens: 50,
+          cached_tokens: 40,
+          cache_creation_tokens: nil
+        )
+
+        result = instance.instrument_llm_call(params) do
+          llm_message
+        end
+
+        expect(result).to eq(llm_message)
+        expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.input_tokens', 110)
+        expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.output_tokens', 50)
+        expect(mock_span).to have_received(:set_attribute).with('gen_ai.usage.total_tokens', 200)
+        expect(mock_span).to have_received(:set_attribute)
+          .with('langfuse.observation.usage_details', '{"input":110,"output":50,"total":200,"cache_read_input_tokens":40}')
       end
 
       it 'sets error attributes when result contains error' do