From 3af6ca626c8f37c931eb51b846b0b3d24afdb615 Mon Sep 17 00:00:00 2001
From: Roger Yang <80478925+RogerHYang@users.noreply.github.com>
Date: Thu, 9 Jan 2025 12:30:11 -0800
Subject: [PATCH] fix(llama-index): coerce token counts to be integers (#1183)

---
 .../instrumentation/llama_index/_handler.py   | 30 +++++++++++++++----
 1 file changed, 24 insertions(+), 6 deletions(-)

diff --git a/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_handler.py b/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_handler.py
index c7f381253..f0e8b0e68 100644
--- a/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_handler.py
+++ b/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_handler.py
@@ -853,22 +853,40 @@ def _get_token_counts(usage: Union[object, Mapping[str, Any]]) -> Iterator[Tuple
 
 def _get_token_counts_from_object(usage: object) -> Iterator[Tuple[str, Any]]:
     if (prompt_tokens := getattr(usage, "prompt_tokens", None)) is not None:
-        yield LLM_TOKEN_COUNT_PROMPT, prompt_tokens
+        try:
+            yield LLM_TOKEN_COUNT_PROMPT, int(prompt_tokens)
+        except BaseException:
+            pass
     if (completion_tokens := getattr(usage, "completion_tokens", None)) is not None:
-        yield LLM_TOKEN_COUNT_COMPLETION, completion_tokens
+        try:
+            yield LLM_TOKEN_COUNT_COMPLETION, int(completion_tokens)
+        except BaseException:
+            pass
     if (total_tokens := getattr(usage, "total_tokens", None)) is not None:
-        yield LLM_TOKEN_COUNT_TOTAL, total_tokens
+        try:
+            yield LLM_TOKEN_COUNT_TOTAL, int(total_tokens)
+        except BaseException:
+            pass
 
 
 def _get_token_counts_from_mapping(
     usage_mapping: Mapping[str, Any],
 ) -> Iterator[Tuple[str, Any]]:
     if (prompt_tokens := usage_mapping.get("prompt_tokens")) is not None:
-        yield LLM_TOKEN_COUNT_PROMPT, prompt_tokens
+        try:
+            yield LLM_TOKEN_COUNT_PROMPT, int(prompt_tokens)
+        except BaseException:
+            pass
     if (completion_tokens := usage_mapping.get("completion_tokens")) is not None:
-        yield LLM_TOKEN_COUNT_COMPLETION, completion_tokens
+        try:
+            yield LLM_TOKEN_COUNT_COMPLETION, int(completion_tokens)
+        except BaseException:
+            pass
     if (total_tokens := usage_mapping.get("total_tokens")) is not None:
-        yield LLM_TOKEN_COUNT_TOTAL, total_tokens
+        try:
+            yield LLM_TOKEN_COUNT_TOTAL, int(total_tokens)
+        except BaseException:
+            pass
 
 
 @singledispatch