diff --git a/enterprise/app/helpers/captain/chat_generation_recorder.rb b/enterprise/app/helpers/captain/chat_generation_recorder.rb index cd631fb16..5a944b4b5 100644 --- a/enterprise/app/helpers/captain/chat_generation_recorder.rb +++ b/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 diff --git a/lib/integrations/llm_base_service.rb b/lib/integrations/llm_base_service.rb index 397888b83..abdd90da3 100644 --- a/lib/integrations/llm_base_service.rb +++ b/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 diff --git a/lib/integrations/llm_instrumentation.rb b/lib/integrations/llm_instrumentation.rb index 326bb901e..74d56f237 100644 --- a/lib/integrations/llm_instrumentation.rb +++ b/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 diff --git a/lib/integrations/llm_instrumentation_completion_helpers.rb b/lib/integrations/llm_instrumentation_completion_helpers.rb index 50e071119..a9b03c0f7 100644 --- a/lib/integrations/llm_instrumentation_completion_helpers.rb +++ b/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 diff --git a/lib/integrations/llm_instrumentation_constants.rb b/lib/integrations/llm_instrumentation_constants.rb index dfe1e7704..54579779c 100644 --- a/lib/integrations/llm_instrumentation_constants.rb +++ b/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 diff --git a/lib/integrations/llm_instrumentation_helpers.rb b/lib/integrations/llm_instrumentation_helpers.rb index 129092ed4..6bc4b79ad 100644 --- a/lib/integrations/llm_instrumentation_helpers.rb +++ b/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 diff --git a/lib/integrations/llm_instrumentation_spans.rb b/lib/integrations/llm_instrumentation_spans.rb index 85ea599f8..90ad0b845 100644 --- a/lib/integrations/llm_instrumentation_spans.rb +++ b/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 diff --git a/lib/integrations/llm_usage_details_builder.rb b/lib/integrations/llm_usage_details_builder.rb new file mode 100644 index 000000000..809cf10c8 --- /dev/null +++ b/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 diff --git a/spec/lib/integrations/llm_instrumentation_spec.rb b/spec/lib/integrations/llm_instrumentation_spec.rb index 0be62f437..526fdf965 100644 --- a/spec/lib/integrations/llm_instrumentation_spec.rb +++ b/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