From 97cf9d5a8e7fe75e891e080e11498e9b67ad24df Mon Sep 17 00:00:00 2001 From: Mathieu Ripert Date: Thu, 5 Jun 2025 08:35:10 +0200 Subject: [PATCH 1/2] fix: Add token usage to ChatGoogleAI message metadata - ChatGoogleAI was extracting token usage and firing callbacks but not adding it to message metadata like other chat models - Modified do_process_response to use TokenUsage.set() to add usage to each processed message, consistent with ChatOpenAI and ChatAnthropic - Added tests for both Message and MessageDelta with token usage - Fixes issue where last_message.metadata.usage was nil for Google AI --- lib/chat_models/chat_google_ai.ex | 9 ++- test/chat_models/chat_google_ai_test.exs | 71 ++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 3 deletions(-) diff --git a/lib/chat_models/chat_google_ai.ex b/lib/chat_models/chat_google_ai.ex index 37addc4a..77790d5f 100644 --- a/lib/chat_models/chat_google_ai.ex +++ b/lib/chat_models/chat_google_ai.ex @@ -603,9 +603,11 @@ defmodule LangChain.ChatModels.ChatGoogleAI do # Google is odd in that it returns token usage for each MessageDelta as it # goes, incrementing the number of generated tokens. I haven't seen anyone # else do this. For now, we fire each and every TokenUsage we receive. - case get_token_usage(data) do - %TokenUsage{} = token_usage -> - Callbacks.fire(model.callbacks, :on_llm_token_usage, [token_usage]) + token_usage = get_token_usage(data) + + case token_usage do + %TokenUsage{} = usage -> + Callbacks.fire(model.callbacks, :on_llm_token_usage, [usage]) :ok nil -> @@ -614,6 +616,7 @@ defmodule LangChain.ChatModels.ChatGoogleAI do candidates |> Enum.map(&do_process_response(model, &1, message_type)) + |> Enum.map(&TokenUsage.set(&1, token_usage)) end # Function Call in a Message diff --git a/test/chat_models/chat_google_ai_test.exs b/test/chat_models/chat_google_ai_test.exs index b912e682..b147cca6 100644 --- a/test/chat_models/chat_google_ai_test.exs +++ b/test/chat_models/chat_google_ai_test.exs @@ -619,6 +619,77 @@ defmodule ChatModels.ChatGoogleAITest do assert error.type == "unexpected_response" assert error.message == "Unexpected response" end + + test "handles receiving a message with token usage", %{model: model} do + response = %{ + "candidates" => [ + %{ + "content" => %{"role" => "model", "parts" => [%{"text" => "Hello User!"}]}, + "finishReason" => "STOP", + "index" => 0 + } + ], + "usageMetadata" => %{ + "promptTokenCount" => 10, + "candidatesTokenCount" => 5, + "totalTokenCount" => 15 + } + } + + assert [%Message{} = struct] = ChatGoogleAI.do_process_response(model, response) + assert struct.role == :assistant + [%ContentPart{type: :text, content: "Hello User!"}] = struct.content + assert struct.index == 0 + assert struct.status == :complete + + # Verify that token usage is properly included in metadata + assert %TokenUsage{} = struct.metadata.usage + assert struct.metadata.usage.input == 10 + assert struct.metadata.usage.output == 5 + assert struct.metadata.usage.raw == %{ + "promptTokenCount" => 10, + "candidatesTokenCount" => 5, + "totalTokenCount" => 15 + } + end + + test "handles receiving MessageDelta with token usage", %{model: model} do + response = %{ + "candidates" => [ + %{ + "content" => %{ + "role" => "model", + "parts" => [%{"text" => "This is a partial message"}] + }, + "finishReason" => "STOP", + "index" => 0 + } + ], + "usageMetadata" => %{ + "promptTokenCount" => 8, + "candidatesTokenCount" => 3, + "totalTokenCount" => 11 + } + } + + assert [%MessageDelta{} = struct] = + ChatGoogleAI.do_process_response(model, response, MessageDelta) + + assert struct.role == :assistant + assert struct.content == "This is a partial message" + assert struct.index == 0 + assert struct.status == :incomplete + + # Verify that token usage is properly included in metadata + assert %TokenUsage{} = struct.metadata.usage + assert struct.metadata.usage.input == 8 + assert struct.metadata.usage.output == 3 + assert struct.metadata.usage.raw == %{ + "promptTokenCount" => 8, + "candidatesTokenCount" => 3, + "totalTokenCount" => 11 + } + end end describe "filter_parts_for_types/2" do From 67fab7dd5a0df7dcbdc7e189ef4a5628eb8914bc Mon Sep 17 00:00:00 2001 From: Mathieu Ripert Date: Thu, 5 Jun 2025 08:40:05 +0200 Subject: [PATCH 2/2] format --- test/chat_models/chat_google_ai_test.exs | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/test/chat_models/chat_google_ai_test.exs b/test/chat_models/chat_google_ai_test.exs index b147cca6..1671f880 100644 --- a/test/chat_models/chat_google_ai_test.exs +++ b/test/chat_models/chat_google_ai_test.exs @@ -646,11 +646,12 @@ defmodule ChatModels.ChatGoogleAITest do assert %TokenUsage{} = struct.metadata.usage assert struct.metadata.usage.input == 10 assert struct.metadata.usage.output == 5 + assert struct.metadata.usage.raw == %{ - "promptTokenCount" => 10, - "candidatesTokenCount" => 5, - "totalTokenCount" => 15 - } + "promptTokenCount" => 10, + "candidatesTokenCount" => 5, + "totalTokenCount" => 15 + } end test "handles receiving MessageDelta with token usage", %{model: model} do @@ -684,11 +685,12 @@ defmodule ChatModels.ChatGoogleAITest do assert %TokenUsage{} = struct.metadata.usage assert struct.metadata.usage.input == 8 assert struct.metadata.usage.output == 3 + assert struct.metadata.usage.raw == %{ - "promptTokenCount" => 8, - "candidatesTokenCount" => 3, - "totalTokenCount" => 11 - } + "promptTokenCount" => 8, + "candidatesTokenCount" => 3, + "totalTokenCount" => 11 + } end end