From 128451fe2f4f5eecaeafd9a6ae7cb7f12b26b76c Mon Sep 17 00:00:00 2001 From: "Eric W. Tramel" Date: Tue, 13 Jan 2026 14:46:01 -0500 Subject: [PATCH] Control retry logic with RunConfig --- src/data_designer/config/run_config.py | 7 +++++++ .../generators/llm_completion.py | 12 ++---------- .../test_llm_completion_generators.py | 17 +++++++++++++---- tests/interface/test_data_designer.py | 10 ++++++++++ 4 files changed, 32 insertions(+), 14 deletions(-) diff --git a/src/data_designer/config/run_config.py b/src/data_designer/config/run_config.py index 8c8d3c70a..9ec1f123f 100644 --- a/src/data_designer/config/run_config.py +++ b/src/data_designer/config/run_config.py @@ -21,11 +21,18 @@ class RunConfig(ConfigBase): early shutdown is enabled. Default is 0.5. shutdown_error_window: Minimum number of completed tasks before error rate monitoring begins. Must be >= 0. Default is 10. + max_conversation_restarts: Maximum number of full conversation restarts permitted when + generation tasks call `ModelFacade.generate(...)`. Must be >= 0. Default is 5. + max_conversation_correction_steps: Maximum number of correction rounds permitted within a + single conversation when generation tasks call `ModelFacade.generate(...)`. Must be >= 0. + Default is 0. """ disable_early_shutdown: bool = False shutdown_error_rate: float = Field(default=0.5, ge=0.0, le=1.0) shutdown_error_window: int = Field(default=10, ge=0) + max_conversation_restarts: int = Field(default=5, ge=0) + max_conversation_correction_steps: int = Field(default=0, ge=0) @model_validator(mode="after") def normalize_shutdown_settings(self) -> Self: diff --git a/src/data_designer/engine/column_generators/generators/llm_completion.py b/src/data_designer/engine/column_generators/generators/llm_completion.py index 97bf2ecc6..8de6b8a52 100644 --- a/src/data_designer/engine/column_generators/generators/llm_completion.py +++ b/src/data_designer/engine/column_generators/generators/llm_completion.py @@ -28,10 +28,6 @@ logger = logging.getLogger(__name__) -DEFAULT_MAX_CONVERSATION_RESTARTS = 5 -DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS = 0 - - class ColumnGeneratorWithModelChatCompletion(ColumnGeneratorWithModel[TaskConfigT]): @functools.cached_property def response_recipe(self) -> ResponseRecipe: @@ -39,11 +35,11 @@ def response_recipe(self) -> ResponseRecipe: @property def max_conversation_correction_steps(self) -> int: - return DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS + return self.resource_provider.run_config.max_conversation_correction_steps @property def max_conversation_restarts(self) -> int: - return DEFAULT_MAX_CONVERSATION_RESTARTS + return self.resource_provider.run_config.max_conversation_restarts @functools.cached_property def prompt_renderer(self) -> RecordBasedPromptRenderer: @@ -129,7 +125,3 @@ def metadata() -> GeneratorMetadata: description="Judge a new dataset cell based on a set of rubrics", generation_strategy=GenerationStrategy.CELL_BY_CELL, ) - - @property - def max_conversation_restarts(self) -> int: - return 2 * DEFAULT_MAX_CONVERSATION_RESTARTS diff --git a/tests/engine/column_generators/generators/test_llm_completion_generators.py b/tests/engine/column_generators/generators/test_llm_completion_generators.py index 317d329f9..2da52d6c4 100644 --- a/tests/engine/column_generators/generators/test_llm_completion_generators.py +++ b/tests/engine/column_generators/generators/test_llm_completion_generators.py @@ -11,9 +11,8 @@ LLMStructuredColumnConfig, LLMTextColumnConfig, ) +from data_designer.config.run_config import RunConfig from data_designer.engine.column_generators.generators.llm_completion import ( - DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS, - DEFAULT_MAX_CONVERSATION_RESTARTS, REASONING_TRACE_COLUMN_POSTFIX, LLMCodeCellGenerator, LLMJudgeCellGenerator, @@ -34,6 +33,10 @@ def _create_generator_with_mocks(config_class=LLMTextColumnConfig, **config_kwar mock_provider = Mock() mock_resource_provider.model_registry = mock_model_registry + mock_resource_provider.run_config = RunConfig( + max_conversation_restarts=7, + max_conversation_correction_steps=2, + ) mock_model_registry.get_model.return_value = mock_model mock_model_registry.get_model_config.return_value = mock_model_config mock_model_registry.get_model_provider.return_value = mock_provider @@ -80,6 +83,8 @@ def test_generate_method(): assert mock_prompt_renderer.render.call_count == 2 mock_model.generate.assert_called_once() + assert mock_model.generate.call_args[1]["max_correction_steps"] == 2 + assert mock_model.generate.call_args[1]["max_conversation_restarts"] == 7 assert result["test_column"] == {"result": "test_output"} assert "test_column" + REASONING_TRACE_COLUMN_POSTFIX not in result @@ -202,6 +207,10 @@ def test_judge_generator_max_conversation_restarts_override(): mock_inference_params = Mock() mock_resource_provider.model_registry = mock_model_registry + mock_resource_provider.run_config = RunConfig( + max_conversation_restarts=7, + max_conversation_correction_steps=2, + ) mock_model_registry.get_model.return_value = mock_model mock_model_registry.get_model_config.return_value = mock_model_config mock_model_config.inference_parameters = mock_inference_params @@ -215,8 +224,8 @@ def test_judge_generator_max_conversation_restarts_override(): generator = LLMJudgeCellGenerator(config=config, resource_provider=mock_resource_provider) - assert generator.max_conversation_restarts == 2 * DEFAULT_MAX_CONVERSATION_RESTARTS - assert generator.max_conversation_correction_steps == DEFAULT_MAX_CONVERSATION_CORRECTION_STEPS + assert generator.max_conversation_restarts == 7 + assert generator.max_conversation_correction_steps == 2 @pytest.mark.parametrize( diff --git a/tests/interface/test_data_designer.py b/tests/interface/test_data_designer.py index 83f4b8b88..09aeb8c1f 100644 --- a/tests/interface/test_data_designer.py +++ b/tests/interface/test_data_designer.py @@ -112,6 +112,8 @@ def test_run_config_setting_persists(stub_artifact_path, stub_model_providers): assert data_designer._run_config.disable_early_shutdown is False assert data_designer._run_config.shutdown_error_rate == 0.5 assert data_designer._run_config.shutdown_error_window == 10 + assert data_designer._run_config.max_conversation_restarts == 5 + assert data_designer._run_config.max_conversation_correction_steps == 0 # Test setting custom values data_designer.set_run_config( @@ -119,11 +121,15 @@ def test_run_config_setting_persists(stub_artifact_path, stub_model_providers): disable_early_shutdown=True, shutdown_error_rate=0.8, shutdown_error_window=25, + max_conversation_restarts=7, + max_conversation_correction_steps=2, ) ) assert data_designer._run_config.disable_early_shutdown is True assert data_designer._run_config.shutdown_error_rate == 1.0 # normalized when disabled assert data_designer._run_config.shutdown_error_window == 25 + assert data_designer._run_config.max_conversation_restarts == 7 + assert data_designer._run_config.max_conversation_correction_steps == 2 # Test updating values data_designer.set_run_config( @@ -131,11 +137,15 @@ def test_run_config_setting_persists(stub_artifact_path, stub_model_providers): disable_early_shutdown=False, shutdown_error_rate=0.3, shutdown_error_window=5, + max_conversation_restarts=9, + max_conversation_correction_steps=1, ) ) assert data_designer._run_config.disable_early_shutdown is False assert data_designer._run_config.shutdown_error_rate == 0.3 assert data_designer._run_config.shutdown_error_window == 5 + assert data_designer._run_config.max_conversation_restarts == 9 + assert data_designer._run_config.max_conversation_correction_steps == 1 def test_run_config_normalizes_error_rate_when_disabled(stub_artifact_path, stub_model_providers):