Class: Llmemory::LLM::TrackingClient

Inherits:
Object
  • Object
show all
Defined in:
lib/llmemory/llm/tracking_client.rb

Overview

Transparent wrapper that records token usage to the per-user ledger.

Instance Method Summary collapse

Constructor Details

#initialize(inner, user_id:, store: nil, api_key: nil) ⇒ TrackingClient

Returns a new instance of TrackingClient.



9
10
11
12
13
14
 
# File 'lib/llmemory/llm/tracking_client.rb', line 9

def initialize(inner, user_id:, store: nil, api_key: nil)
  @inner = inner
  @user_id = user_id
  @store = store
  @api_key = api_key
end

Dynamic Method Handling

This class handles dynamic methods through the method_missing method

#method_missing(method, *args, &block) ⇒ Object 



46
47
48
49
50
51
52
 
# File 'lib/llmemory/llm/tracking_client.rb', line 46

def method_missing(method, *args, &block)
  if inner_client.respond_to?(method)
    inner_client.public_send(method, *args, &block)
  else
    super
  end
end

Instance Method Details

#invoke(prompt) ⇒ Object 



16
17
18
19
20
21
22
23
24
25
26
27
 
# File 'lib/llmemory/llm/tracking_client.rb', line 16

def invoke(prompt)
  response = inner_client.invoke(prompt)
  usage = if response.respond_to?(:usage)
            response.usage
          elsif inner_client.respond_to?(:last_usage)
            inner_client.last_usage
          else
            Usage.zero
          end
  UsageRecorder.record(user_id: @user_id, usage: usage, operation: :invoke, store: @store)
  response
end

#invoke_with_json_schema(prompt, json_schema) ⇒ Object 



29
30
31
32
33
34
 
# File 'lib/llmemory/llm/tracking_client.rb', line 29

def invoke_with_json_schema(prompt, json_schema)
  result = inner_client.invoke_with_json_schema(prompt, json_schema)
  usage = inner_client.respond_to?(:last_usage) ? inner_client.last_usage : Usage.zero
  UsageRecorder.record(user_id: @user_id, usage: usage, operation: :invoke, store: @store)
  result
end

#last_usage ⇒ Object 



36
37
38
39
40
 
# File 'lib/llmemory/llm/tracking_client.rb', line 36

def last_usage
  return inner_client.last_usage if inner_client.respond_to?(:last_usage)

  Usage.zero
end

#respond_to?(method, include_private = false) ⇒ Boolean

Returns:

  • (Boolean) 


42
43
44
 
# File 'lib/llmemory/llm/tracking_client.rb', line 42

def respond_to?(method, include_private = false)
  inner_client.respond_to?(method, include_private) || super
end