chore: add retry logic and structural validation for variations

Author: andrewklatzke

packages/optimization/src/ldai_optimizer/client.py

@@ -61,6 +61,7 @@
     extract_json_from_response,
     interpolate_variables,
     restore_variable_placeholders,
+    validate_variation_response,
 )
 
 logger = logging.getLogger(__name__)
@@ -1179,24 +1180,18 @@ def _apply_new_variation_response(
         :param iteration: Current iteration number for logging
         :return: A new OptimizationContext populated with the updated configuration
         """
-        missing_fields = []
-        if "current_instructions" not in response_data:
-            missing_fields.append("current_instructions")
-        if "current_parameters" not in response_data:
-            missing_fields.append("current_parameters")
-        if "model" not in response_data:
-            missing_fields.append("model")
-
-        if missing_fields:
+        validation_errors = validate_variation_response(response_data)
+        if validation_errors:
             logger.debug(
-                "[Iteration %d] -> Response missing required fields: %s. Received fields: %s. Full response_data: %s",
+                "[Iteration %d] -> Variation response failed validation: %s. "
+                "Received fields: %s. Full response_data: %s",
                 iteration,
-                ", ".join(missing_fields),
+                "; ".join(validation_errors),
                 list(response_data.keys()),
                 json.dumps(response_data, indent=2),
             )
             raise ValueError(
-                f"Response missing required fields: {', '.join(missing_fields)}. "
+                f"Variation response failed validation: {'; '.join(validation_errors)}. "
                 f"Received fields: {list(response_data.keys())}"
             )
 

packages/optimization/src/ldai_optimizer/ld_api_client.py

@@ -2,6 +2,7 @@
 
 import json
 import logging
+import time
 import urllib.error
 import urllib.request
 from typing import Any, Dict, List, Optional, TypedDict
@@ -13,6 +14,14 @@
 
 _BASE_URL = "https://app.launchdarkly.com"
 
+_MAX_RETRIES = 3
+_INITIAL_BACKOFF = 1.0  # seconds; doubles on each attempt (1s, 2s, 4s)
+
+# Status codes that warrant a retry.  Everything else (including 400, 401, 403,
+# 404) is a permanent or auth failure — retrying would not help and could lead
+# to corrupted optimization results if some requests succeed and others fail.
+_RETRYABLE_STATUS_CODES = frozenset({429, 500, 502, 503, 504})
+
 
 class LDApiError(Exception):
     """Raised when the LaunchDarkly REST API returns an error or is unreachable.
@@ -186,35 +195,74 @@ def __repr__(self) -> str:
     def _auth_headers(self) -> Dict[str, str]:
         return {"Authorization": self._api_key}
 
-    def _request(self, method: str, path: str, body: Any = None) -> Any:
+    def _request(
+        self,
+        method: str,
+        path: str,
+        body: Any = None,
+        extra_headers: Optional[Dict[str, str]] = None,
+    ) -> Any:
+        """Execute an HTTP request with automatic retry and exponential backoff.
+
+        Retries up to ``_MAX_RETRIES`` times for transient errors (429, 5xx,
+        network failures) with exponential backoff starting at ``_INITIAL_BACKOFF``
+        seconds.  Non-retryable status codes (400, 401, 403, 404, …) are raised
+        immediately without retrying.
+
+        :param method: HTTP method (GET, POST, PATCH, …).
+        :param path: API path, appended to ``self._base_url``.
+        :param body: Optional request body; serialised to JSON.
+        :param extra_headers: Additional headers merged with the auth header.
+        :raises LDApiError: After all retry attempts are exhausted, or immediately
+            for non-retryable status codes.
+        """
         url = f"{self._base_url}{path}"
+        headers = {**self._auth_headers(), **(extra_headers or {})}
         data = json.dumps(body).encode() if body is not None else None
-        headers = self._auth_headers()
         if data is not None:
             headers["Content-Type"] = "application/json"
-        req = urllib.request.Request(url, data=data, headers=headers, method=method)
-        try:
-            with urllib.request.urlopen(req) as resp:
-                raw = resp.read()
-                return json.loads(raw) if raw else None
-        except urllib.error.HTTPError as exc:
-            body_excerpt = exc.read(500).decode(errors="replace")
-            hint = _HTTP_ERROR_HINTS.get(exc.code, "")
-            detail = f"{hint} (API response: {body_excerpt})" if hint else f"API response: {body_excerpt}"
-            raise LDApiError(
-                f"LaunchDarkly API error {exc.code} {exc.msg} for {method} {path}. {detail}",
-                status_code=exc.code,
-                path=path,
-            ) from exc
-        except urllib.error.URLError as exc:
-            raise LDApiError(
-                f"Could not reach LaunchDarkly API at {url}: {exc.reason}. "
-                "Check your network connection and the base_url setting.",
-                path=path,
-            ) from exc
-
-    def _ai_config_headers(self) -> Dict[str, str]:
-        return {**self._auth_headers(), "LD-API-Version": "beta"}
+
+        last_exc: Optional[LDApiError] = None
+        for attempt in range(_MAX_RETRIES + 1):
+            req = urllib.request.Request(url, data=data, headers=headers, method=method)
+            try:
+                with urllib.request.urlopen(req) as resp:
+                    raw = resp.read()
+                    return json.loads(raw) if raw else None
+            except urllib.error.HTTPError as exc:
+                body_excerpt = exc.read(500).decode(errors="replace")
+                hint = _HTTP_ERROR_HINTS.get(exc.code, "")
+                detail = f"{hint} (API response: {body_excerpt})" if hint else f"API response: {body_excerpt}"
+                api_error = LDApiError(
+                    f"LaunchDarkly API error {exc.code} {exc.msg} for {method} {path}. {detail}",
+                    status_code=exc.code,
+                    path=path,
+                )
+                if exc.code not in _RETRYABLE_STATUS_CODES:
+                    raise api_error from exc
+                last_exc = api_error
+            except urllib.error.URLError as exc:
+                last_exc = LDApiError(
+                    f"Could not reach LaunchDarkly API at {url}: {exc.reason}. "
+                    "Check your network connection and the base_url setting.",
+                    path=path,
+                )
+
+            if attempt < _MAX_RETRIES:
+                delay = _INITIAL_BACKOFF * (2 ** attempt)
+                logger.warning(
+                    "LaunchDarkly API request failed (attempt %d/%d, path=%s), "
+                    "retrying in %.1fs: %s",
+                    attempt + 1,
+                    _MAX_RETRIES + 1,
+                    path,
+                    delay,
+                    last_exc,
+                )
+                time.sleep(delay)
+
+        assert last_exc is not None
+        raise last_exc
 
     def get_model_configs(self, project_key: str) -> List[Dict[str, Any]]:
         """Fetch all AI model configs for a project.
@@ -224,26 +272,8 @@ def get_model_configs(self, project_key: str) -> List[Dict[str, Any]]:
         :raises LDApiError: On non-200 HTTP responses or network errors.
         """
         path = f"/api/v2/projects/{project_key}/ai-configs/model-configs"
-        url = f"{self._base_url}{path}"
-        req = urllib.request.Request(url, headers=self._ai_config_headers(), method="GET")
-        try:
-            with urllib.request.urlopen(req) as resp:
-                raw = resp.read()
-                return json.loads(raw) if raw else []
-        except urllib.error.HTTPError as exc:
-            body_excerpt = exc.read(500).decode(errors="replace")
-            hint = _HTTP_ERROR_HINTS.get(exc.code, "")
-            detail = f"{hint} (API response: {body_excerpt})" if hint else f"API response: {body_excerpt}"
-            raise LDApiError(
-                f"LaunchDarkly API error {exc.code} {exc.msg} for GET {path}. {detail}",
-                status_code=exc.code,
-                path=path,
-            ) from exc
-        except urllib.error.URLError as exc:
-            raise LDApiError(
-                f"Could not reach LaunchDarkly API at {url}: {exc.reason}.",
-                path=path,
-            ) from exc
+        result = self._request("GET", path, extra_headers={"LD-API-Version": "beta"})
+        return result if isinstance(result, list) else []
 
     def get_ai_config(self, project_key: str, config_key: str) -> Any:
         """Fetch a single AI Config by key, including its variations.
@@ -254,27 +284,7 @@ def get_ai_config(self, project_key: str, config_key: str) -> Any:
         :raises LDApiError: On non-200 HTTP responses or network errors.
         """
         path = f"/api/v2/projects/{project_key}/ai-configs/{config_key}"
-        headers = self._ai_config_headers()
-        url = f"{self._base_url}{path}"
-        req = urllib.request.Request(url, headers=headers, method="GET")
-        try:
-            with urllib.request.urlopen(req) as resp:
-                raw = resp.read()
-                return json.loads(raw) if raw else None
-        except urllib.error.HTTPError as exc:
-            body_excerpt = exc.read(500).decode(errors="replace")
-            hint = _HTTP_ERROR_HINTS.get(exc.code, "")
-            detail = f"{hint} (API response: {body_excerpt})" if hint else f"API response: {body_excerpt}"
-            raise LDApiError(
-                f"LaunchDarkly API error {exc.code} {exc.msg} for GET {path}. {detail}",
-                status_code=exc.code,
-                path=path,
-            ) from exc
-        except urllib.error.URLError as exc:
-            raise LDApiError(
-                f"Could not reach LaunchDarkly API at {url}: {exc.reason}.",
-                path=path,
-            ) from exc
+        return self._request("GET", path, extra_headers={"LD-API-Version": "beta"})
 
     def create_ai_config_variation(
         self, project_key: str, config_key: str, payload: Dict[str, Any]
@@ -288,28 +298,7 @@ def create_ai_config_variation(
         :raises LDApiError: On non-200 HTTP responses or network errors.
         """
         path = f"/api/v2/projects/{project_key}/ai-configs/{config_key}/variations"
-        url = f"{self._base_url}{path}"
-        data = json.dumps(payload).encode()
-        headers = {**self._ai_config_headers(), "Content-Type": "application/json"}
-        req = urllib.request.Request(url, data=data, headers=headers, method="POST")
-        try:
-            with urllib.request.urlopen(req) as resp:
-                raw = resp.read()
-                return json.loads(raw) if raw else None
-        except urllib.error.HTTPError as exc:
-            body_excerpt = exc.read(500).decode(errors="replace")
-            hint = _HTTP_ERROR_HINTS.get(exc.code, "")
-            detail = f"{hint} (API response: {body_excerpt})" if hint else f"API response: {body_excerpt}"
-            raise LDApiError(
-                f"LaunchDarkly API error {exc.code} {exc.msg} for POST {path}. {detail}",
-                status_code=exc.code,
-                path=path,
-            ) from exc
-        except urllib.error.URLError as exc:
-            raise LDApiError(
-                f"Could not reach LaunchDarkly API at {url}: {exc.reason}.",
-                path=path,
-            ) from exc
+        return self._request("POST", path, body=payload, extra_headers={"LD-API-Version": "beta"})
 
     def get_agent_optimization(
         self, project_key: str, optimization_key: str

packages/optimization/src/ldai_optimizer/util.py

@@ -290,6 +290,47 @@ def create_variation_tool(model_choices: List[str]) -> ToolDefinition:
     )
 
 
+def validate_variation_response(response_data: Dict[str, Any]) -> List[str]:
+    """Validate the shape of a parsed LLM variation response.
+
+    Checks that the three required fields are present and have the expected
+    types. An empty ``current_parameters`` dict is acceptable; an empty
+    ``current_instructions`` or ``model`` string is flagged as an error
+    because downstream code cannot meaningfully use a blank value.
+
+    :param response_data: Parsed dict from the LLM (output of extract_json_from_response).
+    :return: List of human-readable error strings. Empty list means the response is valid.
+    """
+    errors: List[str] = []
+
+    if "current_instructions" not in response_data:
+        errors.append("missing required field 'current_instructions'")
+    elif not isinstance(response_data["current_instructions"], str):
+        errors.append(
+            f"'current_instructions' must be a string, "
+            f"got {type(response_data['current_instructions']).__name__}"
+        )
+    elif not response_data["current_instructions"].strip():
+        errors.append("'current_instructions' must not be empty")
+
+    if "current_parameters" not in response_data:
+        errors.append("missing required field 'current_parameters'")
+    elif not isinstance(response_data["current_parameters"], dict):
+        errors.append(
+            f"'current_parameters' must be a dict, "
+            f"got {type(response_data['current_parameters']).__name__}"
+        )
+
+    if "model" not in response_data:
+        errors.append("missing required field 'model'")
+    elif not isinstance(response_data["model"], str):
+        errors.append(
+            f"'model' must be a string, got {type(response_data['model']).__name__}"
+        )
+
+    return errors
+
+
 def extract_json_from_response(response_str: str) -> Dict[str, Any]:
     """
     Parse a JSON object from an LLM response string.

packages/optimization/tests/test_ld_api_client.py

@@ -5,7 +5,7 @@
 import urllib.request
 from io import BytesIO
 from typing import Any, Dict
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, call, patch
 
 import pytest
 
@@ -14,6 +14,7 @@
     AgentOptimizationResultPost as OptimizationResultPayload,
     LDApiClient,
     LDApiError,
+    _MAX_RETRIES,
     _parse_agent_optimization,
 )
 
@@ -295,10 +296,76 @@ def test_swallows_http_errors_without_raising(self):
             url="http://x", code=500, msg="Server Error", hdrs=MagicMock(), fp=BytesIO(b"err")
         )
         with patch("urllib.request.urlopen", side_effect=http_error):
-            # must not raise
-            client.post_agent_optimization_result("proj", "opt-id", self._make_payload())
+            with patch("time.sleep"):
+                # must not raise even after all retries are exhausted
+                client.post_agent_optimization_result("proj", "opt-id", self._make_payload())
 
     def test_swallows_url_errors_without_raising(self):
         client = LDApiClient("test-key")
         with patch("urllib.request.urlopen", side_effect=urllib.error.URLError("timeout")):
-            client.post_agent_optimization_result("proj", "opt-id", self._make_payload())
+            with patch("time.sleep"):
+                client.post_agent_optimization_result("proj", "opt-id", self._make_payload())
+
+
+# ---------------------------------------------------------------------------
+# LDApiClient retry behaviour
+# ---------------------------------------------------------------------------
+
+
+class TestLDApiClientRetry:
+    def _http_error(self, code: int) -> urllib.error.HTTPError:
+        return urllib.error.HTTPError(
+            url="http://x", code=code, msg="Error", hdrs=MagicMock(), fp=BytesIO(b"body")
+        )
+
+    def test_retryable_error_retries_max_times(self):
+        """A 429 or 5xx should be retried up to _MAX_RETRIES times then raise."""
+        client = LDApiClient("test-key")
+        with patch("urllib.request.urlopen", side_effect=self._http_error(429)) as mock_open:
+            with patch("time.sleep"):
+                with pytest.raises(LDApiError) as exc_info:
+                    client._request("GET", "/path")
+        assert mock_open.call_count == _MAX_RETRIES + 1
+        assert exc_info.value.status_code == 429
+
+    def test_non_retryable_error_raises_immediately(self):
+        """A 401, 403, or 404 should raise after a single attempt with no retries."""
+        for code in (400, 401, 403, 404):
+            client = LDApiClient("test-key")
+            with patch("urllib.request.urlopen", side_effect=self._http_error(code)) as mock_open:
+                with patch("time.sleep") as mock_sleep:
+                    with pytest.raises(LDApiError) as exc_info:
+                        client._request("GET", "/path")
+            assert mock_open.call_count == 1, f"Expected 1 attempt for {code}, got {mock_open.call_count}"
+            mock_sleep.assert_not_called()
+            assert exc_info.value.status_code == code
+
+    def test_url_error_retries_max_times(self):
+        """Network-level errors should also be retried."""
+        client = LDApiClient("test-key")
+        with patch("urllib.request.urlopen", side_effect=urllib.error.URLError("timeout")) as mock_open:
+            with patch("time.sleep"):
+                with pytest.raises(LDApiError):
+                    client._request("GET", "/path")
+        assert mock_open.call_count == _MAX_RETRIES + 1
+
+    def test_backoff_delays_are_exponential(self):
+        """Sleep durations should double on each retry: 1s, 2s, 4s."""
+        client = LDApiClient("test-key")
+        with patch("urllib.request.urlopen", side_effect=self._http_error(500)):
+            with patch("time.sleep") as mock_sleep:
+                with pytest.raises(LDApiError):
+                    client._request("GET", "/path")
+        sleep_calls = [c.args[0] for c in mock_sleep.call_args_list]
+        assert sleep_calls == [1.0, 2.0, 4.0]
+
+    def test_succeeds_on_retry_after_transient_error(self):
+        """If a retryable error clears, the successful response should be returned."""
+        client = LDApiClient("test-key")
+        ok_response = _mock_urlopen({"result": "ok"})
+        side_effects = [self._http_error(500), ok_response]
+        with patch("urllib.request.urlopen", side_effect=side_effects) as mock_open:
+            with patch("time.sleep"):
+                result = client._request("GET", "/path")
+        assert result == {"result": "ok"}
+        assert mock_open.call_count == 2