Skip to content

Add cost tracking to OpenAi::Client

George Koltsov requested to merge georgekoltsov/track-per-issue-token-spent into master

What does this MR do and why?

This MR adds cost tracking to each request we make to OpenAI. It increment Prometheus counter for each prompt and completion token. A total token amount can be calculated by summing the two, so not adding total tokens count.

As per discussion in https://gitlab.com/gitlab-org/gitlab/-/issues/407334#note_1365929039 & labkit#60 the metric used has a generic name and labels, which is meant to be generic and to be used by other teams when tracking third party vendor costs.

Mentions https://gitlab.com/gitlab-org/gitlab/-/issues/407334

Screenshots or screen recordings

openaimetrics.mov

Screenshots are required for UI changes, and strongly recommended for all other merge requests.

How to set up and validate locally

  1. In order to quickly test in rails console, make the following edit to lib/gitlab/metrics.rb in order to publish metrics from rails console
diff --git a/lib/gitlab/metrics.rb b/lib/gitlab/metrics.rb
index e99761a0459b..d1be3ad3bc67 100644
--- a/lib/gitlab/metrics.rb
+++ b/lib/gitlab/metrics.rb
@@ -33,7 +33,7 @@ def self.add_event(*args)
 
     # Allow access from other metrics related middlewares
     def self.current_transaction
-      WebTransaction.current || BackgroundTransaction.current
+      WebTransaction.current || BackgroundTransaction.current || BackgroundTransaction.new
     end
 
     # Returns the prefix to use for the name of a series.
  1. Open rails console
  2. Set your OpenAI API key Gitlab::CurrentSettings.update!(openai_api_key: "...")
  3. Make a new OpenAI request Gitlab::Llm::OpenAi::Client.new(User.first).chat(content: 'test')
  4. Open /-/metrics page and look for gitlab_cloud_cost_spend_entry_total metric. Observe it's present and it has correct labels and increment count corresponding to prompt/completion tokens amounts

MR acceptance checklist

This checklist encourages us to confirm any changes have been analyzed to reduce risks in quality, performance, reliability, security, and maintainability.

Edited by George Koltsov

Merge request reports

Loading