Record metrics for AI tokens used

Author: colin-sentry

sentry_sdk/integrations/_ai_common.py

@@ -1,7 +1,9 @@
+from sentry_sdk import metrics
 from sentry_sdk._types import TYPE_CHECKING
+from sentry_sdk.consts import SPANDATA
 
 if TYPE_CHECKING:
-    from typing import Any
+    from typing import Any, Optional
 
 from sentry_sdk.tracing import Span
 from sentry_sdk.utils import logger
@@ -30,3 +32,26 @@ def set_data_normalized(span, key, value):
     # type: (Span, str, Any) -> None
     normalized = _normalize_data(value)
     span.set_data(key, normalized)
+
+
+def record_token_usage(
+    span, prompt_tokens=None, completion_tokens=None, total_tokens=None
+):
+    # type: (Span, Optional[int], Optional[int], Optional[int]) -> None
+    if prompt_tokens is not None:
+        span.set_data(SPANDATA.AI_PROMPT_TOKENS_USED, prompt_tokens)
+        metrics.incr(SPANDATA.AI_PROMPT_TOKENS_USED, value=prompt_tokens, unit="tokens")
+    if completion_tokens is not None:
+        span.set_data(SPANDATA.AI_COMPLETION_TOKENS_USED, completion_tokens)
+        metrics.incr(
+            SPANDATA.AI_COMPLETION_TOKENS_USED, value=completion_tokens, unit="tokens"
+        )
+    if (
+        total_tokens is None
+        and prompt_tokens is not None
+        and completion_tokens is not None
+    ):
+        total_tokens = prompt_tokens + completion_tokens
+    if total_tokens is not None:
+        span.set_data(SPANDATA.AI_TOTAL_TOKENS_USED, total_tokens)
+        metrics.incr(SPANDATA.AI_TOTAL_TOKENS_USED, value=total_tokens, unit="tokens")

sentry_sdk/integrations/langchain.py

@@ -4,7 +4,7 @@
 import sentry_sdk
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.consts import OP, SPANDATA
-from sentry_sdk.integrations._ai_common import set_data_normalized
+from sentry_sdk.integrations._ai_common import set_data_normalized, record_token_usage
 from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.tracing import Span
 
@@ -244,31 +244,18 @@ def on_llm_end(self, response, *, run_id, **kwargs):
                 )
 
             if token_usage:
-                span_data.span.set_data(
-                    SPANDATA.AI_PROMPT_TOKENS_USED, token_usage.get("prompt_tokens")
-                )
-                span_data.span.set_data(
-                    SPANDATA.AI_COMPLETION_TOKENS_USED,
+                record_token_usage(
+                    span_data.span,
+                    token_usage.get("prompt_tokens"),
                     token_usage.get("completion_tokens"),
-                )
-                span_data.span.set_data(
-                    SPANDATA.AI_TOTAL_TOKENS_USED, token_usage.get("total_tokens")
+                    token_usage.get("total_tokens"),
                 )
             else:
-                if span_data.num_completion_tokens:
-                    span_data.span.set_data(
-                        SPANDATA.AI_COMPLETION_TOKENS_USED,
-                        span_data.num_completion_tokens,
-                    )
-                if span_data.num_prompt_tokens:
-                    span_data.span.set_data(
-                        SPANDATA.AI_PROMPT_TOKENS_USED, span_data.num_prompt_tokens
-                    )
-                if span_data.num_prompt_tokens and span_data.num_completion_tokens:
-                    span_data.span.set_data(
-                        SPANDATA.AI_TOTAL_TOKENS_USED,
-                        span_data.num_prompt_tokens + span_data.num_completion_tokens,
-                    )
+                record_token_usage(
+                    span_data.span,
+                    span_data.num_prompt_tokens,
+                    span_data.num_completion_tokens,
+                )
 
             span_data.span.__exit__(None, None, None)
             del self.span_map[run_id]

sentry_sdk/integrations/openai.py

@@ -3,7 +3,7 @@
 from sentry_sdk import consts
 from sentry_sdk._types import TYPE_CHECKING
 from sentry_sdk.consts import SPANDATA
-from sentry_sdk.integrations._ai_common import set_data_normalized
+from sentry_sdk.integrations._ai_common import set_data_normalized, record_token_usage
 
 if TYPE_CHECKING:
     from typing import Any, Iterable, List, Optional, Callable, Iterator
@@ -109,15 +109,11 @@ def _calculate_chat_completion_usage(
                 if hasattr(choice, "message"):
                     completion_tokens += count_tokens(choice.message)
 
-    if total_tokens == 0:
-        total_tokens = prompt_tokens + completion_tokens
-
-    if completion_tokens != 0:
-        set_data_normalized(span, SPANDATA.AI_COMPLETION_TOKENS_USED, completion_tokens)
-    if prompt_tokens != 0:
-        set_data_normalized(span, SPANDATA.AI_PROMPT_TOKENS_USED, prompt_tokens)
-    if total_tokens != 0:
-        set_data_normalized(span, SPANDATA.AI_TOTAL_TOKENS_USED, total_tokens)
+    if prompt_tokens == 0:
+        prompt_tokens = None
+    if completion_tokens == 0:
+        completion_tokens = None
+    record_token_usage(span, prompt_tokens, completion_tokens, total_tokens)
 
 
 def _wrap_chat_completion_create(f):
@@ -262,11 +258,7 @@ def new_embeddings_create(*args, **kwargs):
             if prompt_tokens == 0:
                 prompt_tokens = count_tokens(kwargs["input"] or "")
 
-            if total_tokens == 0:
-                total_tokens = prompt_tokens
-
-            set_data_normalized(span, SPANDATA.AI_PROMPT_TOKENS_USED, prompt_tokens)
-            set_data_normalized(span, SPANDATA.AI_TOTAL_TOKENS_USED, total_tokens)
+            record_token_usage(span, prompt_tokens, None, total_tokens or prompt_tokens)
 
             return response