Repository: run-llama/llama_index
Pull request: Fix Google GenAI token counting behavior
# Description

- Add `usage_metadata` as candidate key for usage tracing
- Prevent doubling llm events for `complete` calls
- Add basic retry policy

Fixes #18875 

## New Package?

Did I fill in the `tool.llamahub` section in the `pyproject.toml` and provide a detailed README.md for my new integration or package?

- [ ] Yes
- [ X ] No

## Version Bump?

Did I bump the version in the `pyproject.toml` file of the package I am updating? (Except for the `llama-index-core` package)

- [ ] Yes
- [ X ] No

## Type of Change

Please delete options that are not relevant.

- [ X ] Bug fix (non-breaking change which fixes an issue)
- [ X ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] This change requires a documentation update

## How Has This Been Tested?

Your pull-request will likely not be merged unless it is covered by some form of impactful unit testing.

- [ ] I added new unit tests to cover this change
- [ X ] I believe this change is already covered by existing unit tests

## Suggested Checklist:

- [ X ] I have performed a self-review of my own code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] I have added Google Colab support for the newly added notebooks.
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] I ran `uv run make format; uv run make lint` to appease the lint gods


Problem statement:
[Bug]: Incorrect token counting for Google GenAI
### Bug Description

Incorrect token counting for Google GenAI models when using `complete` call.

### Version

0.12.37

### Steps to Reproduce

- Create custom callback manager and add token counting handler with verbosity mode on
- Trigger `complete` with a prompt
- Notice how the handler logs 2 different token counts

```python
import tiktoken
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.core.callbacks import CallbackManager, TokenCountingHandler

handler = TokenCountingHandler(
            tokenizer=tiktoken.encoding_for_model('gpt-4o').encode,  # Fallback tokenizer
            verbose=True,
        )

llm = GoogleGenAI(
    model='gemini-1.5-flash',
    api_key='API_KEY_HERE',
    callback_manager=CallbackManager([handler]),
)

llm.complete('Answer with "Hello, World" in German')

print(handler.total_llm_token_count)  # Should be 13 found 32
print(handler.prompt_llm_token_count)  # Should be 9 found 22
print(handler.completion_llm_token_count)  # Should be 4 found 10
```

### Relevant Logs/Tracbacks

```shell
LLM Prompt Token Usage: 13
LLM Completion Token Usage: 6
LLM Prompt Token Usage: 9
LLM Completion Token Usage: 4
32
22
10
```

Hints from issue discussion:

Commit message: Add default retry policy and prevent doubling llm events for completion requests

Patch to review:
diff --git a/llama-index-core/llama_index/core/callbacks/token_counting.py b/llama-index-core/llama_index/core/callbacks/token_counting.py
index b44314f44c..20af7e8d0e 100644
--- a/llama-index-core/llama_index/core/callbacks/token_counting.py
+++ b/llama-index-core/llama_index/core/callbacks/token_counting.py
@@ -44,7 +44,7 @@ def get_tokens_from_response(
     if not isinstance(raw_response, dict):
         raw_response = dict(raw_response or {})
 
-    usage = raw_response.get("usage", {})
+    usage = raw_response.get("usage", raw_response.get("usage_metadata", {}))
     if usage is None:
         usage = response.additional_kwargs
 
@@ -54,8 +54,12 @@ def get_tokens_from_response(
     if not isinstance(usage, dict):
         usage = usage.model_dump()
 
-    possible_input_keys = ("prompt_tokens", "input_tokens")
-    possible_output_keys = ("completion_tokens", "output_tokens")
+    possible_input_keys = ("prompt_tokens", "input_tokens", "prompt_token_count")
+    possible_output_keys = (
+        "completion_tokens",
+        "output_tokens",
+        "candidates_token_count",
+    )
 
     prompt_tokens = 0
     for input_key in possible_input_keys:
diff --git a/llama-index-integrations/llms/llama-index-llms-google-genai/llama_index/llms/google_genai/base.py b/llama-index-integrations/llms/llama-index-llms-google-genai/llama_index/llms/google_genai/base.py
index 967f043d4e..3f02dd9dc8 100644
--- a/llama-index-integrations/llms/llama-index-llms-google-genai/llama_index/llms/google_genai/base.py
+++ b/llama-index-integrations/llms/llama-index-llms-google-genai/llama_index/llms/google_genai/base.py
@@ -1,5 +1,6 @@
 """Google's hosted Gemini API."""
 
+import functools
 import os
 import typing
 from typing import (
@@ -13,8 +14,10 @@
     Sequence,
     Type,
     Union,
+    Callable,
 )
 
+
 import llama_index.core.instrumentation as instrument
 from llama_index.core.base.llms.generic_utils import (
     chat_to_completion_decorator,
@@ -48,13 +51,13 @@
     convert_schema_to_function_declaration,
     prepare_chat_params,
     handle_streaming_flexible_model,
+    create_retry_decorator,
 )
 
 import google.genai
 import google.auth
 import google.genai.types as types
 
-
 dispatcher = instrument.get_dispatcher(__name__)
 
 DEFAULT_MODEL = "gemini-2.0-flash"
@@ -69,6 +72,25 @@ class VertexAIConfig(typing.TypedDict):
     location: Optional[str] = None
 
 
+def llm_retry_decorator(f: Callable[..., Any]) -> Callable[..., Any]:
+    @functools.wraps(f)
+    def wrapper(self, *args: Any, **kwargs: Any) -> Any:
+        max_retries = getattr(self, "max_retries", 0)
+        if max_retries <= 0:
+            return f(self, *args, **kwargs)
+
+        retry = create_retry_decorator(
+            max_retries=max_retries,
+            random_exponential=True,
+            stop_after_delay_seconds=60,
+            min_seconds=1,
+            max_seconds=20,
+        )
+        return retry(f)(self, *args, **kwargs)
+
+    return wrapper
+
+
 class GoogleGenAI(FunctionCallingLLM):
     """
     Google GenAI LLM.
@@ -97,6 +119,11 @@ class GoogleGenAI(FunctionCallingLLM):
         default=None,
         description="The context window of the model. If not provided, the default context window 200000 will be used.",
     )
+    max_retries: int = Field(
+        default=3,
+        description="The maximum number of API retries.",
+        ge=0,
+    )
     is_function_calling_model: bool = Field(
         default=True, description="Whether the model is a function calling model."
     )
@@ -113,6 +140,7 @@ def __init__(
         temperature: float = DEFAULT_TEMPERATURE,
         max_tokens: Optional[int] = None,
         context_window: Optional[int] = None,
+        max_retries: int = 3,
         vertexai_config: Optional[VertexAIConfig] = None,
         http_options: Optional[types.HttpOptions] = None,
         debug_config: Optional[google.genai.client.DebugConfig] = None,
@@ -164,6 +192,7 @@ def __init__(
             context_window=context_window,
             callback_manager=callback_manager,
             is_function_calling_model=is_function_calling_model,
+            max_retries=max_retries,
             **kwargs,
         )
 
@@ -208,21 +237,21 @@ def metadata(self) -> LLMMetadata:
     def complete(
         self, prompt: str, formatted: bool = False, **kwargs: Any
     ) -> CompletionResponse:
-        chat_fn = chat_to_completion_decorator(self.chat)
+        chat_fn = chat_to_completion_decorator(self._chat)
         return chat_fn(prompt, **kwargs)
 
     @llm_completion_callback()
     async def acomplete(
         self, prompt: str, formatted: bool = False, **kwargs: Any
     ) -> CompletionResponse:
-        chat_fn = achat_to_completion_decorator(self.achat)
+        chat_fn = achat_to_completion_decorator(self._achat)
         return await chat_fn(prompt, **kwargs)
 
     @llm_completion_callback()
     def stream_complete(
         self, prompt: str, formatted: bool = False, **kwargs: Any
     ) -> CompletionResponseGen:
-        chat_fn = stream_chat_to_completion_decorator(self.stream_chat)
+        chat_fn = stream_chat_to_completion_decorator(self._stream_chat)
         return chat_fn(prompt, **kwargs)
 
     @llm_completion_callback()
@@ -232,8 +261,8 @@ async def astream_complete(
         chat_fn = astream_chat_to_completion_decorator(self.astream_chat)
         return await chat_fn(prompt, **kwargs)
 
-    @llm_chat_callback()
-    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+    @llm_retry_decorator
+    def _chat(self, messages: Sequence[ChatMessage], **kwargs: Any):
         generation_config = {
             **(self._generation_config or {}),
             **kwargs.pop("generation_config", {}),
@@ -244,10 +273,8 @@ def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
         response = chat.send_message(next_msg.parts)
         return chat_from_gemini_response(response)
 
-    @llm_chat_callback()
-    async def achat(
-        self, messages: Sequence[ChatMessage], **kwargs: Any
-    ) -> ChatResponse:
+    @llm_retry_decorator
+    async def _achat(self, messages: Sequence[ChatMessage], **kwargs: Any):
         generation_config = {
             **(self._generation_config or {}),
             **kwargs.pop("generation_config", {}),
@@ -259,7 +286,16 @@ async def achat(
         return chat_from_gemini_response(response)
 
     @llm_chat_callback()
-    def stream_chat(
+    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+        return self._chat(messages, **kwargs)
+
+    @llm_chat_callback()
+    async def achat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponse:
+        return await self._achat(messages, **kwargs)
+
+    def _stream_chat(
         self, messages: Sequence[ChatMessage], **kwargs: Any
     ) -> ChatResponseGen:
         generation_config = {
@@ -294,7 +330,12 @@ def gen() -> ChatResponseGen:
         return gen()
 
     @llm_chat_callback()
-    async def astream_chat(
+    def stream_chat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponseGen:
+        return self._stream_chat(messages, **kwargs)
+
+    async def _astream_chat(
         self, messages: Sequence[ChatMessage], **kwargs: Any
     ) -> ChatResponseAsyncGen:
         generation_config = {
@@ -334,6 +375,12 @@ async def gen() -> ChatResponseAsyncGen:
 
         return gen()
 
+    @llm_chat_callback()
+    async def astream_chat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponseAsyncGen:
+        return await self._astream_chat(messages, **kwargs)
+
     def _prepare_chat_with_tools(
         self,
         tools: Sequence["BaseTool"],
diff --git a/llama-index-integrations/llms/llama-index-llms-google-genai/llama_index/llms/google_genai/utils.py b/llama-index-integrations/llms/llama-index-llms-google-genai/llama_index/llms/google_genai/utils.py
index 9c3f016bde..e8fa2881ec 100644
--- a/llama-index-integrations/llms/llama-index-llms-google-genai/llama_index/llms/google_genai/utils.py
+++ b/llama-index-integrations/llms/llama-index-llms-google-genai/llama_index/llms/google_genai/utils.py
@@ -1,3 +1,4 @@
+import logging
 from collections.abc import Sequence
 from typing import (
     TYPE_CHECKING,
@@ -12,7 +13,9 @@
 
 import google.genai.types as types
 import google.genai
+import httpx
 from google.genai import _transformers
+from google.genai import errors
 
 from llama_index.core.bridge.pydantic import BaseModel, ValidationError
 from llama_index.core.base.llms.types import (
@@ -24,10 +27,22 @@
     DocumentBlock,
 )
 from llama_index.core.program.utils import _repair_incomplete_json
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    retry_if_exception,
+    stop_after_attempt,
+    stop_after_delay,
+    wait_exponential,
+    wait_random_exponential,
+)
+from tenacity.stop import stop_base
 
 if TYPE_CHECKING:
     from llama_index.core.tools.types import BaseTool
 
+logger = logging.getLogger(__name__)
 
 ROLES_TO_GEMINI: dict[MessageRole, MessageRole] = {
     MessageRole.USER: MessageRole.USER,
@@ -364,3 +379,41 @@ def handle_streaming_flexible_model(
                 return None, current_json
 
     return None, current_json
+
+
+def _should_retry(exception: BaseException):
+    match exception:
+        case errors.ClientError(status=429) | errors.ClientError(status=408):
+            return True
+    return False
+
+
+def create_retry_decorator(
+    max_retries: int,
+    random_exponential: bool = False,
+    stop_after_delay_seconds: Optional[float] = None,
+    min_seconds: float = 4,
+    max_seconds: float = 60,
+) -> typing.Callable[[Any], Any]:
+    wait_strategy = (
+        wait_random_exponential(min=min_seconds, max=max_seconds)
+        if random_exponential
+        else wait_exponential(multiplier=1, min=min_seconds, max=max_seconds)
+    )
+
+    stop_strategy: stop_base = stop_after_attempt(max_retries)
+    if stop_after_delay_seconds is not None:
+        stop_strategy = stop_strategy | stop_after_delay(stop_after_delay_seconds)
+
+    return retry(
+        reraise=True,
+        stop=stop_strategy,
+        wait=wait_strategy,
+        retry=(
+            retry_if_exception_type(
+                (errors.ServerError, httpx.ConnectError, httpx.ConnectTimeout)
+            )
+            | retry_if_exception(_should_retry)
+        ),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )


Code review comments:
diff_hunk: @@ -364,3 +379,41 @@ def handle_streaming_flexible_model(
                 return None, current_json
 
     return None, current_json
+
+
+def _should_retry(exception: BaseException):
+    match exception:
path: llama-index-integrations/llms/llama-index-llms-google-genai/llama_index/llms/google_genai/utils.py
text: @user1:
Hey there @author, this whole PR is great (really thanks a lot🙏) but this `_should_retry` function uses a match statement that is not compatible with Python 3.9.

I'm gonna switch this to:

```python
def _should_retry(exception: BaseException):
    if isinstance(exception, errors.ClientError):
        if exception.status in (429, 408):
            return True
    return False
```

In the next commit :)