Be more strict with health check responses

Signed-off-by: Mike Knepper <mknepper@nvidia.com>

Author: mikeknep

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

@@ -8,6 +8,8 @@
 
 from data_designer.config.models import GenerationType, ModelConfig
 from data_designer.engine.model_provider import ModelProvider, ModelProviderRegistry
+from data_designer.engine.models.errors import ModelGenerationValidationFailureError
+from data_designer.engine.models.parsers.errors import ParserException
 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 +29,21 @@
 logger = logging.getLogger(__name__)
 
 
+def _parse_health_check_chat_response(response: str) -> str:
+    if not isinstance(response, str) or not response:
+        raise ParserException("Health check response must be non-empty text.")
+    return response
+
+
+def _validate_health_check_embedding_response(vectors: list[list[float]], *, model_alias: str) -> None:
+    if not isinstance(vectors, list) or len(vectors) != 1 or not isinstance(vectors[0], list) or not vectors[0]:
+        raise ModelGenerationValidationFailureError(
+            f"Health check for model alias {model_alias!r} returned an invalid embedding response.",
+            detail="Expected exactly one non-empty embedding vector.",
+            failure_kind="validation_error",
+        )
+
+
 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)"
@@ -241,15 +258,16 @@ def run_health_check(self, model_aliases: list[str]) -> None:
             )
             try:
                 if model.model_generation_type == GenerationType.EMBEDDING:
-                    model.generate_text_embeddings(
+                    vectors = model.generate_text_embeddings(
                         input_texts=["Hello!"],
                         skip_usage_tracking=True,
                         purpose="running health checks",
                     )
+                    _validate_health_check_embedding_response(vectors, model_alias=model_alias)
                 elif model.model_generation_type == GenerationType.CHAT_COMPLETION:
                     model.generate(
                         prompt="Hello!",
-                        parser=lambda x: x,
+                        parser=_parse_health_check_chat_response,
                         system_prompt="You are a helpful assistant.",
                         max_correction_steps=0,
                         max_conversation_restarts=0,
@@ -286,15 +304,16 @@ async def arun_health_check(self, model_aliases: list[str]) -> None:
             )
             try:
                 if model.model_generation_type == GenerationType.EMBEDDING:
-                    await model.agenerate_text_embeddings(
+                    vectors = await model.agenerate_text_embeddings(
                         input_texts=["Hello!"],
                         skip_usage_tracking=True,
                         purpose="running health checks",
                     )
+                    _validate_health_check_embedding_response(vectors, model_alias=model_alias)
                 elif model.model_generation_type == GenerationType.CHAT_COMPLETION:
                     await model.agenerate(
                         prompt="Hello!",
-                        parser=lambda x: x,
+                        parser=_parse_health_check_chat_response,
                         system_prompt="You are a helpful assistant.",
                         max_correction_steps=0,
                         max_conversation_restarts=0,

packages/data-designer-engine/tests/engine/models/test_model_registry.py

@@ -7,11 +7,12 @@
 
 from data_designer.config.models import ChatCompletionInferenceParams, ModelConfig
 from data_designer.config.run_config import RequestAdmissionTuningConfig, RunConfig
-from data_designer.engine.models.errors import ModelAuthenticationError
+from data_designer.engine.models.errors import ModelAuthenticationError, ModelGenerationValidationFailureError
 from data_designer.engine.models.facade import ModelFacade
 from data_designer.engine.models.factory import create_model_registry
 from data_designer.engine.models.registry import ModelRegistry
 from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenCountSource, TokenUsageStats
+from data_designer.engine.testing import make_stub_completion_response
 from data_designer.logging import LOG_INDENT
 
 
@@ -332,6 +333,8 @@ def test_run_health_check_success(
     mock_generate_image: object,
     stub_model_registry: ModelRegistry,
 ) -> None:
+    mock_completion.return_value = make_stub_completion_response(content="Hello!")
+    mock_generate_text_embeddings.return_value = [[0.1]]
     model_aliases = ["stub-text", "stub-reasoning", "stub-embedding", "stub-image"]
     stub_model_registry.run_health_check(model_aliases)
     assert mock_completion.call_count == 2
@@ -365,6 +368,7 @@ def test_run_health_check_embedding_authentication_error(
     stub_model_registry: ModelRegistry,
 ) -> None:
     auth_error = ModelAuthenticationError("Invalid API key for embedding model")
+    mock_completion.return_value = make_stub_completion_response(content="Hello!")
     mock_generate_text_embeddings.side_effect = auth_error
     model_aliases = ["stub-text", "stub-reasoning", "stub-embedding"]
 
@@ -375,12 +379,39 @@ def test_run_health_check_embedding_authentication_error(
     mock_generate_text_embeddings.assert_called_once()
 
 
+@patch.object(ModelFacade, "completion", autospec=True)
+def test_run_health_check_rejects_empty_completion_response(
+    mock_completion: object,
+    stub_model_registry: ModelRegistry,
+) -> None:
+    mock_completion.return_value = make_stub_completion_response(content="")
+
+    with pytest.raises(ModelGenerationValidationFailureError, match="Health check response must be non-empty text"):
+        stub_model_registry.run_health_check(["stub-text"])
+
+    mock_completion.assert_called_once()
+
+
+@patch.object(ModelFacade, "generate_text_embeddings", autospec=True)
+def test_run_health_check_rejects_empty_embedding_vector(
+    mock_generate_text_embeddings: object,
+    stub_model_registry: ModelRegistry,
+) -> None:
+    mock_generate_text_embeddings.return_value = [[]]
+
+    with pytest.raises(ModelGenerationValidationFailureError, match="invalid embedding response"):
+        stub_model_registry.run_health_check(["stub-embedding"])
+
+    mock_generate_text_embeddings.assert_called_once()
+
+
 @patch.object(ModelFacade, "completion", autospec=True)
 def test_run_health_check_skip_health_check_flag(
     mock_completion: object,
     stub_secrets_resolver: object,
     stub_model_provider_registry: object,
 ) -> None:
+    mock_completion.return_value = make_stub_completion_response(content="Hello!")
     # Create model configs: one with skip_health_check=True, others with default (False)
     model_configs = [
         ModelConfig(
@@ -436,6 +467,7 @@ async def test_arun_health_check_success(
     mock_agenerate_image: AsyncMock,
     stub_model_registry: ModelRegistry,
 ) -> None:
+    mock_agenerate_text_embeddings.return_value = [[0.1]]
     model_aliases = ["stub-text", "stub-reasoning", "stub-embedding", "stub-image"]
     await stub_model_registry.arun_health_check(model_aliases)
     assert mock_agenerate.call_count == 2
@@ -461,6 +493,20 @@ async def test_arun_health_check_authentication_error(
     mock_agenerate_text_embeddings.assert_not_awaited()
 
 
+@patch.object(ModelFacade, "agenerate_text_embeddings", new_callable=AsyncMock)
+@pytest.mark.asyncio
+async def test_arun_health_check_rejects_empty_embedding_vector(
+    mock_agenerate_text_embeddings: AsyncMock,
+    stub_model_registry: ModelRegistry,
+) -> None:
+    mock_agenerate_text_embeddings.return_value = [[]]
+
+    with pytest.raises(ModelGenerationValidationFailureError, match="invalid embedding response"):
+        await stub_model_registry.arun_health_check(["stub-embedding"])
+
+    mock_agenerate_text_embeddings.assert_awaited_once()
+
+
 def test_get_aggregate_max_parallel_requests(stub_model_registry: ModelRegistry) -> None:
     """get_aggregate_max_parallel_requests returns the sum across all model configs."""
     total = stub_model_registry.get_aggregate_max_parallel_requests()