Add support for cached tokens

Author: SamSaffron

app/models/ai_api_audit_log.rb

@@ -43,3 +43,10 @@ def prev_log_id
 #  feature_name         :string(255)
 #  language_model       :string(255)
 #  feature_context      :jsonb
+#  cached_tokens        :integer
+#
+# Indexes
+#
+#  index_ai_api_audit_logs_on_created_at_and_feature_name    (created_at,feature_name)
+#  index_ai_api_audit_logs_on_created_at_and_language_model  (created_at,language_model)
+#

app/serializers/ai_usage_serializer.rb

@@ -4,15 +4,35 @@ class AiUsageSerializer < ApplicationSerializer
   attributes :data, :features, :models, :summary
 
   def data
-    object.tokens_by_period.map {|key, value| [key, value]}
+    object.tokens_by_period.as_json(
+      only: %i[period total_tokens total_cached_tokens total_request_tokens total_response_tokens],
+    )
   end
 
   def features
-    object.feature_breakdown.as_json(only: %i[feature_name usage_count total_tokens])
+    object.feature_breakdown.as_json(
+      only: %i[
+        feature_name
+        usage_count
+        total_tokens
+        total_cached_tokens
+        total_request_tokens
+        total_response_tokens
+      ],
+    )
   end
 
   def models
-    object.model_breakdown.as_json(only: %i[llm usage_count total_tokens])
+    object.model_breakdown.as_json(
+      only: %i[
+        llm
+        usage_count
+        total_tokens
+        total_cached_tokens
+        total_request_tokens
+        total_response_tokens
+      ],
+    )
   end
 
   def summary

assets/javascripts/discourse/components/ai-usage.gjs

@@ -59,11 +59,11 @@ export default class AiUsage extends Component {
     return {
       type: "line",
       data: {
-        labels: this.data.data.map((pair) => pair[0]),
+        labels: this.data.data.map((row) => row.period),
         datasets: [
           {
             label: "Tokens",
-            data: this.data.data.map((pair) => pair[1]),
+            data: this.data.data.map((row) => row.total_tokens),
             fill: false,
             borderColor: "rgb(75, 192, 192)",
             tension: 0.1,

db/migrate/20241128010221_add_cached_tokens_to_ai_api_audit_log.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class AddCachedTokensToAiApiAuditLog < ActiveRecord::Migration[7.2]
+  def change
+    add_column :ai_api_audit_logs, :cached_tokens, :integer
+    add_index :ai_api_audit_logs, [:created_at, :feature_name]
+    add_index :ai_api_audit_logs, [:created_at, :language_model]
+  end
+end

lib/completions/endpoints/open_ai.rb

@@ -97,6 +97,7 @@ def prepare_request(payload)
         def final_log_update(log)
           log.request_tokens = processor.prompt_tokens if processor.prompt_tokens
           log.response_tokens = processor.completion_tokens if processor.completion_tokens
+          log.cached_tokens = processor.cached_tokens if processor.cached_tokens
         end
 
         def decode(response_raw)

lib/completions/open_ai_message_processor.rb

@@ -1,13 +1,14 @@
 # frozen_string_literal: true
 module DiscourseAi::Completions
   class OpenAiMessageProcessor
-    attr_reader :prompt_tokens, :completion_tokens
+    attr_reader :prompt_tokens, :completion_tokens, :cached_tokens
 
     def initialize(partial_tool_calls: false)
       @tool = nil
       @tool_arguments = +""
       @prompt_tokens = nil
       @completion_tokens = nil
+      @cached_tokens = nil
       @partial_tool_calls = partial_tool_calls
     end
 
@@ -121,6 +122,7 @@ def process_arguments
     def update_usage(json)
       @prompt_tokens ||= json.dig(:usage, :prompt_tokens)
       @completion_tokens ||= json.dig(:usage, :completion_tokens)
+      @cached_tokens ||= json.dig(:usage, :prompt_tokens_details, :cached_tokens)
     end
   end
 end

lib/completions/report.rb

@@ -28,8 +28,12 @@ def guess_period(period)
 
       def tokens_by_period(period = nil)
         period = guess_period(period)
-        base_query.group("DATE_TRUNC('#{period}', created_at)").sum(
-          "request_tokens + response_tokens",
+        base_query.group("DATE_TRUNC('#{period}', created_at)").select(
+          "DATE_TRUNC('#{period}', created_at) as period",
+          "SUM(request_tokens + response_tokens) as total_tokens",
+          "SUM(COALESCE(cached_tokens,0)) as total_cached_tokens",
+          "SUM(request_tokens) as total_request_tokens",
+          "SUM(response_tokens) as total_response_tokens",
         )
       end
 
@@ -38,6 +42,9 @@ def feature_breakdown
           "feature_name",
           "COUNT(*) as usage_count",
           "SUM(request_tokens + response_tokens) as total_tokens",
+          "SUM(COALESCE(cached_tokens,0)) as total_cached_tokens",
+          "SUM(request_tokens) as total_request_tokens",
+          "SUM(response_tokens) as total_response_tokens",
         )
       end
 
@@ -46,6 +53,9 @@ def model_breakdown
           "language_model as llm",
           "COUNT(*) as usage_count",
           "SUM(request_tokens + response_tokens) as total_tokens",
+          "SUM(COALESCE(cached_tokens,0)) as total_cached_tokens",
+          "SUM(request_tokens) as total_request_tokens",
+          "SUM(response_tokens) as total_response_tokens",
         )
       end
 

spec/fixtures/bot/openai_artifact_call.txt

@@ -292,7 +292,7 @@ data: {"id":"chatcmpl-ATimVYagKnCWQ0VXY0Hn2SDjRuN6B","object":"chat.completion.c
 
 data: {"id":"chatcmpl-ATimVYagKnCWQ0VXY0Hn2SDjRuN6B","object":"chat.completion.chunk","created":1731647015,"model":"gpt-4o-2024-08-06","system_fingerprint":"fp_45cf54deae","choices":[{"index":0,"delta":{},"logprobs":null,"finish_reason":"tool_calls"}],"usage":null}
 
-data: {"id":"chatcmpl-ATimVYagKnCWQ0VXY0Hn2SDjRuN6B","object":"chat.completion.chunk","created":1731647015,"model":"gpt-4o-2024-08-06","system_fingerprint":"fp_45cf54deae","choices":[],"usage":{"prompt_tokens":735,"completion_tokens":156,"total_tokens":891,"prompt_tokens_details":{"cached_tokens":0,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":0,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}}}
+data: {"id":"chatcmpl-ATimVYagKnCWQ0VXY0Hn2SDjRuN6B","object":"chat.completion.chunk","created":1731647015,"model":"gpt-4o-2024-08-06","system_fingerprint":"fp_45cf54deae","choices":[],"usage":{"prompt_tokens":735,"completion_tokens":156,"total_tokens":891,"prompt_tokens_details":{"cached_tokens":33,"audio_tokens":0},"completion_tokens_details":{"reasoning_tokens":0,"audio_tokens":0,"accepted_prediction_tokens":0,"rejected_prediction_tokens":0}}}
 
 data: [DONE]
 

spec/lib/completions/endpoints/open_ai_spec.rb

@@ -657,6 +657,9 @@ def request_body(prompt, stream: false, tool_call: false)
               end
             end
           end
+
+          audit_log = AiApiAuditLog.order("id desc").first
+          expect(audit_log.cached_tokens).to eq(33)
         end
 
         it "properly handles spaces in tools payload and partial tool calls" do