Class: LlmCostTracker::Capture::StreamCollector

Inherits:
Object
  • Object
show all
Defined in:
lib/llm_cost_tracker/capture/stream_collector.rb

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(provider:, model:, latency_ms: nil, provider_response_id: nil, provider_project_id: nil, provider_api_key_id: nil, provider_workspace_id: nil, batch: nil, pricing_mode: nil, metadata: {}, context_tags: nil, request: nil) ⇒ StreamCollector

Returns a new instance of StreamCollector.



16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
 
# File 'lib/llm_cost_tracker/capture/stream_collector.rb', line 16

def initialize(provider:, model:, latency_ms: nil, provider_response_id: nil, provider_project_id: nil,
               provider_api_key_id: nil, provider_workspace_id: nil, batch: nil, pricing_mode: nil,
               metadata: {}, context_tags: nil, request: nil)
  @provider = provider.to_s
  @model = model
  @latency_ms = latency_ms
  @provider_response_id = provider_response_id
  @provider_project_id = provider_project_id
  @provider_api_key_id = provider_api_key_id
  @provider_workspace_id = provider_workspace_id
  @batch = batch
  @pricing_mode = pricing_mode
  @metadata = (metadata || {}).deep_dup
  @context_tags = (context_tags || LlmCostTracker::Tags::Context.tags).deep_dup
  @request = request
  @events = []
  @captured_bytes = 0
  @overflowed = false
  @explicit_usage = nil
  @started_at = LlmCostTracker::Timing.now_monotonic
  @finished = false
  @recording = false
  @mutex = Mutex.new
end

Instance Attribute Details

#provider ⇒ Object (readonly)

Returns the value of attribute provider.



14
15
16
 
# File 'lib/llm_cost_tracker/capture/stream_collector.rb', line 14

def provider
  @provider
end

Instance Method Details

#event(data, type: nil) ⇒ Object 



67
68
69
70
71
72
 
# File 'lib/llm_cost_tracker/capture/stream_collector.rb', line 67

def event(data, type: nil)
  @mutex.synchronize do
    ensure_open!
    capture_event(data, type: type) unless data.nil?
  end
end

#finish!(errored: false) ⇒ Object 



91
92
93
94
95
96
 
# File 'lib/llm_cost_tracker/capture/stream_collector.rb', line 91

def finish!(errored: false)
  snapshot = claim_recording_slot
  return if snapshot.nil?

  record_snapshot(snapshot, errored: errored)
end

#metadata ⇒ Object 



45
46
47
 
# File 'lib/llm_cost_tracker/capture/stream_collector.rb', line 45

def metadata
  @mutex.synchronize { @metadata.deep_dup }
end

#model ⇒ Object 



41
42
43
 
# File 'lib/llm_cost_tracker/capture/stream_collector.rb', line 41

def model
  @mutex.synchronize { @model }
end

#model=(value) ⇒ Object 



53
54
55
56
57
58
 
# File 'lib/llm_cost_tracker/capture/stream_collector.rb', line 53

def model=(value)
  @mutex.synchronize do
    ensure_open!
    @model = value
  end
end

#provider_response_id ⇒ Object 



49
50
51
 
# File 'lib/llm_cost_tracker/capture/stream_collector.rb', line 49

def provider_response_id
  @mutex.synchronize { @provider_response_id }
end

#provider_response_id=(value) ⇒ Object 



60
61
62
63
64
65
 
# File 'lib/llm_cost_tracker/capture/stream_collector.rb', line 60

def provider_response_id=(value)
  @mutex.synchronize do
    ensure_open!
    @provider_response_id = value
  end
end

#usage(input_tokens:, output_tokens:, **extra) ⇒ Object 



74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
 
# File 'lib/llm_cost_tracker/capture/stream_collector.rb', line 74

def usage(input_tokens:, output_tokens:, **extra)
  @mutex.synchronize do
    ensure_open!
    @provider_response_id = extra.delete(:provider_response_id) || @provider_response_id
    @provider_project_id = extra.delete(:provider_project_id) || @provider_project_id
    @provider_api_key_id = extra.delete(:provider_api_key_id) || @provider_api_key_id
    @provider_workspace_id = extra.delete(:provider_workspace_id) || @provider_workspace_id
    batch = extra.delete(:batch)
    @batch = batch unless batch.nil?
    @explicit_usage = TokenUsage.build(
      **extra.slice(*TokenUsage.members),
      input_tokens: input_tokens,
      output_tokens: output_tokens
    )
  end
end