fix(langchain): Stop double-counting usage in multi-candidate responses

scripts/populate_tox/config.py

@@ -254,6 +254,7 @@
             "<=0.1": ["httpx<0.28.0"],
             ">=0.3": ["langchain-community"],
             ">=1.0": ["langchain-classic"],
+            ">=1.1.2": ["google-genai", "langchain-google-genai>=4"],
         },
         "python": {
             "<1.0": "<3.14",  # https://github.com/langchain-ai/langchain/issues/33449#issuecomment-3408876631
@@ -267,6 +268,7 @@
             "<=0.1": ["httpx<0.28.0"],
             ">=0.3": ["langchain-community"],
             ">=1.0": ["langchain-classic"],
+            ">=1.1.2": ["google-genai", "langchain-google-genai>=4"],
         },
         "python": {
             "<1.0": "<3.14",  # https://github.com/langchain-ai/langchain/issues/33449#issuecomment-3408876631

sentry_sdk/integrations/langchain.py

@@ -51,7 +51,13 @@
         manager,
     )
     from langchain_core.messages import BaseMessage
-    from langchain_core.outputs import ChatGeneration, LLMResult
+    from langchain_core.outputs import (
+        ChatGeneration,
+        ChatGenerationChunk,
+        Generation,
+        GenerationChunk,
+        LLMResult,
+    )
 
 except ImportError:
     raise DidNotEnable("langchain not installed")
@@ -713,7 +719,7 @@ def _extract_tokens(
 
 
 def _extract_tokens_from_generations(
-    generations: "Any",
+    generations: "list[list[Generation | ChatGeneration | GenerationChunk | ChatGenerationChunk]]",
 ) -> "tuple[Optional[int], Optional[int], Optional[int]]":
     """Extract token usage from response.generations structure."""
     if not generations:
@@ -724,12 +730,14 @@ def _extract_tokens_from_generations(
     total_total = 0
 
     for gen_list in generations:
-        for gen in gen_list:
-            token_usage = _get_token_usage(gen)
-            input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage)
-            total_input += input_tokens if input_tokens is not None else 0
-            total_output += output_tokens if output_tokens is not None else 0
-            total_total += total_tokens if total_tokens is not None else 0
+        if not gen_list:
+            continue
+
+        token_usage = _get_token_usage(gen_list[0])
+        input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage)
+        total_input += input_tokens if input_tokens is not None else 0
+        total_output += output_tokens if output_tokens is not None else 0
+        total_total += total_tokens if total_tokens is not None else 0
 
     return (
         total_input if total_input > 0 else None,
@@ -766,7 +774,9 @@ def _get_token_usage(obj: "Any") -> "Optional[Dict[str, Any]]":
     return None
 
 
-def _record_token_usage(span: "Union[Span, StreamedSpan]", response: "Any") -> None:
+def _record_token_usage(
+    span: "Union[Span, StreamedSpan]", response: "LLMResult"
+) -> None:
     token_usage = _get_token_usage(response)
     if token_usage:
         input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage)

tests/integrations/langchain/test_langchain.py

@@ -21,6 +21,13 @@
 from langchain_core.outputs import ChatGenerationChunk, ChatResult
 from langchain_core.runnables import RunnableConfig
 
+try:
+    import google
+    from langchain_google_genai import ChatGoogleGenerativeAI
+except ImportError:
+    ChatGoogleGenerativeAI = None
+    google = None
+
 import sentry_sdk
 from sentry_sdk import start_transaction
 from sentry_sdk.integrations.langchain import (
@@ -231,6 +238,43 @@ def inner():
     return inner
 
 
+@pytest.fixture
+def nonstreaming_multi_candidate_google_genai_model_response():
+    return google.genai.types.GenerateContentResponse(
+        response_id="resp_123",
+        candidates=[
+            google.genai.types.Candidate(
+                content=google.genai.types.Content(
+                    role="model",
+                    parts=[
+                        google.genai.types.Part(
+                            text="Hello, how can I help you?",
+                        )
+                    ],
+                ),
+                finish_reason="STOP",
+            ),
+            google.genai.types.Candidate(
+                content=google.genai.types.Content(
+                    role="model",
+                    parts=[
+                        google.genai.types.Part(
+                            text="Hello, how are you?",
+                        )
+                    ],
+                ),
+                finish_reason="STOP",
+            ),
+        ],
+        model_version="gemini/gemini-pro",
+        usage_metadata=google.genai.types.GenerateContentResponseUsageMetadata(
+            prompt_token_count=10,
+            candidates_token_count=20,
+            total_token_count=30,
+        ),
+    )
+
+
 @tool
 def get_word_length(word: str) -> int:
     """Returns the length of a word."""
@@ -487,6 +531,89 @@ def test_langchain_chat_with_run_name(
         )
 
 
+@pytest.mark.skipif(
+    ChatGoogleGenerativeAI is None,
+    reason="Requires langchain-google-genai.",
+)
+@pytest.mark.parametrize("span_streaming", [True, False])
+@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
+def test_langchain_multi_choice_response(
+    sentry_init,
+    capture_events,
+    capture_items,
+    get_model_response,
+    nonstreaming_multi_candidate_google_genai_model_response,
+    stream_gen_ai_spans,
+    span_streaming,
+):
+    sentry_init(
+        integrations=[
+            LangchainIntegration(
+                include_prompts=True,
+            )
+        ],
+        disabled_integrations=[StdlibIntegration],
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+        stream_gen_ai_spans=stream_gen_ai_spans,
+        trace_lifecycle="stream" if span_streaming else "static",
+    )
+
+    model_response = get_model_response(
+        nonstreaming_multi_candidate_google_genai_model_response,
+        serialize_pydantic=True,
+    )
+
+    llm = ChatGoogleGenerativeAI(
+        model="gemini/gemini-pro",
+        temperature=0,
+        google_api_key="badkey",
+    )
+
+    if span_streaming or stream_gen_ai_spans:
+        items = capture_items("span")
+
+        with patch.object(
+            llm.client._api_client._httpx_client,
+            "send",
+            return_value=model_response,
+        ) as _, start_transaction():
+            llm.invoke(
+                "How many letters in the word eudca",
+            )
+
+        sentry_sdk.flush()
+        spans = [item.payload for item in items]
+        chat_spans = list(
+            x for x in spans if x["attributes"]["sentry.op"] == "gen_ai.chat"
+        )
+        assert len(chat_spans) == 1
+
+        assert chat_spans[0]["attributes"]["gen_ai.usage.input_tokens"] == 10
+        assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens"] == 20
+        assert chat_spans[0]["attributes"]["gen_ai.usage.total_tokens"] == 30
+    else:
+        events = capture_events()
+
+        with patch.object(
+            llm.client._api_client._httpx_client,
+            "send",
+            return_value=model_response,
+        ) as _, start_transaction():
+            llm.invoke(
+                "How many letters in the word eudca",
+            )
+
+        tx = events[0]
+
+        chat_spans = list(x for x in tx["spans"] if x["op"] == "gen_ai.chat")
+        assert len(chat_spans) == 1
+
+        assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 10
+        assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 20
+        assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 30
+
+
 @pytest.mark.parametrize("span_streaming", [True, False])
 @pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
 def test_langchain_tool_call_with_run_name(