retry schema validation error

Author: chiang-daniel

libs/core/kiln_ai/adapters/eval/eval_runner.py

@@ -263,21 +263,38 @@ async def run_job(self, job: EvalJob) -> bool:
             eval_run.save_to_file()
 
             return True
-        except (
-            litellm.RateLimitError,
-            litellm.APIConnectionError,
-            litellm.InternalServerError,
-            litellm.ServiceUnavailableError,
-            litellm.BadGatewayError,
-        ) as e:
-            logger.error(
-                f"Transient error running eval job for dataset item {job.item.id}: {e}",
-                exc_info=True,
-            )
-            raise RetryableError(str(e)) from e
         except Exception as e:
+            if _is_retryable_error(e):
+                logger.error(
+                    f"Transient error running eval job for dataset item {job.item.id}: {e}",
+                    exc_info=True,
+                )
+                raise RetryableError(str(e)) from e
             logger.error(
                 f"Error running eval job for dataset item {job.item.id}: {e}",
                 exc_info=True,
             )
             raise
+
+
+def _is_retryable_error(e: BaseException) -> bool:
+    if isinstance(
+        e,
+        (
+            litellm.RateLimitError,
+            litellm.APIConnectionError,
+            litellm.InternalServerError,
+            litellm.ServiceUnavailableError,
+            litellm.BadGatewayError,
+            litellm.JSONSchemaValidationError,
+        ),
+    ):
+        return True
+
+    # ValueError thrown by Kiln's adapter when structured output doesn't match schema
+    if isinstance(
+        e, ValueError
+    ) and "This task requires a specific output schema" in str(e):
+        return True
+
+    return False

libs/core/kiln_ai/adapters/eval/test_eval_runner.py

@@ -1,10 +1,11 @@
 from typing import Dict
 from unittest.mock import AsyncMock, patch
 
+import litellm
 import pytest
 
 from kiln_ai.adapters.eval.base_eval import BaseEval
-from kiln_ai.adapters.eval.eval_runner import EvalJob, EvalRunner
+from kiln_ai.adapters.eval.eval_runner import EvalJob, EvalRunner, _is_retryable_error
 from kiln_ai.adapters.ml_model_list import ModelProviderName
 from kiln_ai.datamodel import (
     DataSource,
@@ -831,3 +832,34 @@ async def run_task_and_eval(self, eval_job_item: TaskRun):
     # For full_trace evals, None trace should fail and not save a run
     eval_runs = mock_eval_config.runs()
     assert len(eval_runs) == 0
+
+
+@pytest.mark.parametrize(
+    "error",
+    [
+        litellm.RateLimitError("rate limited", "provider", "model", None),
+        litellm.APIConnectionError("connection failed", "provider", "model", None),
+        litellm.InternalServerError("server error", "provider", "model", None),
+        litellm.ServiceUnavailableError("unavailable", "provider", "model", None),
+        litellm.BadGatewayError("bad gateway", "provider", "model", None),
+        litellm.JSONSchemaValidationError("schema error", "provider", "model", None),
+        ValueError(
+            "This task requires a specific output schema. While the model produced JSON, that JSON didn't meet the schema."
+        ),
+    ],
+)
+def test_is_retryable_error_returns_true(error):
+    assert _is_retryable_error(error) is True
+
+
+@pytest.mark.parametrize(
+    "error",
+    [
+        ValueError("some other value error"),
+        RuntimeError("runtime error"),
+        KeyError("missing key"),
+        TypeError("type error"),
+    ],
+)
+def test_is_retryable_error_returns_false(error):
+    assert _is_retryable_error(error) is False