From 8d82ed95fb929f1ccf7b8526c223ab767fe3de1d Mon Sep 17 00:00:00 2001
From: John Wilger <john@johnwilger.com>
Date: Wed, 21 Feb 2024 17:36:28 -0800
Subject: [PATCH] Add token usage data to RunCompleted event

---
 lib/gpt_agent.ex                      |  7 +++++--
 lib/gpt_agent/events/run_completed.ex |  3 +++
 test/gpt_agent_test.exs               | 12 ++++++++++--
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/lib/gpt_agent.ex b/lib/gpt_agent.ex
index 8cd50af..2572915 100644
--- a/lib/gpt_agent.ex
+++ b/lib/gpt_agent.ex
@@ -359,7 +359,7 @@ defmodule GptAgent do
     handle_run_status(status, id, response, state)
   end
 
-  defp handle_run_status("completed", id, _response, %__MODULE__{} = state) do
+  defp handle_run_status("completed", id, response, %__MODULE__{} = state) do
     log("Run ID #{inspect(id)} completed")
 
     state
@@ -368,7 +368,10 @@ defmodule GptAgent do
       RunCompleted.new!(
         id: id,
         thread_id: state.thread_id,
-        assistant_id: state.assistant_id
+        assistant_id: state.assistant_id,
+        prompt_tokens: response |> Map.get("usage", %{}) |> Map.get("prompt_tokens", 0),
+        completion_tokens: response |> Map.get("usage", %{}) |> Map.get("completion_tokens", 0),
+        total_tokens: response |> Map.get("usage", %{}) |> Map.get("total_tokens", 0)
       )
     )
     |> noreply({:continue, :read_messages})
diff --git a/lib/gpt_agent/events/run_completed.ex b/lib/gpt_agent/events/run_completed.ex
index bd66e64..ae4e54c 100644
--- a/lib/gpt_agent/events/run_completed.ex
+++ b/lib/gpt_agent/events/run_completed.ex
@@ -10,5 +10,8 @@ defmodule GptAgent.Events.RunCompleted do
     field :id, Types.run_id()
     field :thread_id, Types.thread_id()
     field :assistant_id, Types.assistant_id()
+    field :prompt_tokens, non_neg_integer()
+    field :completion_tokens, non_neg_integer()
+    field :total_tokens, non_neg_integer()
   end
 end
diff --git a/test/gpt_agent_test.exs b/test/gpt_agent_test.exs
index 7c25e85..39f0e41 100644
--- a/test/gpt_agent_test.exs
+++ b/test/gpt_agent_test.exs
@@ -632,7 +632,12 @@ defmodule GptAgentTest do
             "instructions" => nil,
             "tools" => [],
             "file_ids" => [],
-            "metadata" => %{}
+            "metadata" => %{},
+            "usage" => %{
+              "prompt_tokens" => 1,
+              "completion_tokens" => 2,
+              "total_tokens" => 3
+            }
           })
         )
       end)
@@ -643,7 +648,10 @@ defmodule GptAgentTest do
                       %RunCompleted{
                         id: ^run_id,
                         thread_id: ^thread_id,
-                        assistant_id: ^assistant_id
+                        assistant_id: ^assistant_id,
+                        prompt_tokens: 1,
+                        completion_tokens: 2,
+                        total_tokens: 3
                       }},
                      5_000
     end