From aa615680ce0c995ef69a20b195edb2963f038ea9 Mon Sep 17 00:00:00 2001 From: Sam Saffron Date: Wed, 22 Jan 2025 11:12:56 +1100 Subject: [PATCH 1/2] FEATURE: track duration of AI calls --- db/migrate/20250122003035_add_duration_to_ai_api_log.rb | 6 ++++++ lib/completions/endpoints/base.rb | 4 ++++ spec/lib/completions/endpoints/open_ai_spec.rb | 1 + 3 files changed, 11 insertions(+) create mode 100644 db/migrate/20250122003035_add_duration_to_ai_api_log.rb diff --git a/db/migrate/20250122003035_add_duration_to_ai_api_log.rb b/db/migrate/20250122003035_add_duration_to_ai_api_log.rb new file mode 100644 index 000000000..a8e99b4a1 --- /dev/null +++ b/db/migrate/20250122003035_add_duration_to_ai_api_log.rb @@ -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 diff --git a/lib/completions/endpoints/base.rb b/lib/completions/endpoints/base.rb index faf9820ae..439b54f55 100644 --- a/lib/completions/endpoints/base.rb +++ b/lib/completions/endpoints/base.rb @@ -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) @@ -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? diff --git a/spec/lib/completions/endpoints/open_ai_spec.rb b/spec/lib/completions/endpoints/open_ai_spec.rb index f8e5d303b..a58061cdd 100644 --- a/spec/lib/completions/endpoints/open_ai_spec.rb +++ b/spec/lib/completions/endpoints/open_ai_spec.rb @@ -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( From 4be621b2e2ae5fff0d3e9bff296943390596767a Mon Sep 17 00:00:00 2001 From: Sam Saffron Date: Thu, 23 Jan 2025 08:44:03 +1100 Subject: [PATCH 2/2] annotate --- app/models/ai_api_audit_log.rb | 1 + 1 file changed, 1 insertion(+) diff --git a/app/models/ai_api_audit_log.rb b/app/models/ai_api_audit_log.rb index 075f14442..69f722677 100644 --- a/app/models/ai_api_audit_log.rb +++ b/app/models/ai_api_audit_log.rb @@ -46,6 +46,7 @@ def prev_log_id # language_model :string(255) # feature_context :jsonb # cached_tokens :integer +# duration_msecs :integer # # Indexes #