Fill raw of TokenUsage in chat models

nallwhy · Jan 16, 2025 · 0784df4 · 0784df4
1 parent c9081ff
commit 0784df4
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 5 deletions.
diff --git a/lib/chat_models/chat_anthropic.ex b/lib/chat_models/chat_anthropic.ex
@@ -1021,7 +1021,8 @@ defmodule LangChain.ChatModels.ChatAnthropic do
     #  https://platform.openai.com/docs/api-reference/chat/object#chat/object-usage
     TokenUsage.new!(%{
       input: Map.get(usage, "input_tokens"),
-      output: Map.get(usage, "output_tokens")
+      output: Map.get(usage, "output_tokens"),
+      raw: usage
     })
   end
 

diff --git a/lib/chat_models/chat_bumblebee.ex b/lib/chat_models/chat_bumblebee.ex
@@ -322,9 +322,9 @@ defmodule LangChain.ChatModels.ChatBumblebee do
     [result]
   end
 
-  defp fire_token_usage_callback(model, %{input: input, output: output} = _token_summary) do
+  defp fire_token_usage_callback(model, %{input: input, output: output} = token_summary) do
     Callbacks.fire(model.callbacks, :on_llm_token_usage, [
-      TokenUsage.new!(%{input: input, output: output})
+      TokenUsage.new!(%{input: input, output: output, raw: token_summary})
     ])
   end
 

diff --git a/lib/chat_models/chat_google_ai.ex b/lib/chat_models/chat_google_ai.ex
@@ -739,7 +739,8 @@ defmodule LangChain.ChatModels.ChatGoogleAI do
     # extract out the reported response token usage
     TokenUsage.new!(%{
       input: Map.get(usage, "promptTokenCount", 0),
-      output: Map.get(usage, "candidatesTokenCount", 0)
+      output: Map.get(usage, "candidatesTokenCount", 0),
+      raw: usage
     })
   end
 

diff --git a/lib/chat_models/chat_open_ai.ex b/lib/chat_models/chat_open_ai.ex
@@ -1046,7 +1046,8 @@ defmodule LangChain.ChatModels.ChatOpenAI do
     #  https://platform.openai.com/docs/api-reference/chat/object#chat/object-usage
     TokenUsage.new!(%{
       input: Map.get(usage, "prompt_tokens"),
-      output: Map.get(usage, "completion_tokens")
+      output: Map.get(usage, "completion_tokens"),
+      raw: usage
     })
   end