From 43cd4aad532b575047cf5bdbc7b08125883dd285 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Tue, 15 Sep 2026 15:01:03 +0200 Subject: [PATCH] ref(langgraph): Stop setting tokens on Invoke Agent spans --- sentry_sdk/integrations/langgraph.py | 32 -- .../integrations/langgraph/test_langgraph.py | 475 ------------------ 2 files changed, 507 deletions(-) diff --git a/sentry_sdk/integrations/langgraph.py b/sentry_sdk/integrations/langgraph.py index 32f9f32a10..75b0c90b7b 100644 --- a/sentry_sdk/integrations/langgraph.py +++ b/sentry_sdk/integrations/langgraph.py @@ -268,37 +268,6 @@ def _extract_tool_calls(messages: "Optional[List[Any]]") -> "Optional[List[Any]] return tool_calls if tool_calls else None -def _set_usage_data(span: "Span", messages: "Any") -> None: - input_tokens = 0 - output_tokens = 0 - total_tokens = 0 - - for message in messages: - response_metadata = message.get("response_metadata") - if response_metadata is None: - continue - - token_usage = response_metadata.get("token_usage") - if not token_usage: - continue - - input_tokens += int(token_usage.get("prompt_tokens", 0)) - output_tokens += int(token_usage.get("completion_tokens", 0)) - total_tokens += int(token_usage.get("total_tokens", 0)) - - if input_tokens > 0: - span.set_attribute(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens) - - if output_tokens > 0: - span.set_attribute(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens) - - if total_tokens > 0: - span.set_attribute( - SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, - total_tokens, - ) - - def _set_response_model_name(span: "Span", messages: "Any") -> None: if len(messages) == 0: return @@ -327,7 +296,6 @@ def _set_response_attributes( if new_messages is None: return - _set_usage_data(span, new_messages) _set_response_model_name(span, new_messages) if _should_record_outputs(integration): diff --git a/tests/integrations/langgraph/test_langgraph.py b/tests/integrations/langgraph/test_langgraph.py index c0d6c182ca..51a32990a9 100644 --- a/tests/integrations/langgraph/test_langgraph.py +++ b/tests/integrations/langgraph/test_langgraph.py @@ -473,342 +473,6 @@ def original_invoke(self, *args, **kwargs): assert SPANDATA.GEN_AI_AGENT_NAME not in invoke_span.get("attributes", {}) -def test_pregel_invoke_span_includes_usage_data( - sentry_init, - capture_items, -): - """ - Test that invoke_agent spans include aggregated usage data from context_wrapper. - This verifies the new functionality added to track token usage in invoke_agent spans. - """ - sentry_init( - integrations=[LanggraphIntegration()], - traces_sample_rate=1.0, - ) - - test_state = { - "messages": [ - MockMessage("Hello, can you help me?", name="user"), - MockMessage("Of course! How can I assist you?", name="assistant"), - ] - } - - pregel = MockPregelInstance("test_graph") - - expected_assistant_response = "I'll help you with that task!" - expected_tool_calls = [ - { - "id": "call_test_123", - "type": "function", - "function": {"name": "search_tool", "arguments": '{"query": "help"}'}, - } - ] - - def original_invoke(self, *args, **kwargs): - input_messages = args[0].get("messages", []) - new_messages = input_messages + [ - MockMessage( - content=expected_assistant_response, - name="assistant", - tool_calls=expected_tool_calls, - response_metadata={ - "token_usage": { - "total_tokens": 30, - "prompt_tokens": 10, - "completion_tokens": 20, - }, - "model_name": "gpt-4.1-2025-04-14", - }, - ) - ] - return {"messages": new_messages} - - items = capture_items("span") - - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span has usage data - assert invoke_agent_span["name"] == "invoke_agent test_graph" - assert "gen_ai.usage.input_tokens" in invoke_agent_span["attributes"] - assert "gen_ai.usage.output_tokens" in invoke_agent_span["attributes"] - assert "gen_ai.usage.total_tokens" in invoke_agent_span["attributes"] - - # The usage should match the mock_usage values (aggregated across all calls) - assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 10 - assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 30 - - -def test_pregel_ainvoke_span_includes_usage_data( - sentry_init, - capture_items, -): - """ - Test that invoke_agent spans include aggregated usage data from context_wrapper. - This verifies the new functionality added to track token usage in invoke_agent spans. - """ - sentry_init( - integrations=[LanggraphIntegration()], - traces_sample_rate=1.0, - ) - - test_state = { - "messages": [ - MockMessage("Hello, can you help me?", name="user"), - MockMessage("Of course! How can I assist you?", name="assistant"), - ] - } - - pregel = MockPregelInstance("test_graph") - - expected_assistant_response = "I'll help you with that task!" - expected_tool_calls = [ - { - "id": "call_test_123", - "type": "function", - "function": {"name": "search_tool", "arguments": '{"query": "help"}'}, - } - ] - - async def original_ainvoke(self, *args, **kwargs): - input_messages = args[0].get("messages", []) - new_messages = input_messages + [ - MockMessage( - content=expected_assistant_response, - name="assistant", - tool_calls=expected_tool_calls, - response_metadata={ - "token_usage": { - "total_tokens": 30, - "prompt_tokens": 10, - "completion_tokens": 20, - }, - "model_name": "gpt-4.1-2025-04-14", - }, - ) - ] - return {"messages": new_messages} - - async def run_test(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - result = await wrapped_ainvoke(pregel, test_state) - return result - - items = capture_items("span") - - result = asyncio.run(run_test()) - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span has usage data - assert invoke_agent_span["name"] == "invoke_agent test_graph" - assert "gen_ai.usage.input_tokens" in invoke_agent_span["attributes"] - assert "gen_ai.usage.output_tokens" in invoke_agent_span["attributes"] - assert "gen_ai.usage.total_tokens" in invoke_agent_span["attributes"] - - # The usage should match the mock_usage values (aggregated across all calls) - assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 10 - assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 30 - - -def test_pregel_invoke_multiple_llm_calls_aggregate_usage( - sentry_init, - capture_items, -): - """ - Test that invoke_agent spans show aggregated usage across multiple LLM calls - (e.g., when tools are used and multiple API calls are made). - """ - sentry_init( - integrations=[LanggraphIntegration()], - traces_sample_rate=1.0, - ) - - test_state = { - "messages": [ - MockMessage("Hello, can you help me?", name="user"), - MockMessage("Of course! How can I assist you?", name="assistant"), - ] - } - - pregel = MockPregelInstance("test_graph") - - expected_assistant_response = "I'll help you with that task!" - expected_tool_calls = [ - { - "id": "call_test_123", - "type": "function", - "function": {"name": "search_tool", "arguments": '{"query": "help"}'}, - } - ] - - def original_invoke(self, *args, **kwargs): - input_messages = args[0].get("messages", []) - new_messages = input_messages + [ - MockMessage( - content=expected_assistant_response, - name="assistant", - tool_calls=expected_tool_calls, - response_metadata={ - "token_usage": { - "total_tokens": 15, - "prompt_tokens": 10, - "completion_tokens": 5, - }, - }, - ), - MockMessage( - content=expected_assistant_response, - name="assistant", - tool_calls=expected_tool_calls, - response_metadata={ - "token_usage": { - "total_tokens": 35, - "prompt_tokens": 20, - "completion_tokens": 15, - }, - }, - ), - ] - return {"messages": new_messages} - - items = capture_items("span") - - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - result = wrapped_invoke(pregel, test_state) - - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span has aggregated usage from both API calls - # Total: 10 + 20 = 30 input tokens, 5 + 15 = 20 output tokens, 15 + 35 = 50 total - assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 30 - assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 50 - - -def test_pregel_ainvoke_multiple_llm_calls_aggregate_usage( - sentry_init, - capture_items, -): - """ - Test that invoke_agent spans show aggregated usage across multiple LLM calls - (e.g., when tools are used and multiple API calls are made). - """ - sentry_init( - integrations=[LanggraphIntegration()], - traces_sample_rate=1.0, - ) - - test_state = { - "messages": [ - MockMessage("Hello, can you help me?", name="user"), - MockMessage("Of course! How can I assist you?", name="assistant"), - ] - } - - pregel = MockPregelInstance("test_graph") - - expected_assistant_response = "I'll help you with that task!" - expected_tool_calls = [ - { - "id": "call_test_123", - "type": "function", - "function": {"name": "search_tool", "arguments": '{"query": "help"}'}, - } - ] - - async def original_ainvoke(self, *args, **kwargs): - input_messages = args[0].get("messages", []) - new_messages = input_messages + [ - MockMessage( - content=expected_assistant_response, - name="assistant", - tool_calls=expected_tool_calls, - response_metadata={ - "token_usage": { - "total_tokens": 15, - "prompt_tokens": 10, - "completion_tokens": 5, - }, - }, - ), - MockMessage( - content=expected_assistant_response, - name="assistant", - tool_calls=expected_tool_calls, - response_metadata={ - "token_usage": { - "total_tokens": 35, - "prompt_tokens": 20, - "completion_tokens": 15, - }, - }, - ), - ] - return {"messages": new_messages} - - async def run_test(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - result = await wrapped_ainvoke(pregel, test_state) - return result - - items = capture_items("span") - - result = asyncio.run(run_test()) - assert result is not None - - sentry_sdk.flush() - spans = [item.payload for item in items] - invoke_spans = [ - span - for span in spans - if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT - ] - assert len(invoke_spans) == 1 - invoke_agent_span = invoke_spans[0] - - # Verify invoke_agent span has aggregated usage from both API calls - # Total: 10 + 20 = 30 input tokens, 5 + 15 = 20 output tokens, 15 + 35 = 50 total - assert invoke_agent_span["attributes"]["gen_ai.usage.input_tokens"] == 30 - assert invoke_agent_span["attributes"]["gen_ai.usage.output_tokens"] == 20 - assert invoke_agent_span["attributes"]["gen_ai.usage.total_tokens"] == 50 - - def test_pregel_invoke_span_includes_response_model( sentry_init, capture_items, @@ -1646,142 +1310,3 @@ async def run_test(): else: assert SPANDATA.GEN_AI_RESPONSE_TEXT not in data assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in data - - -def test_pregel_invoke_message_delta_ignores_gen_ai_inputs_setting( - sentry_init, - capture_items, -): - sentry_init( - integrations=[LanggraphIntegration()], - traces_sample_rate=1.0, - _experiments={ - "data_collection": {"gen_ai": {"inputs": False, "outputs": True}} - }, - ) - - prior_response = "Of course! How can I assist you?" - test_state = { - "messages": [ - MockMessage("Hello, can you help me?", name="user"), - MockMessage( - prior_response, - name="assistant", - response_metadata={ - "token_usage": { - "total_tokens": 300, - "prompt_tokens": 100, - "completion_tokens": 200, - }, - "model_name": "gpt-3.5-turbo", - }, - ), - ] - } - pregel = MockPregelInstance("test_graph") - expected_assistant_response = "I'll help you with that task!" - - def original_invoke(self, *args, **kwargs): - return { - "messages": args[0].get("messages", []) - + [ - MockMessage( - content=expected_assistant_response, - name="assistant", - response_metadata={ - "token_usage": { - "total_tokens": 30, - "prompt_tokens": 10, - "completion_tokens": 20, - }, - "model_name": "gpt-4.1-2025-04-14", - }, - ) - ] - } - - captured = capture_items("span") - - wrapped_invoke = _wrap_pregel_invoke(original_invoke) - wrapped_invoke(pregel, test_state) - - data = _invoke_span_data(captured) - - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data - assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response - assert prior_response not in data[SPANDATA.GEN_AI_RESPONSE_TEXT] - assert data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - assert data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 - assert data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 - assert data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4.1-2025-04-14" - - -def test_pregel_ainvoke_message_delta_ignores_gen_ai_inputs_setting( - sentry_init, - capture_items, -): - sentry_init( - integrations=[LanggraphIntegration()], - traces_sample_rate=1.0, - _experiments={ - "data_collection": {"gen_ai": {"inputs": False, "outputs": True}} - }, - ) - - prior_response = "It is sunny in Berlin." - test_state = { - "messages": [ - MockMessage("What is the weather?", name="user"), - MockMessage( - prior_response, - name="assistant", - response_metadata={ - "token_usage": { - "total_tokens": 300, - "prompt_tokens": 100, - "completion_tokens": 200, - }, - "model_name": "gpt-3.5-turbo", - }, - ), - ] - } - pregel = MockPregelInstance("async_graph") - expected_assistant_response = "Let me check the weather for you!" - - async def original_ainvoke(self, *args, **kwargs): - return { - "messages": args[0].get("messages", []) - + [ - MockMessage( - content=expected_assistant_response, - name="assistant", - response_metadata={ - "token_usage": { - "total_tokens": 30, - "prompt_tokens": 10, - "completion_tokens": 20, - }, - "model_name": "gpt-4.1-2025-04-14", - }, - ) - ] - } - - async def run_test(): - wrapped_ainvoke = _wrap_pregel_ainvoke(original_ainvoke) - return await wrapped_ainvoke(pregel, test_state) - - captured = capture_items("span") - - asyncio.run(run_test()) - - data = _invoke_span_data(captured) - - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in data - assert data[SPANDATA.GEN_AI_RESPONSE_TEXT] == expected_assistant_response - assert prior_response not in data[SPANDATA.GEN_AI_RESPONSE_TEXT] - assert data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 - assert data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 - assert data[SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS] == 30 - assert data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4.1-2025-04-14"