Module: DSPy::Instrumentation::TokenTracker

Extended by:
TokenTracker
Included in:
TokenTracker
Defined in:
lib/dspy/instrumentation/token_tracker.rb

Overview

Utility for extracting token usage from different LM adapters Uses actual token counts from API responses for accuracy

Instance Method Summary collapse

Instance Method Details

#extract_token_usage(response, provider) ⇒ Object

Extract actual token usage from API responses



13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
 
# File 'lib/dspy/instrumentation/token_tracker.rb', line 13

def extract_token_usage(response, provider)
  return {} unless response&.usage
  
  # Handle Usage struct
  if response.usage.is_a?(DSPy::LM::Usage) || response.usage.is_a?(DSPy::LM::OpenAIUsage)
    return {
      input_tokens: response.usage.input_tokens,
      output_tokens: response.usage.output_tokens,
      total_tokens: response.usage.total_tokens
    }
  end
  
  # Fallback to legacy hash handling
  case provider.to_s.downcase
  when 'openai'
    extract_openai_tokens(response)
  when 'anthropic'
    extract_anthropic_tokens(response)
  else
    {} # No token information for other providers
  end
end