Switch to a metric for total tokens used

Author: colin-sentry

sentry_sdk/consts.py

@@ -97,24 +97,6 @@ class SPANDATA:
     Example: [{"role": "user", "message": "hello"}]
     """
 
-    AI_COMPLETION_TOKENS_USED = "ai.completion_tokens.used"
-    """
-    The number of tokens used to respond to an AI model request
-    Example: 10
-    """
-
-    AI_PROMPT_TOKENS_USED = "ai.prompt_tokens.used"
-    """
-    The number of tokens used to process the input text to an AI model request
-    Example: 20
-    """
-
-    AI_TOTAL_TOKENS_USED = "ai.total_tokens.used"
-    """
-    The number of tokens used in total to process an AI model request
-    Example: 30
-    """
-
     AI_MODEL_ID = "ai.model_id"
     """
     The unique descriptor of the model being execugted

sentry_sdk/integrations/_ai_common.py

@@ -1,6 +1,4 @@
-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, Optional
@@ -39,19 +37,14 @@ def record_token_usage(
 ):
     # 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")
+        span.set_measurement("ai_prompt_tokens_used", value=prompt_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"
-        )
+        span.set_measurement("ai_completion_tokens_used", value=completion_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")
+        span.set_measurement("ai_total_tokens_used", total_tokens)

sentry_sdk/integrations/langchain.py

@@ -127,6 +127,9 @@ def _normalize_langchain_message(self, message):
     def _create_span(self, run_id, parent_id, **kwargs):
         # type: (SentryLangchainCallback, UUID, Optional[Any], Any) -> Span
 
+        if "origin" not in kwargs:
+            kwargs["origin"] = "auto.ai.langchain"
+
         span = None  # type: Optional[Span]
         if parent_id:
             parent_span = self.span_map[parent_id]  # type: Optional[WatchedSpan]

sentry_sdk/integrations/openai.py

@@ -140,7 +140,9 @@ def new_chat_completion(*args, **kwargs):
         streaming = kwargs.get("stream")
 
         span = sentry_sdk.start_span(
-            op=consts.OP.OPENAI_CHAT_COMPLETIONS_CREATE, description="Chat Completion"
+            op=consts.OP.OPENAI_CHAT_COMPLETIONS_CREATE,
+            origin="auto.ai.openai",
+            description="Chat Completion",
         )
         span.__enter__()
         try:
@@ -223,6 +225,7 @@ def new_embeddings_create(*args, **kwargs):
         # type: (*Any, **Any) -> Any
         with sentry_sdk.start_span(
             op=consts.OP.OPENAI_EMBEDDINGS_CREATE,
+            origin="auto.ai.openai",
             description="OpenAI Embedding Creation",
         ) as span:
             integration = sentry_sdk.get_client().get_integration(OpenAIIntegration)

tests/integrations/openai/test_openai.py

@@ -7,7 +7,6 @@
 from openai.types.create_embedding_response import Usage as EmbeddingTokenUsage
 
 from sentry_sdk import start_transaction
-from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations.openai import OpenAIIntegration
 
 from unittest import mock  # python 3.3 and above
@@ -74,9 +73,9 @@ def test_nonstreaming_chat_completion(
         assert "ai.input_messages" not in span["data"]
         assert "ai.responses" not in span["data"]
 
-    assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 10
-    assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 20
-    assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
+    assert span["measurements"]["ai_completion_tokens_used"]["value"] == 10
+    assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 20
+    assert span["measurements"]["ai_total_tokens_used"]["value"] == 30
 
 
 # noinspection PyTypeChecker
@@ -156,9 +155,9 @@ def test_streaming_chat_completion(
     try:
         import tiktoken  # type: ignore # noqa # pylint: disable=unused-import
 
-        assert span["data"][SPANDATA.AI_COMPLETION_TOKENS_USED] == 2
-        assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 1
-        assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 3
+        assert span["measurements"]["ai_completion_tokens_used"]["value"] == 2
+        assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 1
+        assert span["measurements"]["ai_total_tokens_used"]["value"] == 3
     except ImportError:
         pass  # if tiktoken is not installed, we can't guarantee token usage will be calculated properly
 
@@ -223,5 +222,5 @@ def test_embeddings_create(
     else:
         assert "ai.input_messages" not in span["data"]
 
-    assert span["data"][SPANDATA.AI_PROMPT_TOKENS_USED] == 20
-    assert span["data"][SPANDATA.AI_TOTAL_TOKENS_USED] == 30
+    assert span["measurements"]["ai_prompt_tokens_used"]["value"] == 20
+    assert span["measurements"]["ai_total_tokens_used"]["value"] == 30