From 7e4a8a4d27415ae9f3c4ff2be66fa40d6ec8f1a8 Mon Sep 17 00:00:00 2001 From: stepwise-ai-dev Date: Tue, 14 Jul 2026 15:49:35 -0400 Subject: [PATCH 1/5] docs: plan #411 truncation guidance Signed-off-by: stepwise-ai-dev --- ...truncation-aware-parse-failure-guidance.md | 157 ++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 plans/411/truncation-aware-parse-failure-guidance.md diff --git a/plans/411/truncation-aware-parse-failure-guidance.md b/plans/411/truncation-aware-parse-failure-guidance.md new file mode 100644 index 000000000..718647e33 --- /dev/null +++ b/plans/411/truncation-aware-parse-failure-guidance.md @@ -0,0 +1,157 @@ +--- +date: 2026-07-13 +authors: + - stepwise-ai-dev +issue: https://github.com/NVIDIA-NeMo/DataDesigner/issues/411 +pull_request: https://github.com/NVIDIA-NeMo/DataDesigner/pull/754 +--- + +# Plan: Truncation-aware parse failure guidance + +## Problem + +When structured model output cannot be parsed, DataDesigner reports a generic validation failure. That message is incomplete when the response ended because the model exhausted its output-token budget or context window: the parser is operating on truncated content, and schema changes or retries alone may not fix the problem. + +PR #754 added useful guidance, but it was written before the legacy sync engine was removed. Its boolean truncation flag now crosses more layers than the current async error flow requires, and it only represents the final failed parse attempt. The current implementation also treats Anthropic `max_tokens` without distinguishing `model_context_window_exceeded`, even though the remedies differ. + +## Current state + +- OpenAI-compatible responses already normalize `choices[*].finish_reason` into `ChatCompletionChoice.finish_reason`. +- The Anthropic adapter preserves its response in `raw` but does not populate canonical choice finish reasons. +- `ModelFacade.generate()` and `agenerate()` own correction and conversation-restart loops. +- `GenerationValidationFailureError` is the internal validation boundary; `handle_llm_exceptions()` converts it into the public `ModelGenerationValidationFailureError`. +- The async scheduler already logs the formatted public exception before dropping the affected row. No scheduler-specific truncation metadata is needed. + +## Goals + +- Detect parse failures preceded by output-token or context-window termination. +- Preserve that reason across correction attempts and conversation restarts. +- Give users remediation specific to the detected reason. +- Normalize Anthropic termination metadata at the adapter boundary. +- Keep termination metadata on the internal validation error only. +- Let the existing async scheduler logging path surface the formatted guidance. +- Address the current maintainer feedback with a small current-main implementation. + +## Non-goals + +- Do not restore or modify the removed sync engine. +- Do not add truncation state to `ModelGenerationValidationFailureError`. +- Do not change retry counts, correction control flow, row-drop policy, or scheduler retryability. +- Do not build a general provider termination taxonomy beyond the two reasons needed by #411 and the cited Anthropic behavior. +- Do not refactor unrelated model-error messages or adapter response parsing. +- Do not add public APIs or configuration fields. + +## Design + +### Canonical adapter metadata + +`parse_anthropic_response()` will place Anthropic's string `stop_reason` into `choices[0].finish_reason`, matching the canonical response representation already used by OpenAI-compatible adapters. Existing message, usage, and raw response fields remain unchanged. + +Expected canonical values: + +| Provider response | Canonical finish reason | Internal reason | +|---|---|---| +| OpenAI `finish_reason="length"` | `length` | `MAX_TOKENS` | +| Anthropic `stop_reason="max_tokens"` | `max_tokens` | `MAX_TOKENS` | +| Anthropic `stop_reason="model_context_window_exceeded"` | `model_context_window_exceeded` | `MODEL_CONTEXT_WINDOW_EXCEEDED` | +| Normal stop/tool use | provider value | none | + +### Internal reason type + +Add a small `GenerationTruncationReason` enum in `models/errors.py` with: + +- `MAX_TOKENS` +- `MODEL_CONTEXT_WINDOW_EXCEEDED` + +`GenerationValidationFailureError` gains an optional `truncation_reason`. The public `ModelGenerationValidationFailureError` remains unchanged; it continues to expose only normalized detail and failure kind. + +This keeps provider-independent classification available to `handle_llm_exceptions()` without threading a boolean through the public error or the scheduler. + +### Detection and compatibility fallback + +`ModelFacade` will classify the canonical `choices[0].finish_reason` first. A short raw-response fallback remains for custom or future adapters that do not yet populate canonical choices. + +The fallback will reuse `get_value_from()` and `get_first_value_or_none()` from `models/clients/parsing.py`. It will not introduce duplicate response-access helpers. + +### Accumulation and precedence + +The facade keeps one local `GenerationTruncationReason | None` for each `generate()` or `agenerate()` call. Whenever parsing fails, it merges the current response's reason into the accumulated reason before deciding whether to correct, restart, or raise. + +Precedence is deterministic: + +1. `MODEL_CONTEXT_WINDOW_EXCEEDED` +2. `MAX_TOKENS` +3. no classified reason + +Context-window exhaustion wins because advising the user to increase `max_tokens` would be wrong after the model has already exhausted its total context window. A successful parse returns normally regardless of earlier failed attempts; accumulated state matters only if generation ultimately fails. + +### User-facing messages + +`handle_llm_exceptions()` formats the reason-specific message: + +- `MAX_TOKENS`: explain that the response was cut off at the output-token limit and recommend increasing `inference_parameters.max_tokens`. +- `MODEL_CONTEXT_WINDOW_EXCEEDED`: explain that the model exhausted its context window and recommend reducing prompt/context/schema size or selecting a model with a larger context window. +- no classified reason: preserve the existing generic validation guidance. + +The scheduler receives and logs the existing public exception normally. No `getattr()`, public boolean, or scheduler formatting branch is introduced. + +## Files + +Production changes: + +- `packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py` +- `packages/data-designer-engine/src/data_designer/engine/models/errors.py` +- `packages/data-designer-engine/src/data_designer/engine/models/facade.py` + +Focused tests: + +- `packages/data-designer-engine/tests/engine/models/clients/test_anthropic.py` +- `packages/data-designer-engine/tests/engine/models/test_facade.py` +- `packages/data-designer-engine/tests/engine/models/test_model_errors.py` +- `packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py` + +Explicitly excluded from the implementation: + +- `packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py` +- `packages/data-designer-engine/tests/engine/dataset_builders/test_dataset_builder.py` + +## Test strategy + +All new behavior is implemented test-first. + +1. Adapter normalization: + - Anthropic `end_turn`, `tool_use`, `max_tokens`, and `model_context_window_exceeded` become canonical finish reasons. +2. One-attempt behavior: + - sync and async parse failures classify canonical OpenAI and Anthropic reasons; + - raw fallback remains covered; + - unclassified stops preserve the generic message. +3. Accumulation: + - an early `max_tokens` response followed by an unclassified final failure still reports output-token truncation; + - correction and restart paths are both covered across sync/async tests; + - context-window exhaustion takes precedence when both reasons occur. +4. Error formatting: + - each typed reason produces its specific remediation; + - generic validation behavior and public error attributes remain unchanged. +5. Scheduler integration: + - the existing non-retryable dropped-row warning includes the already-formatted reason-specific message without scheduler production changes. + +Validation order: + +1. Run each new focused test and observe the expected failure before implementation. +2. Run the four touched test modules. +3. Run `make check-engine`. +4. Run `make check-all-fix` and review any automatic changes. +5. Run the repository-required full `make test` suite before requesting review. +6. Record whether E2E coverage is applicable; run applicable E2E checks or explicitly report why they were not run. + +## Delivery strategy + +Delivery proceeds in reviewable stages from current `main`: + +1. submit this plan on the existing PR and confirm the direction with maintainers; +2. add Anthropic normalization and its focused adapter tests; +3. add internal truncation classification, accumulation, formatting, and the scheduler regression; +4. run the focused tests and repository-required validation; +5. address review findings in additional commits that name the resolved concern. + +Production implementation remains paused until the plan has been submitted to the PR and its direction approved, as required by `CONTRIBUTING.md`. From 1c1ddb0d942cd5e699639e50305227e1901c4f4e Mon Sep 17 00:00:00 2001 From: schultzjack Date: Tue, 21 Jul 2026 16:15:45 -0400 Subject: [PATCH 2/5] fix: normalize Anthropic stop reasons Signed-off-by: schultzjack --- .../clients/adapters/anthropic_translation.py | 12 +++++++++++- .../models/clients/test_anthropic_translation.py | 16 ++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py b/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py index 5ba186add..1449cb1a5 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/clients/adapters/anthropic_translation.py @@ -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( diff --git a/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py b/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py index 450a32b0d..dfef87687 100644 --- a/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py +++ b/packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py @@ -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"}) From 197c462b089e1d9ae6efb0dd93ad34f0669f8fe0 Mon Sep 17 00:00:00 2001 From: schultzjack Date: Tue, 21 Jul 2026 16:22:10 -0400 Subject: [PATCH 3/5] fix: format truncation-specific parse guidance Signed-off-by: schultzjack --- .../src/data_designer/engine/models/errors.py | 16 +++++- .../tests/engine/models/test_model_errors.py | 52 +++++++++++++++++++ 2 files changed, 67 insertions(+), 1 deletion(-) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/errors.py b/packages/data-designer-engine/src/data_designer/engine/models/errors.py index c289d5c62..68b02932b 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/errors.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/errors.py @@ -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,10 +42,16 @@ 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, @@ -52,10 +59,12 @@ def __init__( *, 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, diff --git a/packages/data-designer-engine/tests/engine/models/test_model_errors.py b/packages/data-designer-engine/tests/engine/models/test_model_errors.py index 8873aceba..9646a2326 100644 --- a/packages/data-designer-engine/tests/engine/models/test_model_errors.py +++ b/packages/data-designer-engine/tests/engine/models/test_model_errors.py @@ -12,6 +12,7 @@ from data_designer.engine.models.clients.errors import ProviderError, ProviderErrorKind from data_designer.engine.models.errors import ( DataDesignerError, + GenerationTruncationReason, GenerationValidationFailureError, ModelAPIConnectionError, ModelAPIError, @@ -231,6 +232,57 @@ def test_handle_llm_exceptions( handle_llm_exceptions(exception, stub_model_name, stub_model_provider_name, stub_purpose) +def test_generation_validation_failure_error_stores_truncation_reason() -> None: + exception = GenerationValidationFailureError( + " Generation validation failure ", + detail=" invalid\nJSON ", + failure_kind="parse_error", + truncation_reason=GenerationTruncationReason.MAX_TOKENS, + ) + + assert exception.summary == "Generation validation failure" + assert exception.detail == "invalid JSON" + assert exception.failure_kind == "parse_error" + assert exception.truncation_reason is GenerationTruncationReason.MAX_TOKENS + + +@pytest.mark.parametrize( + ("reason", "expected_solution"), + [ + pytest.param( + GenerationTruncationReason.MAX_TOKENS, + "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.", + id="max-tokens", + ), + pytest.param( + GenerationTruncationReason.MODEL_CONTEXT_WINDOW_EXCEEDED, + "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.", + id="context-window", + ), + ], +) +def test_handle_llm_exceptions_formats_truncation_reason( + reason: GenerationTruncationReason, + expected_solution: str, +) -> None: + with pytest.raises(ModelGenerationValidationFailureError, match=re.escape(expected_solution)) as exc_info: + handle_llm_exceptions( + GenerationValidationFailureError( + "Generation validation failure", + detail="invalid JSON", + failure_kind="parse_error", + truncation_reason=reason, + ), + stub_model_name, + stub_model_provider_name, + stub_purpose, + ) + + assert exc_info.value.detail == "invalid JSON" + assert exc_info.value.failure_kind == "parse_error" + assert not hasattr(exc_info.value, "truncation_reason") + + def test_handle_llm_exceptions_preserves_generation_failure_kind() -> None: with pytest.raises(ModelGenerationValidationFailureError) as exc_info: handle_llm_exceptions( From 103a6adc1edf6af65a5e52dfda6e0c2254702441 Mon Sep 17 00:00:00 2001 From: schultzjack Date: Tue, 21 Jul 2026 16:32:46 -0400 Subject: [PATCH 4/5] fix: surface truncation-aware parse failures Signed-off-by: schultzjack --- .../src/data_designer/engine/models/facade.py | 58 ++++- .../dataset_builders/test_async_scheduler.py | 44 ++++ .../tests/engine/models/test_facade.py | 216 +++++++++++++++++- ...truncation-aware-parse-failure-guidance.md | 2 +- 4 files changed, 316 insertions(+), 4 deletions(-) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/facade.py b/packages/data-designer-engine/src/data_designer/engine/models/facade.py index d2c35255b..c6d38e620 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/facade.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/facade.py @@ -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,40 @@ 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) + + 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 +105,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 +391,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 +445,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 +473,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 +506,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 +557,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 +584,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: diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py b/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py index 8de0e4cee..b9fad7e13 100644 --- a/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py +++ b/packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py @@ -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")) diff --git a/packages/data-designer-engine/tests/engine/models/test_facade.py b/packages/data-designer-engine/tests/engine/models/test_facade.py index b851dade7..529badc4e 100644 --- a/packages/data-designer-engine/tests/engine/models/test_facade.py +++ b/packages/data-designer-engine/tests/engine/models/test_facade.py @@ -12,6 +12,7 @@ from data_designer.engine.models.clients.errors import ProviderError, ProviderErrorKind from data_designer.engine.models.clients.types import ( AssistantMessage, + ChatCompletionChoice, ChatCompletionRequest, ChatCompletionResponse, EmbeddingResponse, @@ -33,9 +34,51 @@ from data_designer.engine.testing import StubMCPFacade, StubMCPRegistry, make_stub_completion_response -def _make_response(content: str | None = None, **kwargs: Any) -> ChatCompletionResponse: +def _make_response( + content: str | None = None, + *, + finish_reason: str | None = None, + raw: Any | None = None, + **kwargs: Any, +) -> ChatCompletionResponse: """Shorthand for creating a ChatCompletionResponse in tests.""" - return make_stub_completion_response(content=content, **kwargs) + response = make_stub_completion_response(content=content, **kwargs) + first_choice: ChatCompletionChoice = response.choices[0] + first_choice.finish_reason = finish_reason + response.raw = raw + return response + + +_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." +) +_CONTEXT_WINDOW_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." +) +_GENERIC_VALIDATION_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." +) +_PARSE_FAILURE_DETAIL = "invalid JSON response" + + +def _failing_parser(_response: str) -> str: + raise ParserException(_PARSE_FAILURE_DETAIL) + + +def _assert_public_parse_failure( + error: ModelGenerationValidationFailureError, + *, + expected_solution: str, +) -> None: + assert f"Solution: {expected_solution}" in str(error) + assert error.detail == _PARSE_FAILURE_DETAIL + assert error.failure_kind == "parse_error" def _assert_no_multi_choice_request( @@ -261,6 +304,175 @@ def _failing_parser(response: str) -> str: assert exc_info.value.failure_kind == "schema_validation" +@pytest.mark.parametrize( + "finish_reason", + [ + pytest.param("length", id="openai-length"), + pytest.param("max_tokens", id="anthropic-max-tokens"), + ], +) +def test_generate_truncation_from_canonical_finish_reason_reports_max_tokens_solution( + stub_model_facade: ModelFacade, + stub_model_client: MagicMock, + finish_reason: str, +) -> None: + stub_model_client.completion.return_value = _make_response( + "truncated response", + finish_reason=finish_reason, + ) + + with pytest.raises(ModelGenerationValidationFailureError) as exc_info: + stub_model_facade.generate(prompt="test", parser=_failing_parser) + + _assert_public_parse_failure(exc_info.value, expected_solution=_MAX_TOKENS_SOLUTION) + + +@pytest.mark.asyncio +async def test_agenerate_truncation_from_canonical_finish_reason_reports_context_window_solution( + stub_model_facade: ModelFacade, + stub_model_client: MagicMock, +) -> None: + stub_model_client.acompletion = AsyncMock( + return_value=_make_response( + "truncated response", + finish_reason="model_context_window_exceeded", + ) + ) + + with pytest.raises(ModelGenerationValidationFailureError) as exc_info: + await stub_model_facade.agenerate(prompt="test", parser=_failing_parser) + + _assert_public_parse_failure(exc_info.value, expected_solution=_CONTEXT_WINDOW_SOLUTION) + + +@pytest.mark.parametrize( + "raw", + [ + pytest.param({"choices": [{"finish_reason": "length"}]}, id="openai"), + pytest.param({"stop_reason": "max_tokens"}, id="anthropic"), + ], +) +def test_generate_truncation_uses_raw_finish_reason_fallback( + stub_model_facade: ModelFacade, + stub_model_client: MagicMock, + raw: dict[str, Any], +) -> None: + stub_model_client.completion.return_value = _make_response("truncated response", raw=raw) + + with pytest.raises(ModelGenerationValidationFailureError) as exc_info: + stub_model_facade.generate(prompt="test", parser=_failing_parser) + + _assert_public_parse_failure(exc_info.value, expected_solution=_MAX_TOKENS_SOLUTION) + + +@pytest.mark.asyncio +async def test_agenerate_truncation_prefers_populated_canonical_reason_over_raw_fallback( + stub_model_facade: ModelFacade, + stub_model_client: MagicMock, +) -> None: + stub_model_client.acompletion = AsyncMock( + return_value=_make_response( + "complete but invalid response", + finish_reason="stop", + raw={ + "choices": [{"finish_reason": "length"}], + "stop_reason": "max_tokens", + }, + ) + ) + + with pytest.raises(ModelGenerationValidationFailureError) as exc_info: + await stub_model_facade.agenerate(prompt="test", parser=_failing_parser) + + _assert_public_parse_failure(exc_info.value, expected_solution=_GENERIC_VALIDATION_SOLUTION) + assert _MAX_TOKENS_SOLUTION not in str(exc_info.value) + + +def test_generate_truncation_accumulates_max_tokens_across_correction_attempts( + stub_model_facade: ModelFacade, + stub_model_client: MagicMock, +) -> None: + stub_model_client.completion.side_effect = [ + _make_response("first invalid response", finish_reason="max_tokens"), + _make_response("final invalid response", finish_reason="stop"), + ] + + with pytest.raises(ModelGenerationValidationFailureError) as exc_info: + stub_model_facade.generate( + prompt="test", + parser=_failing_parser, + max_correction_steps=1, + ) + + _assert_public_parse_failure(exc_info.value, expected_solution=_MAX_TOKENS_SOLUTION) + assert stub_model_client.completion.call_count == 2 + + +@pytest.mark.parametrize( + "finish_reasons", + [ + pytest.param( + ("model_context_window_exceeded", "max_tokens"), + id="context-window-then-max-tokens", + ), + pytest.param( + ("max_tokens", "model_context_window_exceeded"), + id="max-tokens-then-context-window", + ), + ], +) +@pytest.mark.asyncio +async def test_agenerate_truncation_context_window_wins_across_conversation_restart( + stub_model_facade: ModelFacade, + stub_model_client: MagicMock, + finish_reasons: tuple[str, str], +) -> None: + stub_model_client.acompletion = AsyncMock( + side_effect=[ + _make_response("first invalid response", finish_reason=finish_reasons[0]), + _make_response("final invalid response", finish_reason=finish_reasons[1]), + ] + ) + + with pytest.raises(ModelGenerationValidationFailureError) as exc_info: + await stub_model_facade.agenerate( + prompt="test", + parser=_failing_parser, + max_conversation_restarts=1, + ) + + _assert_public_parse_failure(exc_info.value, expected_solution=_CONTEXT_WINDOW_SOLUTION) + assert stub_model_client.acompletion.await_count == 2 + + +@pytest.mark.asyncio +async def test_agenerate_truncation_then_success_returns_normally( + stub_model_facade: ModelFacade, + stub_model_client: MagicMock, +) -> None: + stub_model_client.acompletion = AsyncMock( + side_effect=[ + _make_response("invalid response", finish_reason="max_tokens"), + _make_response("valid response", finish_reason="stop"), + ] + ) + + def _parse_valid_response(response: str) -> dict[str, str]: + if response != "valid response": + raise ParserException(_PARSE_FAILURE_DETAIL) + return {"parsed": response} + + result, messages = await stub_model_facade.agenerate( + prompt="test", + parser=_parse_valid_response, + max_correction_steps=1, + ) + + assert result == {"parsed": "valid response"} + assert messages[-1] == ChatMessage.as_assistant(content="valid response") + assert stub_model_client.acompletion.await_count == 2 + + @pytest.mark.parametrize( "raw_content,expected", [ diff --git a/plans/411/truncation-aware-parse-failure-guidance.md b/plans/411/truncation-aware-parse-failure-guidance.md index 718647e33..0fe887fda 100644 --- a/plans/411/truncation-aware-parse-failure-guidance.md +++ b/plans/411/truncation-aware-parse-failure-guidance.md @@ -105,7 +105,7 @@ Production changes: Focused tests: -- `packages/data-designer-engine/tests/engine/models/clients/test_anthropic.py` +- `packages/data-designer-engine/tests/engine/models/clients/test_anthropic_translation.py` - `packages/data-designer-engine/tests/engine/models/test_facade.py` - `packages/data-designer-engine/tests/engine/models/test_model_errors.py` - `packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py` From 7642a68ce3a032fadc94b07e5d29813a1f0c45ed Mon Sep 17 00:00:00 2001 From: schultzjack Date: Tue, 21 Jul 2026 21:06:46 -0400 Subject: [PATCH 5/5] fix: address #411 review follow-ups Signed-off-by: schultzjack --- .../src/data_designer/engine/models/facade.py | 2 ++ plans/411/truncation-aware-parse-failure-guidance.md | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/facade.py b/packages/data-designer-engine/src/data_designer/engine/models/facade.py index c6d38e620..f73a8dcc7 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/facade.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/facade.py @@ -77,6 +77,8 @@ def _classify_generation_truncation_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") diff --git a/plans/411/truncation-aware-parse-failure-guidance.md b/plans/411/truncation-aware-parse-failure-guidance.md index 0fe887fda..71da8ef09 100644 --- a/plans/411/truncation-aware-parse-failure-guidance.md +++ b/plans/411/truncation-aware-parse-failure-guidance.md @@ -154,4 +154,4 @@ Delivery proceeds in reviewable stages from current `main`: 4. run the focused tests and repository-required validation; 5. address review findings in additional commits that name the resolved concern. -Production implementation remains paused until the plan has been submitted to the PR and its direction approved, as required by `CONTRIBUTING.md`. +Implementation began only after Andrea approved this direction on 2026-07-15, satisfying the `CONTRIBUTING.md` plan gate.