feat: track reasoning token usage (#670)

* feat: track reasoning token usage

Capture provider-reported reasoning-token breakdowns alongside output tokens without changing output token totals. Carry the field through model usage aggregation and add coverage for parsing, facade tracking, and deltas.

Refs #665

* fix: show reasoning tokens in usage summary

Include reasoning token counts in the local model usage summary while preserving output and total token semantics. Telemetry remains unchanged.

Refs #665

* fix: estimate missing reasoning token counts

When providers return reasoning content without a numeric usage breakdown, estimate reasoning tokens from that content while preserving provider-reported output and total token counts.

Refs #665

* fix: track reasoning token count source

* fix: simplify reasoning token source

* fix: omit unknown reasoning tokens from logs

* refactor: clarify reasoning token count helpers

* test: move token counting tests

* fix: enforce reasoning token source

* fix: address reasoning usage review

Author: nabinchha

packages/data-designer-config/src/data_designer/lazy_heavy_imports.py

@@ -40,6 +40,7 @@
     "jsonschema": "jsonschema",
     "PIL": "PIL",
     "Image": "PIL.Image",
+    "tiktoken": "tiktoken",
 }
 
 

packages/data-designer-engine/src/data_designer/engine/analysis/utils/column_statistics_calculations.py

@@ -4,12 +4,9 @@
 from __future__ import annotations
 
 import logging
-from functools import lru_cache
 from numbers import Number
 from typing import TYPE_CHECKING, Any
 
-import tiktoken
-
 import data_designer.lazy_heavy_imports as lazy
 from data_designer.config.analysis.column_statistics import (
     CategoricalDistribution,
@@ -25,6 +22,7 @@
     RecordBasedPromptRenderer,
     create_response_recipe,
 )
+from data_designer.engine.utils.token_counting import count_text_tokens
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -38,12 +36,6 @@
 logger = logging.getLogger(__name__)
 
 
-@lru_cache(maxsize=1)
-def _get_tokenizer() -> tiktoken.Encoding:
-    """Lazily initialize tokenizer to avoid import-time side effects."""
-    return tiktoken.get_encoding("cl100k_base")
-
-
 def calculate_column_distribution(
     column_name: str, df: pd.DataFrame, distribution_type: ColumnDistributionType
 ) -> dict[str, CategoricalDistribution | NumericalDistribution | MissingValue | None]:
@@ -106,7 +98,6 @@ def calculate_input_token_stats(
     column_config: LLMTextColumnConfig, df: pd.DataFrame
 ) -> dict[str, float | MissingValue]:
     try:
-        tokenizer = _get_tokenizer()
         num_tokens = []
         num_samples = min(MAX_PROMPT_SAMPLE_SIZE, len(df))
         renderer = RecordBasedPromptRenderer(response_recipe=create_response_recipe(column_config))
@@ -118,7 +109,7 @@ def calculate_input_token_stats(
                 prompt_template=column_config.prompt, record=record, prompt_type=PromptType.USER_PROMPT
             )
             concatenated_prompt = str(system_prompt + "\n\n" + prompt)
-            num_tokens.append(len(tokenizer.encode(concatenated_prompt, disallowed_special=())))
+            num_tokens.append(count_text_tokens(concatenated_prompt))
     except Exception as e:
         logger.warning(f"{WARNING_PREFIX} failed to calculate input token stats for column {column_config.name!r}: {e}")
         return {
@@ -137,10 +128,7 @@ def calculate_output_token_stats(
     column_config: LLMTextColumnConfig, df: pd.DataFrame
 ) -> dict[str, float | MissingValue]:
     try:
-        tokenizer = _get_tokenizer()
-        tokens_per_record = df[column_config.name].apply(
-            lambda value: len(tokenizer.encode(str(value), disallowed_special=()))
-        )
+        tokens_per_record = df[column_config.name].apply(lambda value: count_text_tokens(str(value)))
         return {
             "output_tokens_mean": tokens_per_record.mean(),
             "output_tokens_median": tokens_per_record.median(),

packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py

@@ -7,7 +7,7 @@
 import re
 from typing import Any
 
-from data_designer.engine.models.clients.parsing import extract_usage
+from data_designer.engine.models.clients.parsing import extract_usage, fill_reasoning_token_count_from_content
 from data_designer.engine.models.clients.types import (
     AssistantMessage,
     ChatCompletionRequest,
@@ -100,6 +100,7 @@ def parse_anthropic_response(response_json: dict[str, Any]) -> ChatCompletionRes
     usage: Usage | None = None
     if raw_usage:
         usage = extract_usage(raw_usage)
+        usage = fill_reasoning_token_count_from_content(usage, message.reasoning_content)
 
     return ChatCompletionResponse(message=message, usage=usage, raw=response_json)
 

packages/data-designer-engine/src/data_designer/engine/models/clients/parsing.py

@@ -8,6 +8,7 @@
 import json
 import logging
 import uuid
+from dataclasses import replace
 from typing import Any
 
 from data_designer.config.utils.image_helpers import (
@@ -23,6 +24,8 @@
     ToolCall,
     Usage,
 )
+from data_designer.engine.models.usage import TokenCountSource
+from data_designer.engine.utils.token_counting import count_text_tokens
 
 logger = logging.getLogger(__name__)
 
@@ -44,6 +47,7 @@ def parse_chat_completion_response(response: Any) -> ChatCompletionResponse:
         images=images,
     )
     usage = extract_usage(get_value_from(response, "usage"), generated_images=len(images) if images else None)
+    usage = fill_reasoning_token_count_from_content(usage, assistant_message.reasoning_content)
     return ChatCompletionResponse(message=assistant_message, usage=usage, raw=response)
 
 
@@ -59,6 +63,7 @@ async def aparse_chat_completion_response(response: Any) -> ChatCompletionRespon
         images=images,
     )
     usage = extract_usage(get_value_from(response, "usage"), generated_images=len(images) if images else None)
+    usage = fill_reasoning_token_count_from_content(usage, assistant_message.reasoning_content)
     return ChatCompletionResponse(message=assistant_message, usage=usage, raw=response)
 
 
@@ -260,6 +265,7 @@ def extract_usage(raw_usage: Any, generated_images: int | None = None) -> Usage
     input_tokens = get_value_from(raw_usage, "prompt_tokens")
     output_tokens = get_value_from(raw_usage, "completion_tokens")
     total_tokens = get_value_from(raw_usage, "total_tokens")
+    reasoning_token_count = extract_reasoning_token_count(raw_usage)
 
     if input_tokens is None:
         input_tokens = get_value_from(raw_usage, "input_tokens")
@@ -269,6 +275,7 @@ def extract_usage(raw_usage: Any, generated_images: int | None = None) -> Usage
     input_tokens = coerce_to_int_or_none(input_tokens)
     output_tokens = coerce_to_int_or_none(output_tokens)
     total_tokens = coerce_to_int_or_none(total_tokens)
+    reasoning_token_count_source = TokenCountSource.PROVIDER if reasoning_token_count is not None else None
 
     if total_tokens is None and input_tokens is not None and output_tokens is not None:
         total_tokens = input_tokens + output_tokens
@@ -280,17 +287,60 @@ def extract_usage(raw_usage: Any, generated_images: int | None = None) -> Usage
 
     generated_images = coerce_to_int_or_none(generated_images)
 
-    if input_tokens is None and output_tokens is None and total_tokens is None and generated_images is None:
+    if (
+        input_tokens is None
+        and output_tokens is None
+        and total_tokens is None
+        and reasoning_token_count is None
+        and generated_images is None
+    ):
         return None
 
     return Usage(
         input_tokens=input_tokens,
         output_tokens=output_tokens,
         total_tokens=total_tokens,
+        reasoning_tokens=reasoning_token_count,
+        reasoning_token_count_source=reasoning_token_count_source,
         generated_images=generated_images,
     )
 
 
+def extract_reasoning_token_count(raw_usage: Any) -> int | None:
+    if raw_usage is None:
+        return None
+
+    top_level = get_value_from(raw_usage, "reasoning_tokens")
+    if top_level is not None:
+        return coerce_to_int_or_none(top_level)
+
+    for details_key in ("completion_tokens_details", "output_tokens_details"):
+        details = get_value_from(raw_usage, details_key)
+        reasoning_token_count = get_value_from(details, "reasoning_tokens")
+        if reasoning_token_count is not None:
+            return coerce_to_int_or_none(reasoning_token_count)
+
+    return None
+
+
+def fill_reasoning_token_count_from_content(usage: Usage | None, reasoning_content: str | None) -> Usage | None:
+    if usage is None:
+        return None
+    if usage.reasoning_tokens is not None or not reasoning_content:
+        return usage
+
+    try:
+        reasoning_token_count = count_text_tokens(reasoning_content)
+    except Exception:
+        logger.debug("Failed to estimate reasoning token count", exc_info=True)
+        return usage
+    return replace(
+        usage,
+        reasoning_tokens=reasoning_token_count,
+        reasoning_token_count_source=TokenCountSource.ESTIMATED,
+    )
+
+
 def extract_embedding_vector(item: Any) -> list[float]:
     value = get_value_from(item, "embedding")
     if isinstance(value, list):

packages/data-designer-engine/src/data_designer/engine/models/clients/types.py

@@ -6,6 +6,8 @@
 from dataclasses import dataclass, field, fields
 from typing import Any, ClassVar, Protocol
 
+from data_designer.engine.models.usage import TokenCountSource
+
 
 class HttpResponse(Protocol):
     """Structural type for HTTP response objects (httpx, requests, etc.)."""
@@ -21,8 +23,16 @@ class Usage:
     input_tokens: int | None = None
     output_tokens: int | None = None
     total_tokens: int | None = None
+    reasoning_tokens: int | None = None
+    reasoning_token_count_source: TokenCountSource | None = None
     generated_images: int | None = None
 
+    def __post_init__(self) -> None:
+        if self.reasoning_tokens is None and self.reasoning_token_count_source is not None:
+            raise ValueError("reasoning_token_count_source requires reasoning_tokens")
+        if self.reasoning_tokens is not None and self.reasoning_token_count_source is None:
+            raise ValueError("reasoning_tokens requires reasoning_token_count_source")
+
 
 @dataclass
 class ImagePayload:

packages/data-designer-engine/src/data_designer/engine/models/facade.py

@@ -36,7 +36,12 @@
 )
 from data_designer.engine.models.parsers.errors import ParserException
 from data_designer.engine.models.telemetry import TELEMETRY_ENABLED
-from data_designer.engine.models.usage import ImageUsageStats, ModelUsageStats, RequestUsageStats, TokenUsageStats
+from data_designer.engine.models.usage import (
+    ImageUsageStats,
+    ModelUsageStats,
+    RequestUsageStats,
+    TokenUsageStats,
+)
 from data_designer.engine.models.utils import ChatMessage, prompt_to_messages
 
 if TYPE_CHECKING:
@@ -814,6 +819,8 @@ def _track_usage(self, usage: Usage | None, *, is_request_successful: bool) -> N
             token_usage = TokenUsageStats(
                 input_tokens=usage.input_tokens,
                 output_tokens=usage.output_tokens or 0,
+                reasoning_tokens=usage.reasoning_tokens,
+                reasoning_token_count_source=usage.reasoning_token_count_source,
             )
 
         self._usage_stats.extend(

packages/data-designer-engine/src/data_designer/engine/models/registry.py

@@ -8,7 +8,7 @@
 
 from data_designer.config.models import GenerationType, ModelConfig
 from data_designer.engine.model_provider import ModelProvider, ModelProviderRegistry
-from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenUsageStats
+from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenCountSource, TokenUsageStats
 from data_designer.engine.secret_resolver import SecretResolver
 from data_designer.logging import LOG_INDENT
 
@@ -27,6 +27,18 @@
 logger = logging.getLogger(__name__)
 
 
+def format_reasoning_token_count(reasoning_token_count: int, source: TokenCountSource | str | None) -> str:
+    if source == TokenCountSource.ESTIMATED or source == TokenCountSource.ESTIMATED.value:
+        return f"{reasoning_token_count} (estimated)"
+    return str(reasoning_token_count)
+
+
+def get_token_count_delta(current: int | None, previous: int | None) -> int | None:
+    if current is None:
+        return None
+    return current - (previous or 0)
+
+
 class ModelRegistry:
     def __init__(
         self,
@@ -115,9 +127,17 @@ def log_model_usage(self, total_time_elapsed: float) -> None:
             output_tokens = token_usage["output_tokens"]
             total_tokens = token_usage["total_tokens"]
             tokens_per_second = stats["tokens_per_second"]
-            logger.info(
-                f"{LOG_INDENT}tokens: input={input_tokens}, output={output_tokens}, total={total_tokens}, tps={tokens_per_second}"
-            )
+            token_parts = [f"input={input_tokens}", f"output={output_tokens}"]
+            if (reasoning_token_count := token_usage.get("reasoning_tokens")) is not None:
+                formatted_reasoning_token_count = format_reasoning_token_count(
+                    reasoning_token_count,
+                    token_usage.get("reasoning_token_count_source"),
+                )
+                token_parts.append(f"reasoning={formatted_reasoning_token_count}")
+            token_parts.extend([f"total={total_tokens}", f"tps={tokens_per_second}"])
+            logger.info(f"{LOG_INDENT}tokens: {', '.join(token_parts)}")
+            if token_usage.get("reasoning_token_count_source") == TokenCountSource.ESTIMATED.value:
+                logger.info(f"{LOG_INDENT}reasoning token count estimated with tiktoken")
 
             request_usage = stats["request_usage"]
             successful_requests = request_usage["successful_requests"]
@@ -160,14 +180,31 @@ def get_usage_deltas(self, snapshot: dict[str, ModelUsageStats]) -> dict[str, Mo
             prev = snapshot.get(model_name)
             delta_input = current.token_usage.input_tokens - (prev.token_usage.input_tokens if prev else 0)
             delta_output = current.token_usage.output_tokens - (prev.token_usage.output_tokens if prev else 0)
+            delta_reasoning_token_count = get_token_count_delta(
+                current.token_usage.reasoning_tokens,
+                prev.token_usage.reasoning_tokens if prev else None,
+            )
             delta_successful = current.request_usage.successful_requests - (
                 prev.request_usage.successful_requests if prev else 0
             )
             delta_failed = current.request_usage.failed_requests - (prev.request_usage.failed_requests if prev else 0)
 
-            if delta_input > 0 or delta_output > 0 or delta_successful > 0 or delta_failed > 0:
+            if (
+                delta_input > 0
+                or delta_output > 0
+                or (delta_reasoning_token_count is not None and delta_reasoning_token_count > 0)
+                or delta_successful > 0
+                or delta_failed > 0
+            ):
                 deltas[model_name] = ModelUsageStats(
-                    token_usage=TokenUsageStats(input_tokens=delta_input, output_tokens=delta_output),
+                    token_usage=TokenUsageStats(
+                        input_tokens=delta_input,
+                        output_tokens=delta_output,
+                        reasoning_tokens=delta_reasoning_token_count,
+                        reasoning_token_count_source=current.token_usage.reasoning_token_count_source
+                        if delta_reasoning_token_count is not None
+                        else None,
+                    ),
                     request_usage=RequestUsageStats(successful_requests=delta_successful, failed_requests=delta_failed),
                 )
         return deltas