fix and adapt tests for token extraction

Author: shellmayr

tests/integrations/langchain/test_langchain.py

@@ -54,15 +54,6 @@ def _llm_type(self) -> str:
         return llm_type
 
 
-def tiktoken_encoding_if_installed():
-    try:
-        import tiktoken  # type: ignore # noqa # pylint: disable=unused-import
-
-        return "cl100k_base"
-    except ImportError:
-        return None
-
-
 @pytest.mark.parametrize(
     "send_default_pii, include_prompts, use_unknown_llm_type",
     [
@@ -82,7 +73,6 @@ def test_langchain_agent(
         integrations=[
             LangchainIntegration(
                 include_prompts=include_prompts,
-                tiktoken_encoding_name=tiktoken_encoding_if_installed(),
             )
         ],
         traces_sample_rate=1.0,
@@ -144,15 +134,33 @@ def test_langchain_agent(
                 ),
                 ChatGenerationChunk(
                     type="ChatGenerationChunk",
-                    message=AIMessageChunk(content="5"),
+                    message=AIMessageChunk(
+                        content="5",
+                        usage_metadata={
+                            "input_tokens": 142,
+                            "output_tokens": 50,
+                            "total_tokens": 192,
+                            "input_token_details": {"audio": 0, "cache_read": 0},
+                            "output_token_details": {"audio": 0, "reasoning": 0},
+                        },
+                    ),
                     generation_info={"finish_reason": "function_call"},
                 ),
             ],
             [
                 ChatGenerationChunk(
                     text="The word eudca has 5 letters.",
                     type="ChatGenerationChunk",
-                    message=AIMessageChunk(content="The word eudca has 5 letters."),
+                    message=AIMessageChunk(
+                        content="The word eudca has 5 letters.",
+                        usage_metadata={
+                            "input_tokens": 89,
+                            "output_tokens": 28,
+                            "total_tokens": 117,
+                            "input_token_details": {"audio": 0, "cache_read": 0},
+                            "output_token_details": {"audio": 0, "reasoning": 0},
+                        },
+                    ),
                 ),
                 ChatGenerationChunk(
                     type="ChatGenerationChunk",
@@ -182,15 +190,22 @@ def test_langchain_agent(
     assert len(chat_spans) == 2
 
     # We can't guarantee anything about the "shape" of the langchain execution graph
-    assert len(list(x for x in tx["spans"] if x["op"] == "gen_ai.run")) > 0
+    assert len(list(x for x in tx["spans"] if x["op"] == "gen_ai.chat")) > 0
 
-    if use_unknown_llm_type:
-        assert "gen_ai.usage.input_tokens" in chat_spans[0]["data"]
-        assert "gen_ai.usage.total_tokens" in chat_spans[0]["data"]
-    else:
-        # important: to avoid double counting, we do *not* measure
-        # tokens used if we have an explicit integration (e.g. OpenAI)
-        assert "measurements" not in chat_spans[0]
+    assert "gen_ai.usage.input_tokens" in chat_spans[0]["data"]
+    assert "gen_ai.usage.output_tokens" in chat_spans[0]["data"]
+    assert "gen_ai.usage.total_tokens" in chat_spans[0]["data"]
+
+    assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 142
+    assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 50
+    assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 192
+
+    assert "gen_ai.usage.input_tokens" in chat_spans[1]["data"]
+    assert "gen_ai.usage.output_tokens" in chat_spans[1]["data"]
+    assert "gen_ai.usage.total_tokens" in chat_spans[1]["data"]
+    assert chat_spans[1]["data"]["gen_ai.usage.input_tokens"] == 89
+    assert chat_spans[1]["data"]["gen_ai.usage.output_tokens"] == 28
+    assert chat_spans[1]["data"]["gen_ai.usage.total_tokens"] == 117
 
     if send_default_pii and include_prompts:
         assert (
@@ -311,15 +326,33 @@ def test_span_origin(sentry_init, capture_events):
                 ),
                 ChatGenerationChunk(
                     type="ChatGenerationChunk",
-                    message=AIMessageChunk(content="5"),
+                    message=AIMessageChunk(
+                        content="5",
+                        usage_metadata={
+                            "input_tokens": 142,
+                            "output_tokens": 50,
+                            "total_tokens": 192,
+                            "input_token_details": {"audio": 0, "cache_read": 0},
+                            "output_token_details": {"audio": 0, "reasoning": 0},
+                        },
+                    ),
                     generation_info={"finish_reason": "function_call"},
                 ),
             ],
             [
                 ChatGenerationChunk(
                     text="The word eudca has 5 letters.",
                     type="ChatGenerationChunk",
-                    message=AIMessageChunk(content="The word eudca has 5 letters."),
+                    message=AIMessageChunk(
+                        content="The word eudca has 5 letters.",
+                        usage_metadata={
+                            "input_tokens": 89,
+                            "output_tokens": 28,
+                            "total_tokens": 117,
+                            "input_token_details": {"audio": 0, "cache_read": 0},
+                            "output_token_details": {"audio": 0, "reasoning": 0},
+                        },
+                    ),
                 ),
                 ChatGenerationChunk(
                     type="ChatGenerationChunk",