Add token usage data to RunCompleted event

thisisartium · Feb 22, 2024 · 8d82ed9 · 8d82ed9
1 parent c31231f
commit 8d82ed9
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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