Skip to content

Commit de6ec8e

Browse files
alexander-alderman-webbalexander-alderman-webb
committed
feat(langchain): Record reasoning token usage attribute
1 parent b1c3907 commit de6ec8e

3 files changed

Lines changed: 70 additions & 15 deletions

File tree

‎sentry_sdk/consts.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -814,6 +814,12 @@ class SPANDATA:
814814
Example: 100
815815
"""
816816

817+
GEN_AI_USAGE_REASONING_OUTPUT_TOKENS = "gen_ai.usage.reasoning.output_tokens"
818+
"""
819+
The number of tokens used for reasoning to create the AI output.
820+
Example: 75
821+
"""
822+
817823
GEN_AI_USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
818824
"""
819825
The number of tokens in the output.

‎sentry_sdk/integrations/langchain.py‎

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ class TokenUsage(NamedTuple):
6969
total_tokens: "Optional[int]"
7070
cache_read: "Optional[int]"
7171
cache_creation: "Optional[int]"
72+
reasoning: "Optional[int]"
7273

7374

7475
try:
@@ -735,6 +736,7 @@ def _extract_tokens_from_generations(
735736
total_total = 0
736737
total_cache_read = None
737738
total_cache_creation = None
739+
reasoning = 0
738740

739741
for gen_list in generations:
740742
if not gen_list:
@@ -761,25 +763,31 @@ def _extract_tokens_from_generations(
761763
continue
762764

763765
input_token_details = usage_metadata.get("input_token_details")
764-
if not isinstance(input_token_details, dict):
765-
continue
766+
if isinstance(input_token_details, dict):
767+
if isinstance(input_token_details.get("cache_read"), int):
768+
total_cache_read = (total_cache_read or 0) + input_token_details[
769+
"cache_read"
770+
]
766771

767-
if isinstance(input_token_details.get("cache_read"), int):
768-
total_cache_read = (total_cache_read or 0) + input_token_details[
769-
"cache_read"
770-
]
772+
if isinstance(input_token_details.get("cache_creation"), int):
773+
total_cache_creation = (
774+
total_cache_creation or 0
775+
) + input_token_details["cache_creation"]
771776

772-
if isinstance(input_token_details.get("cache_creation"), int):
773-
total_cache_creation = (total_cache_creation or 0) + input_token_details[
774-
"cache_creation"
775-
]
777+
output_token_details = usage_metadata.get("output_token_details")
778+
if not isinstance(output_token_details, dict):
779+
continue
780+
781+
if isinstance(output_token_details.get("reasoning"), int):
782+
reasoning = (reasoning or 0) + output_token_details["reasoning"]
776783

777784
return TokenUsage(
778785
total_input if total_input > 0 else None,
779786
total_output if total_output > 0 else None,
780787
total_total if total_total > 0 else None,
781788
total_cache_read,
782789
total_cache_creation,
790+
reasoning,
783791
)
784792

785793

@@ -819,6 +827,7 @@ def _record_token_usage(
819827
total_tokens = None
820828
cache_read_tokens = None
821829
cache_creation_tokens = None
830+
reasoning = None
822831

823832
# Legacy that reads provider-specific token information.
824833
token_usage = _get_token_usage(response)
@@ -838,6 +847,8 @@ def _record_token_usage(
838847
cache_read_tokens = token_usage.cache_read
839848
if token_usage.cache_creation is not None:
840849
cache_creation_tokens = token_usage.cache_creation
850+
if token_usage.reasoning is not None:
851+
reasoning = token_usage.reasoning
841852

842853
set_on_span = (
843854
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
@@ -860,6 +871,9 @@ def _record_token_usage(
860871
SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, cache_creation_tokens
861872
)
862873

874+
if reasoning is not None:
875+
set_on_span(SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, reasoning)
876+
863877

864878
def _get_request_data(
865879
obj: "Any", args: "Any", kwargs: "Any"

‎tests/integrations/langchain/test_langchain.py‎

Lines changed: 40 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -271,6 +271,7 @@ def nonstreaming_multi_candidate_google_genai_model_response():
271271
cached_content_token_count=4,
272272
prompt_token_count=10,
273273
candidates_token_count=20,
274+
thoughts_token_count=5,
274275
total_token_count=30,
275276
),
276277
)
@@ -591,7 +592,7 @@ def test_langchain_multi_choice_response(
591592
assert len(chat_spans) == 1
592593

593594
assert chat_spans[0]["attributes"]["gen_ai.usage.input_tokens"] == 10
594-
assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens"] == 20
595+
assert chat_spans[0]["attributes"]["gen_ai.usage.output_tokens"] == 25
595596
assert chat_spans[0]["attributes"]["gen_ai.usage.total_tokens"] == 30
596597

597598
assert (
@@ -616,7 +617,7 @@ def test_langchain_multi_choice_response(
616617
assert len(chat_spans) == 1
617618

618619
assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 10
619-
assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 20
620+
assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 25
620621
assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 30
621622

622623
assert chat_spans[0]["data"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] == 4
@@ -817,6 +818,11 @@ def test_langchain_create_agent(
817818
== 6
818819
)
819820

821+
assert (
822+
chat_spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]
823+
== 5
824+
)
825+
820826
if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):
821827
assert (
822828
chat_spans[0]["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4"
@@ -904,6 +910,11 @@ def test_langchain_create_agent(
904910
== 6
905911
)
906912

913+
assert (
914+
chat_spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]
915+
== 5
916+
)
917+
907918
if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):
908919
assert (
909920
chat_spans[0]["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4"
@@ -970,7 +981,7 @@ def test_langchain_create_agent(
970981
assert chat_spans[0]["data"]["gen_ai.agent.name"] == "word_length_agent"
971982

972983
assert chat_spans[0]["data"]["gen_ai.usage.input_tokens"] == 10
973-
assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 20
984+
assert chat_spans[0]["data"]["gen_ai.usage.output_tokens"] == 25
974985
assert chat_spans[0]["data"]["gen_ai.usage.total_tokens"] == 30
975986

976987
assert chat_spans[0]["data"][SPANDATA.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] == 4
@@ -979,6 +990,8 @@ def test_langchain_create_agent(
979990
== 6
980991
)
981992

993+
assert chat_spans[0]["data"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS] == 5
994+
982995
if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):
983996
assert chat_spans[0]["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-4"
984997

@@ -1054,7 +1067,7 @@ def test_tool_execution_span(
10541067
),
10551068
output_tokens=50,
10561069
output_tokens_details=OutputTokensDetails(
1057-
reasoning_tokens=0,
1070+
reasoning_tokens=10,
10581071
),
10591072
total_tokens=192,
10601073
),
@@ -1066,7 +1079,7 @@ def test_tool_execution_span(
10661079
),
10671080
output_tokens=28,
10681081
output_tokens_details=OutputTokensDetails(
1069-
reasoning_tokens=0,
1082+
reasoning_tokens=11,
10701083
),
10711084
total_tokens=117,
10721085
),
@@ -1153,6 +1166,10 @@ def test_tool_execution_span(
11531166
]
11541167
== 31
11551168
)
1169+
assert (
1170+
chat_spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]
1171+
== 10
1172+
)
11561173
assert chat_spans[0]["attributes"]["gen_ai.system"] == "openai-chat"
11571174

11581175
assert chat_spans[1]["attributes"]["gen_ai.usage.input_tokens"] == 89
@@ -1168,6 +1185,10 @@ def test_tool_execution_span(
11681185
]
11691186
== 10
11701187
)
1188+
assert (
1189+
chat_spans[1]["attributes"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]
1190+
== 11
1191+
)
11711192
assert chat_spans[1]["attributes"]["gen_ai.system"] == "openai-chat"
11721193

11731194
if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):
@@ -1286,6 +1307,10 @@ def test_tool_execution_span(
12861307
]
12871308
== 31
12881309
)
1310+
assert (
1311+
chat_spans[0]["attributes"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]
1312+
== 10
1313+
)
12891314
assert chat_spans[0]["attributes"]["gen_ai.system"] == "openai-chat"
12901315

12911316
assert chat_spans[1]["attributes"]["gen_ai.usage.input_tokens"] == 89
@@ -1301,6 +1326,10 @@ def test_tool_execution_span(
13011326
]
13021327
== 10
13031328
)
1329+
assert (
1330+
chat_spans[1]["attributes"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS]
1331+
== 11
1332+
)
13041333
assert chat_spans[1]["attributes"]["gen_ai.system"] == "openai-chat"
13051334

13061335
if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):
@@ -1416,6 +1445,9 @@ def test_tool_execution_span(
14161445
chat_spans[0]["data"][SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]
14171446
== 31
14181447
)
1448+
assert (
1449+
chat_spans[0]["data"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS] == 10
1450+
)
14191451
assert chat_spans[0]["data"]["gen_ai.system"] == "openai-chat"
14201452

14211453
assert chat_spans[1]["data"]["gen_ai.usage.input_tokens"] == 89
@@ -1428,6 +1460,9 @@ def test_tool_execution_span(
14281460
chat_spans[1]["data"][SPANDATA.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS]
14291461
== 10
14301462
)
1463+
assert (
1464+
chat_spans[1]["data"][SPANDATA.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS] == 11
1465+
)
14311466
assert chat_spans[1]["data"]["gen_ai.system"] == "openai-chat"
14321467

14331468
if LANGCHAIN_OPENAI_VERSION >= (0, 3, 13):

0 commit comments

Comments
 (0)