Skip to content

Commit a748a71

Browse files
committed
feat(output_tokens_details): include also output_tokens_details
1 parent 484ee01 commit a748a71

File tree

2 files changed

+30
-10
lines changed

2 files changed

+30
-10
lines changed

packages/nemo-evaluator/src/nemo_evaluator/adapters/interceptors/reasoning_interceptor.py

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -340,19 +340,30 @@ def _process_reasoning_message(
340340
else:
341341
reasoning_finished = True
342342
if usage:
343+
# First try to get reasoning_tokens directly from usage
343344
reasoning_tokens = usage.get("reasoning_tokens", "unknown")
344345
updated_content_tokens = usage.get("content_tokens", "unknown")
345346

346-
# Check if reasoning_tokens is in completion_tokens_details
347-
if (
348-
reasoning_tokens == "unknown"
349-
and "completion_tokens_details" in usage
350-
):
351-
completion_details = usage["completion_tokens_details"]
352-
if isinstance(completion_details, dict):
353-
reasoning_tokens = completion_details.get(
354-
"reasoning_tokens", "unknown"
355-
)
347+
# If not found, check in completion_tokens_details and output_tokens_details
348+
if reasoning_tokens == "unknown":
349+
for key in ["completion_tokens_details", "output_tokens_details"]:
350+
if key in usage:
351+
details = usage[key]
352+
if isinstance(details, dict):
353+
reasoning_tokens = details.get(
354+
"reasoning_tokens", "unknown"
355+
)
356+
if reasoning_tokens != "unknown":
357+
self.logger.debug(
358+
f"Found reasoning_tokens in {key}: {reasoning_tokens}"
359+
)
360+
break
361+
362+
# Log if reasoning tokens were found
363+
if reasoning_tokens != "unknown":
364+
self.logger.debug(f"Reasoning tokens extracted: {reasoning_tokens}")
365+
else:
366+
self.logger.debug("No reasoning tokens found in usage data")
356367

357368
else:
358369
reasoning_finished = False

packages/nemo-evaluator/tests/unit_tests/adapters/interceptors/test_reasoning.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1004,6 +1004,15 @@ def test_cached_response_reasoning_behavior(
10041004
20,
10051005
40,
10061006
),
1007+
# Format 3: reasoning_tokens in output_tokens_details
1008+
(
1009+
{
1010+
"output_tokens_details": {"reasoning_tokens": 25},
1011+
"content_tokens": 50,
1012+
},
1013+
25,
1014+
50,
1015+
),
10071016
],
10081017
)
10091018
def test_reasoning_tokens_different_formats(

0 commit comments

Comments
 (0)