added total and thinking tokens to callback (#923)

Author: lkuligin

libs/vertexai/langchain_google_vertexai/callbacks.py

@@ -13,6 +13,8 @@ class VertexAICallbackHandler(BaseCallbackHandler):
     completion_tokens: int = 0
     completion_characters: int = 0
     successful_requests: int = 0
+    total_tokens: int = 0
+    cached_tokens: int = 0
 
     def __init__(self) -> None:
         super().__init__()
@@ -24,6 +26,8 @@ def __repr__(self) -> str:
             f"\tPrompt characters: {self.prompt_characters}\n"
             f"\tCompletion tokens: {self.completion_tokens}\n"
             f"\tCompletion characters: {self.completion_characters}\n"
+            f"\tCached tokens: {self.cached_tokens}\n"
+            f"\tTotal tokens: {self.total_tokens}\n"
             f"Successful requests: {self.successful_requests}\n"
         )
 
@@ -44,7 +48,7 @@ def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
 
     def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         """Collects token usage."""
-        completion_tokens, prompt_tokens = 0, 0
+        completion_tokens, prompt_tokens, total_tokens, cached_tokens = 0, 0, 0, 0
         completion_characters, prompt_characters = 0, 0
         for generations in response.generations:
             if len(generations) > 0 and generations[0].generation_info:
@@ -53,6 +57,8 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
                 )
                 completion_tokens += usage_metadata.get("candidates_token_count", 0)
                 prompt_tokens += usage_metadata.get("prompt_token_count", 0)
+                total_tokens += usage_metadata.get("total_token_count", 0)
+                cached_tokens += usage_metadata.get("cached_content_token_count", 0)
                 completion_characters += usage_metadata.get(
                     "candidates_billable_characters", 0
                 )
@@ -64,3 +70,5 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
             self.completion_characters += completion_characters
             self.completion_tokens += completion_tokens
             self.successful_requests += 1
+            self.total_tokens += total_tokens
+            self.cached_tokens += cached_tokens