From 9f2d058a2a1e4eb99333687dbd9b7e622314ab70 Mon Sep 17 00:00:00 2001
From: Neeraj G <108350739+NeerajG03@users.noreply.github.com>
Date: Thu, 29 Aug 2024 21:56:21 +0530
Subject: [PATCH] support to get usage for ChatGoogleGenerativeAI

---
 langfuse/callback/langchain.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/langfuse/callback/langchain.py b/langfuse/callback/langchain.py
index 7ae8402b..4e18c1ba 100644
--- a/langfuse/callback/langchain.py
+++ b/langfuse/callback/langchain.py
@@ -964,6 +964,11 @@ def _parse_usage(response: LLMResult):
                     break
 
                 message_chunk = getattr(generation_chunk, "message", {})
+
+                if hasattr(message_chunk, "usage_metadata") and message_chunk.usage_metadata is not None: #for ChatGoogleGenerativeAI
+                    llm_usage = _parse_usage_model(message_chunk.usage_metadata)
+                    break
+                
                 response_metadata = getattr(message_chunk, "response_metadata", {})
 
                 chunk_usage = (