diff --git a/lib/chat_models/chat_anthropic.ex b/lib/chat_models/chat_anthropic.ex index 40c1ff17..7398e9f7 100644 --- 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 index c7d3185f..a9196176 100644 --- 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 index 5caeeb9a..257c580c 100644 --- 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 index 2435f847..6f4cc113 100644 --- 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 diff --git a/lib/token_usage.ex b/lib/token_usage.ex index 472dbb96..0fb7cf1a 100644 --- a/lib/token_usage.ex +++ b/lib/token_usage.ex @@ -6,7 +6,10 @@ defmodule LangChain.TokenUsage do %TokenUsage{ input: 30, - output: 15 + output: 15, + raw: %{ + "total_tokens" => 29 + } } Input is the tokens from the prompt. Output is the completion or generated @@ -23,11 +26,12 @@ defmodule LangChain.TokenUsage do embedded_schema do field :input, :integer field :output, :integer + field :raw, :map, default: %{} end @type t :: %TokenUsage{} - @create_fields [:input, :output] + @create_fields [:input, :output, :raw] # Anthropic returns only the output token count when streaming deltas @required_fields [:output] diff --git a/test/token_usage_test.exs b/test/token_usage_test.exs index 4b133868..8269cb37 100644 --- a/test/token_usage_test.exs +++ b/test/token_usage_test.exs @@ -7,10 +7,12 @@ defmodule LangChain.TokenUsageTest do describe "new/1" do test "accepts valid data" do - assert {:ok, %TokenUsage{} = usage} = TokenUsage.new(%{"input" => 1, "output" => 2}) + assert {:ok, %TokenUsage{} = usage} = + TokenUsage.new(%{"input" => 1, "output" => 2, "raw" => %{"total_tokens" => 29}}) assert usage.input == 1 assert usage.output == 2 + assert usage.raw == %{"total_tokens" => 29} end test "returns error when invalid" do