Module: Apipurse::WrapOpenAI

Defined in:
lib/apipurse/wrap_openai.rb

Class Method Summary collapse

Class Method Details

.call(client, apipurse, on_track: nil) ⇒ Object 



30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
 
# File 'lib/apipurse/wrap_openai.rb', line 30

def self.call(client, apipurse, on_track: nil)
  completions = client.chat.completions
  original = completions.method(:create)

  completions.define_singleton_method(:create) do |*args, **kwargs, &block|
    result = original.call(*args, **kwargs, &block)

    if result.respond_to?(:usage)
      Apipurse::WrapOpenAI.report(result, apipurse, on_track)
      result
    elsif result.respond_to?(:each)
      Apipurse::WrapOpenAI.wrap_stream(result, apipurse, on_track)
    else
      result
    end
  end

  client
end

.report(completion, apipurse, on_track) ⇒ Object 



50
51
52
53
54
55
56
57
58
59
60
61
62
 
# File 'lib/apipurse/wrap_openai.rb', line 50

def self.report(completion, apipurse, on_track)
  usage = completion.respond_to?(:usage) ? completion.usage : nil
  return unless usage

  result = apipurse.track_completion(
    model: completion.model,
    input_tokens: (usage.respond_to?(:prompt_tokens) ? usage.prompt_tokens : nil) || 0,
    output_tokens: (usage.respond_to?(:completion_tokens) ? usage.completion_tokens : nil) || 0
  )
  on_track&.call(result)
rescue StandardError
  nil
end

.wrap_stream(stream, apipurse, on_track) ⇒ Object 



64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
 
# File 'lib/apipurse/wrap_openai.rb', line 64

def self.wrap_stream(stream, apipurse, on_track)
  Enumerator.new do |yielder|
    last_model = nil
    last_usage = nil

    stream.each do |chunk|
      last_model = chunk.model if chunk.respond_to?(:model) && chunk.model
      last_usage = chunk.usage if chunk.respond_to?(:usage) && chunk.usage
      yielder << chunk
    end

    next unless last_model && last_usage

    begin
      result = apipurse.track_completion(
        model: last_model,
        input_tokens: (last_usage.respond_to?(:prompt_tokens) ? last_usage.prompt_tokens : nil) || 0,
        output_tokens: (last_usage.respond_to?(:completion_tokens) ? last_usage.completion_tokens : nil) || 0
      )
      on_track&.call(result)
    rescue StandardError
      nil
    end
  end
end