Skip to content

Commit

Permalink
FEATURE: track duration of AI calls (#1082)
Browse files Browse the repository at this point in the history
* FEATURE: track duration of AI calls

* annotate
  • Loading branch information
SamSaffron authored Jan 23, 2025
1 parent faa8e6e commit 8bf3502
Show file tree
Hide file tree
Showing 4 changed files with 12 additions and 0 deletions.
1 change: 1 addition & 0 deletions app/models/ai_api_audit_log.rb
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ def prev_log_id
# language_model :string(255)
# feature_context :jsonb
# cached_tokens :integer
# duration_msecs :integer
#
# Indexes
#
Expand Down
6 changes: 6 additions & 0 deletions db/migrate/20250122003035_add_duration_to_ai_api_log.rb
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
# frozen_string_literal: true
class AddDurationToAiApiLog < ActiveRecord::Migration[7.2]
def change
add_column :ai_api_audit_logs, :duration_msecs, :integer
end
end
4 changes: 4 additions & 0 deletions lib/completions/endpoints/base.rb
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ def perform_completion!(
&blk
)
LlmQuota.check_quotas!(@llm_model, user)
start_time = Time.now

@partial_tool_calls = partial_tool_calls
model_params = normalize_model_params(model_params)
Expand Down Expand Up @@ -212,6 +213,9 @@ def perform_completion!(
log.raw_response_payload = response_raw
final_log_update(log)
log.response_tokens = tokenizer.size(partials_raw) if log.response_tokens.blank?
log.created_at = start_time
log.updated_at = Time.now
log.duration_msecs = (Time.now - start_time) * 1000
log.save!
LlmQuota.log_usage(@llm_model, user, log.request_tokens, log.response_tokens)
if Rails.env.development?
Expand Down
1 change: 1 addition & 0 deletions spec/lib/completions/endpoints/open_ai_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -319,6 +319,7 @@ def request_body(prompt, stream: false, tool_call: false)
log = AiApiAuditLog.order(:id).last
expect(log.request_tokens).to eq(55)
expect(log.response_tokens).to eq(13)
expect(log.duration_msecs).not_to be_nil

expected =
DiscourseAi::Completions::ToolCall.new(
Expand Down

0 comments on commit 8bf3502

Please sign in to comment.