NVIDIA-NeMo · stepwise-ai-dev · Jul 14, 2026 · Jul 22, 2026 · Jul 21, 2026 · Jul 21, 2026
@@ -15,6 +15,7 @@
 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,
+    ChatCompletionChoice,
     ChatCompletionRequest,
     ChatCompletionResponse,
     ToolCall,
@@ -124,7 +125,16 @@ def parse_anthropic_response(response_json: dict[str, Any]) -> ChatCompletionRes
         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)
+    choice = ChatCompletionChoice(
+        message=message,
+        finish_reason=response_json.get("stop_reason"),
+    )
+    return ChatCompletionResponse(
+        message=message,
+        usage=usage,
+        raw=response_json,
+        choices=[choice],
+    )
 
 
 def translate_request_messages(

@@ -5,6 +5,7 @@
 
 import logging
 from collections.abc import Callable
+from enum import Enum
 from functools import wraps
 from typing import Any, NoReturn
 
@@ -41,21 +42,29 @@ def get_exception_primary_cause(exception: BaseException) -> BaseException:
     return get_exception_primary_cause(exception.__cause__)
 
 
+class GenerationTruncationReason(str, Enum):
+    MAX_TOKENS = "max_tokens"
+    MODEL_CONTEXT_WINDOW_EXCEEDED = "model_context_window_exceeded"
+
+
 class GenerationValidationFailureError(Exception):
     summary: str
     detail: str | None
     failure_kind: str
+    truncation_reason: GenerationTruncationReason | None
 
     def __init__(
         self,
         summary: str,
         *,
         detail: str | None = None,
         failure_kind: str = "validation_error",
+        truncation_reason: GenerationTruncationReason | None = None,
     ) -> None:
         self.summary = summary.strip()
         self.detail = _normalize_error_detail(detail)
         self.failure_kind = failure_kind
+        self.truncation_reason = truncation_reason
 
         message = self.summary
         if self.detail is not None:
@@ -216,13 +225,18 @@ def handle_llm_exceptions(
         case GenerationValidationFailureError():
             detail_text = exception.detail.rstrip(".") if exception.detail is not None else None
             validation_detail = f" Validation detail: {detail_text}." if detail_text is not None else ""
+            solution = "This is most likely temporary as we make additional attempts. If you continue to see more of this, simplify or modify the output schema for structured output and try again. If you are attempting token-intensive tasks like generations with high-reasoning effort, ensure that max_tokens in the model config is high enough to reach completion."
+            if exception.truncation_reason == GenerationTruncationReason.MAX_TOKENS:
+                solution = "The response appears to have been cut off because it reached max_tokens, causing the parse failure. Increase inference_parameters.max_tokens in the model config and try again."
+            elif exception.truncation_reason == GenerationTruncationReason.MODEL_CONTEXT_WINDOW_EXCEEDED:
+                solution = "The response appears to have been cut off because the model context window was exhausted, causing the parse failure. Reduce the prompt, context, or schema size, or use a model with a larger context window and try again; increasing inference_parameters.max_tokens will not help."
             raise ModelGenerationValidationFailureError(
                 FormattedLLMErrorMessage(
                     cause=(
                         f"The model output from {model_name!r} could not be parsed into the requested format "
                         f"while {purpose}.{validation_detail}"
                     ),
-                    solution="This is most likely temporary as we make additional attempts. If you continue to see more of this, simplify or modify the output schema for structured output and try again. If you are attempting token-intensive tasks like generations with high-reasoning effort, ensure that max_tokens in the model config is high enough to reach completion.",
+                    solution=solution,
                 ),
                 detail=exception.detail,
                 failure_kind=exception.failure_kind,

@@ -18,6 +18,7 @@
 from data_designer.config.utils.media_helpers import is_image_diffusion_model
 from data_designer.engine.mcp.errors import MCPConfigurationError
 from data_designer.engine.model_provider import ModelProviderRegistry
+from data_designer.engine.models.clients.parsing import get_first_value_or_none, get_value_from
 from data_designer.engine.models.clients.types import (
     ChatCompletionRequest,
     ChatCompletionResponse,
@@ -28,6 +29,7 @@
     Usage,
 )
 from data_designer.engine.models.errors import (
+    GenerationTruncationReason,
     GenerationValidationFailureError,
     ImageGenerationError,
     acatch_llm_exceptions,
@@ -60,6 +62,42 @@ def _identity(x: Any) -> Any:
 logger = logging.getLogger(__name__)
 
 
+_TRUNCATION_REASON_BY_FINISH_REASON = {
+    "length": GenerationTruncationReason.MAX_TOKENS,
+    "max_tokens": GenerationTruncationReason.MAX_TOKENS,
+    "model_context_window_exceeded": GenerationTruncationReason.MODEL_CONTEXT_WINDOW_EXCEEDED,
+}
+
+
+def _classify_generation_truncation_reason(
+    response: ChatCompletionResponse,
+) -> GenerationTruncationReason | None:
+    first_choice = get_first_value_or_none(response.choices)
+    canonical_reason = get_value_from(first_choice, "finish_reason")
+    if isinstance(canonical_reason, str):
+        return _TRUNCATION_REASON_BY_FINISH_REASON.get(canonical_reason)
+
+    # Keep a raw fallback for custom or future adapters that have not populated
+    # the canonical choice finish reason yet.
+    raw_choices = get_value_from(response.raw, "choices")
+    raw_first_choice = get_first_value_or_none(raw_choices)
+    raw_reason = get_value_from(raw_first_choice, "finish_reason")
+    if not isinstance(raw_reason, str):
+        raw_reason = get_value_from(response.raw, "stop_reason")
+    if not isinstance(raw_reason, str):
+        return None
+    return _TRUNCATION_REASON_BY_FINISH_REASON.get(raw_reason)
+
+
+def _merge_generation_truncation_reason(
+    accumulated: GenerationTruncationReason | None,
+    current: GenerationTruncationReason | None,
+) -> GenerationTruncationReason | None:
+    if GenerationTruncationReason.MODEL_CONTEXT_WINDOW_EXCEEDED in (accumulated, current):
+        return GenerationTruncationReason.MODEL_CONTEXT_WINDOW_EXCEEDED
+    return current or accumulated
+
+
 def _classify_generation_failure_kind(exc: ParserException) -> str:
     detail = " ".join(str(get_exception_primary_cause(exc)).split()).lower()
     if "response_schema" in detail or "model_validate" in detail:
@@ -69,11 +107,17 @@ def _classify_generation_failure_kind(exc: ParserException) -> str:
     return "parse_error"
 
 
-def _build_generation_validation_error(summary: str, exc: ParserException) -> GenerationValidationFailureError:
+def _build_generation_validation_error(
+    summary: str,
+    exc: ParserException,
+    *,
+    truncation_reason: GenerationTruncationReason | None = None,
+) -> GenerationValidationFailureError:
     return GenerationValidationFailureError(
         summary,
         detail=str(get_exception_primary_cause(exc)),
         failure_kind=_classify_generation_failure_kind(exc),
+        truncation_reason=truncation_reason,
     )
 
 
@@ -349,6 +393,7 @@ def generate(
         # restart-or-raise. Reset to 0 on each conversation restart.
         parse_attempts = 0
         curr_num_restarts = 0
+        truncation_reason: GenerationTruncationReason | None = None
 
         mcp_facade = self._get_mcp_facade(tool_alias)
 
@@ -402,10 +447,15 @@ def generate(
                 output_obj = parser(response)  # type: ignore - if not a string will cause a ParserException below
                 break
             except ParserException as exc:
+                truncation_reason = _merge_generation_truncation_reason(
+                    truncation_reason,
+                    _classify_generation_truncation_reason(completion_response),
+                )
                 if max_correction_steps == 0 and max_conversation_restarts == 0:
                     raise _build_generation_validation_error(
                         "Unsuccessful generation attempt. No retries were attempted.",
                         exc,
+                        truncation_reason=truncation_reason,
                     ) from exc
 
                 if parse_attempts <= max_correction_steps:
@@ -425,6 +475,7 @@ def generate(
                             f"and {max_conversation_restarts} conversation restarts."
                         ),
                         exc,
+                        truncation_reason=truncation_reason,
                     ) from exc
 
         if not skip_usage_tracking and mcp_facade is not None:
@@ -457,6 +508,7 @@ async def agenerate(
         # See `generate` for a description of the parse-attempts counter semantics.
         parse_attempts = 0
         curr_num_restarts = 0
+        truncation_reason: GenerationTruncationReason | None = None
 
         mcp_facade = self._get_mcp_facade(tool_alias)
 
@@ -507,10 +559,15 @@ async def agenerate(
                 output_obj = parser(response)
                 break
             except ParserException as exc:
+                truncation_reason = _merge_generation_truncation_reason(
+                    truncation_reason,
+                    _classify_generation_truncation_reason(completion_response),
+                )
                 if max_correction_steps == 0 and max_conversation_restarts == 0:
                     raise _build_generation_validation_error(
                         "Unsuccessful generation attempt. No retries were attempted.",
                         exc,
+                        truncation_reason=truncation_reason,
                     ) from exc
 
                 if parse_attempts <= max_correction_steps:
@@ -529,6 +586,7 @@ async def agenerate(
                             f"and {max_conversation_restarts} conversation restarts."
                         ),
                         exc,
+                        truncation_reason=truncation_reason,
                     ) from exc
 
         if not skip_usage_tracking and mcp_facade is not None:

@@ -47,6 +47,8 @@
 from data_designer.engine.dataset_builders.utils.row_group_buffer import RowGroupBufferManager
 from data_designer.engine.models.errors import (
     RETRYABLE_MODEL_ERRORS,
+    FormattedLLMErrorMessage,
+    ModelGenerationValidationFailureError,
     ModelInternalServerError,
     ModelRateLimitError,
     ModelRequestAdmissionTimeoutError,
@@ -204,6 +206,22 @@ def generate(self, data: dict) -> dict:
         return data
 
 
+class MockMaxTokensValidationFailureGenerator(ColumnGenerator[ExpressionColumnConfig]):
+    """Cell generator that surfaces formatted max-token parse guidance."""
+
+    @staticmethod
+    def get_generation_strategy() -> GenerationStrategy:
+        return GenerationStrategy.CELL_BY_CELL
+
+    def generate(self, _data: dict) -> dict:
+        raise ModelGenerationValidationFailureError(
+            FormattedLLMErrorMessage(
+                cause="The model output could not be parsed because it reached max_tokens.",
+                solution="Increase inference_parameters.max_tokens in the model config and try again.",
+            )
+        )
+
+
 class MockBuggyGenerator(ColumnGenerator[ExpressionColumnConfig]):
     """Generator that raises a bare built-in exception from generator code."""
 
@@ -781,6 +799,32 @@ async def test_scheduler_non_retryable_failure_drops_row() -> None:
     assert tracker.is_row_group_complete(0, 2, ["seed", "fail_col"])
 
 
+@pytest.mark.asyncio(loop_scope="session")
+async def test_scheduler_characterization_logs_max_tokens_guidance_when_non_retryable_failure_drops_row(
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Characterize existing warning propagation for formatted public model errors."""
+    provider = _mock_provider()
+    scheduler, tracker = _build_simple_pipeline(
+        num_records=1,
+        generators={
+            "seed": MockSeedGenerator(config=_expr_config("seed"), resource_provider=provider),
+            "cell_out": MockMaxTokensValidationFailureGenerator(
+                config=_expr_config("cell_out"),
+                resource_provider=provider,
+            ),
+        },
+    )
+
+    with caplog.at_level(logging.WARNING, logger="data_designer.engine.dataset_builders.async_scheduler"):
+        await scheduler.run()
+
+    assert tracker.is_dropped(0, 0)
+    assert "Non-retryable failure on cell_out[rg=0, row=0]:" in caplog.text
+    assert "max_tokens" in caplog.text
+    assert "inference_parameters.max_tokens" in caplog.text
+
+
 def test_scheduler_internal_bug_classifier_preserves_scheduler_builtin_failures() -> None:
     scheduler, tracker = _build_simple_pipeline(num_records=1)
     assert scheduler._is_internal_bug(KeyError("missing internal key"))

@@ -269,6 +269,22 @@ def count_reasoning_text(text: str) -> int:
     assert json.loads(response.message.tool_calls[0].arguments_json) == {"query": "weather"}
 
 
+@pytest.mark.parametrize(
+    "stop_reason",
+    ["end_turn", "tool_use", "max_tokens", "model_context_window_exceeded"],
+)
+def test_parse_anthropic_response_maps_stop_reason_to_canonical_choice(stop_reason: str) -> None:
+    raw = {
+        "content": [{"type": "text", "text": "partial response"}],
+        "stop_reason": stop_reason,
+    }
+
+    response = parse_anthropic_response(raw)
+
+    assert response.choices[0].finish_reason == stop_reason
+    assert response.raw is raw
+
+
 def test_translate_non_tool_message_rejects_unsupported_role() -> None:
     with pytest.raises(ValueError, match="does not support message role"):
         translate_non_tool_message({"role": "system", "content": "Nope"})