diff --git a/packages/optimization/pyproject.toml b/packages/optimization/pyproject.toml index 3584bf4..c833a78 100644 --- a/packages/optimization/pyproject.toml +++ b/packages/optimization/pyproject.toml @@ -21,6 +21,7 @@ classifiers = [ ] dependencies = [ "launchdarkly-server-sdk-ai>=0.16.0", + "coolname>=2.0.0", ] [project.urls] diff --git a/packages/optimization/src/ldai_optimization/client.py b/packages/optimization/src/ldai_optimization/client.py index 9f51e5c..3f33dd0 100644 --- a/packages/optimization/src/ldai_optimization/client.py +++ b/packages/optimization/src/ldai_optimization/client.py @@ -9,13 +9,14 @@ import uuid from typing import Any, Dict, List, Literal, Optional, Union +from coolname import generate_slug + from ldai import AIAgentConfig, AIJudgeConfig, AIJudgeConfigDefault, LDAIClient from ldai.models import LDMessage, ModelConfig from ldclient import Context from ldai_optimization.dataclasses import ( AIJudgeCallConfig, - AutoCommitConfig, GroundTruthOptimizationOptions, GroundTruthSample, JudgeResult, @@ -113,6 +114,8 @@ class OptimizationClient: def __init__(self, ldClient: LDAIClient) -> None: self._ldClient = ldClient + self._last_run_succeeded: bool = False + self._last_succeeded_context: Optional[OptimizationContext] = None if os.environ.get("LAUNCHDARKLY_API_KEY"): self._has_api_key = True @@ -819,10 +822,28 @@ async def optimize_from_options( :param options: Optimization options. :return: Optimization result. """ + if options.auto_commit: + if not self._has_api_key: + raise ValueError( + "auto_commit requires LAUNCHDARKLY_API_KEY to be set" + ) + if not options.project_key: + raise ValueError( + "auto_commit requires project_key to be set on OptimizationOptions" + ) self._agent_key = agent_key context = random.choice(options.context_choices) agent_config = await self._get_agent_config(agent_key, context) - return await self._run_optimization(agent_config, options) + result = await self._run_optimization(agent_config, options) + if options.auto_commit and self._last_run_succeeded and self._last_succeeded_context: + self._commit_variation( + self._last_succeeded_context, + project_key=options.project_key, # type: ignore[arg-type] + ai_config_key=agent_key, + output_key=options.output_key, + base_url=options.base_url, + ) + return result async def optimize_from_ground_truth_options( self, agent_key: str, options: GroundTruthOptimizationOptions @@ -839,10 +860,28 @@ async def optimize_from_ground_truth_options( :param options: Ground truth optimization options including the ordered sample list. :return: List of OptimizationContexts from the final attempt (one per sample). """ + if options.auto_commit: + if not self._has_api_key: + raise ValueError( + "auto_commit requires LAUNCHDARKLY_API_KEY to be set" + ) + if not options.project_key: + raise ValueError( + "auto_commit requires project_key to be set on GroundTruthOptimizationOptions" + ) self._agent_key = agent_key context = random.choice(options.context_choices) agent_config = await self._get_agent_config(agent_key, context) - return await self._run_ground_truth_optimization(agent_config, options) + result = await self._run_ground_truth_optimization(agent_config, options) + if options.auto_commit and self._last_run_succeeded and self._last_succeeded_context: + self._commit_variation( + self._last_succeeded_context, + project_key=options.project_key, # type: ignore[arg-type] + ai_config_key=agent_key, + output_key=options.output_key, + base_url=options.base_url, + ) + return result async def _run_ground_truth_optimization( self, @@ -874,6 +913,8 @@ async def _run_ground_truth_optimization( ) self._options = bridge self._agent_config = agent_config + self._last_run_succeeded = False + self._last_succeeded_context = None self._initialize_class_members_from_config(agent_config) # Seed from the first model choice on the first iteration @@ -995,6 +1036,8 @@ async def _run_ground_truth_optimization( attempt, n, ) + self._last_run_succeeded = True + self._last_succeeded_context = last_ctx self._safe_status_update("success", last_ctx, last_ctx.iteration) if self._options.on_passing_result: try: @@ -1011,6 +1054,8 @@ async def _run_ground_truth_optimization( "[GT Optimization] -> Failed after %d attempt(s) — not all samples passed", attempt, ) + self._last_run_succeeded = False + self._last_succeeded_context = None self._safe_status_update("failure", last_ctx, last_ctx.iteration) if self._options.on_failing_result: try: @@ -1037,6 +1082,8 @@ async def _run_ground_truth_optimization( logger.exception( "[GT Attempt %d] -> Variation generation failed", attempt ) + self._last_run_succeeded = False + self._last_succeeded_context = None self._safe_status_update("failure", last_ctx, last_ctx.iteration) if self._options.on_failing_result: try: @@ -1282,6 +1329,10 @@ async def optimize_from_config( raise ValueError( "LAUNCHDARKLY_API_KEY is not set, so optimize_from_config is not available" ) + if options.auto_commit and not self._has_api_key: + raise ValueError( + "auto_commit requires LAUNCHDARKLY_API_KEY to be set" + ) assert self._api_key is not None api_client = LDApiClient( @@ -1303,8 +1354,19 @@ async def optimize_from_config( config, options, api_client, optimization_id, run_id ) if isinstance(optimization_options, GroundTruthOptimizationOptions): - return await self._run_ground_truth_optimization(agent_config, optimization_options) - return await self._run_optimization(agent_config, optimization_options) + result = await self._run_ground_truth_optimization(agent_config, optimization_options) + else: + result = await self._run_optimization(agent_config, optimization_options) + + if options.auto_commit and self._last_run_succeeded and self._last_succeeded_context: + self._commit_variation( + self._last_succeeded_context, + project_key=options.project_key, + ai_config_key=config["aiConfigKey"], + output_key=options.output_key, + api_client=api_client, + ) + return result def _build_options_from_config( self, @@ -1621,6 +1683,8 @@ def _handle_success( :return: The passing OptimizationContext """ logger.info("[Iteration %d] -> Optimization succeeded", iteration) + self._last_run_succeeded = True + self._last_succeeded_context = optimize_context self._safe_status_update("success", optimize_context, iteration) if self._options.on_passing_result: try: @@ -1647,6 +1711,8 @@ def _handle_failure( logger.warning( "[Optimization] -> Optimization failed after %d attempt(s)", iteration ) + self._last_run_succeeded = False + self._last_succeeded_context = None self._safe_status_update("failure", optimize_context, iteration) if self._options.on_failing_result: try: @@ -1657,6 +1723,93 @@ def _handle_failure( ) return optimize_context + def _commit_variation( + self, + optimize_context: OptimizationContext, + project_key: str, + ai_config_key: str, + output_key: Optional[str], + api_client: Optional[LDApiClient] = None, + base_url: Optional[str] = None, + ) -> str: + """Commit the winning optimization context as a new AI Config variation. + + Determines a unique variation key (from output_key or an auto-generated + adjective-noun slug), checks for collisions against existing variation keys, + appends a random hex suffix if the key is taken, then POSTs the new variation + with up to 2 retries before raising on persistent failure. + + :param optimize_context: The winning OptimizationContext. + :param project_key: LaunchDarkly project key. + :param ai_config_key: The AI Config key to add the variation to. + :param output_key: Desired variation key/name; auto-generated if None. + :param api_client: Optional pre-built LDApiClient to reuse (e.g. from optimize_from_config). + :param base_url: Optional base URL override forwarded to a newly created LDApiClient. + :return: The created variation key. + :raises LDApiError: If the variation cannot be created after retries. + """ + if api_client is None: + assert self._api_key is not None + api_client = LDApiClient( + self._api_key, + **({"base_url": base_url} if base_url else {}), + ) + + candidate = output_key if output_key else generate_slug(2) + + try: + ai_config = api_client.get_ai_config(project_key, ai_config_key) + existing_keys = {v["key"] for v in ai_config.get("variations", [])} + except Exception: + logger.warning( + "Could not fetch AI Config to check variation key collisions; proceeding with candidate key." + ) + existing_keys = set() + + if candidate in existing_keys: + suffix = "%04x" % random.randint(0, 0xFFFF) + candidate = f"{candidate}-{suffix}" + logger.info("Variation key collision detected; using '%s' instead.", candidate) + + model_name = optimize_context.current_model or "" + model_config_key = model_name # fallback if lookup fails + try: + model_configs = api_client.get_model_configs(project_key) + match = next((mc for mc in model_configs if mc.get("id") == model_name), None) + if match: + model_config_key = match["key"] + else: + logger.debug( + "No model config found for model id '%s'; using model name as key.", model_name + ) + except Exception as exc: + logger.debug("Could not fetch model configs to resolve modelConfigKey: %s", exc) + + payload: Dict[str, Any] = { + "key": candidate, + "name": candidate, + "mode": "agent", + "instructions": optimize_context.current_instructions, + "modelConfigKey": model_config_key, + } + + last_exc: Optional[Exception] = None + for attempt in range(1, 4): + try: + api_client.create_ai_config_variation(project_key, ai_config_key, payload) + logger.info( + "Auto-committed variation '%s' to AI Config '%s'.", candidate, ai_config_key + ) + return candidate + except Exception as exc: + last_exc = exc + if attempt < 3: + logger.warning( + "Failed to create variation (attempt %d/3): %s. Retrying...", attempt, exc + ) + + raise last_exc # type: ignore[misc] + async def _run_validation_phase( self, passing_context: OptimizationContext, @@ -1816,6 +1969,8 @@ async def _run_optimization( """ self._options = options self._agent_config = agent_config + self._last_run_succeeded = False + self._last_succeeded_context = None self._initialize_class_members_from_config(agent_config) # If the LD flag doesn't carry a model name, seed from the first model choice diff --git a/packages/optimization/src/ldai_optimization/dataclasses.py b/packages/optimization/src/ldai_optimization/dataclasses.py index 02b6a7d..edcbd8b 100644 --- a/packages/optimization/src/ldai_optimization/dataclasses.py +++ b/packages/optimization/src/ldai_optimization/dataclasses.py @@ -152,14 +152,6 @@ class OptimizationJudge: acceptance_statement: Optional[str] = None -@dataclass -class AutoCommitConfig: - """Configuration for auto-committing optimization results to LaunchDarkly.""" - - enabled: bool = False - project_key: Optional[str] = None - - @dataclass class OptimizationContext: """Context for a single optimization iteration.""" @@ -291,10 +283,11 @@ class OptimizationOptions: on_turn: Optional[Callable[[OptimizationContext], bool]] = ( None # if you want manual control of pass/fail ) - # Results - Optional - auto_commit: Optional[AutoCommitConfig] = ( - None # configuration for automatically saving results back to LaunchDarkly - ) + # Auto-commit - Optional + auto_commit: bool = False + project_key: Optional[str] = None # required when auto_commit=True + output_key: Optional[str] = None # variation key/name; auto-generated if omitted + base_url: Optional[str] = None # override to target a non-default LD instance on_passing_result: Optional[Callable[[OptimizationContext], None]] = None on_failing_result: Optional[Callable[[OptimizationContext], None]] = None # called to provide status updates during the optimization flow @@ -379,6 +372,11 @@ class GroundTruthOptimizationOptions: None, ] ] = None + # Auto-commit - Optional + auto_commit: bool = False + project_key: Optional[str] = None # required when auto_commit=True + output_key: Optional[str] = None # variation key/name; auto-generated if omitted + base_url: Optional[str] = None # override to target a non-default LD instance def __post_init__(self): """Validate required options.""" @@ -425,6 +423,9 @@ class OptimizationFromConfigOptions: on_failing_result: Optional[Callable[["OptimizationContext"], None]] = None on_status_update: Optional[Callable[[_StatusLiteral, "OptimizationContext"], None]] = None base_url: Optional[str] = None + # Auto-commit defaults to True for config-driven runs; set False to disable + auto_commit: bool = True + output_key: Optional[str] = None # variation key/name; auto-generated if omitted def __post_init__(self): """Validate required options.""" diff --git a/packages/optimization/src/ldai_optimization/ld_api_client.py b/packages/optimization/src/ldai_optimization/ld_api_client.py index 34f5921..a807629 100644 --- a/packages/optimization/src/ldai_optimization/ld_api_client.py +++ b/packages/optimization/src/ldai_optimization/ld_api_client.py @@ -178,6 +178,9 @@ def __init__(self, api_key: str, base_url: str = _BASE_URL) -> None: self._api_key = api_key self._base_url = base_url.rstrip("/") + def __repr__(self) -> str: + return f"LDApiClient(base_url={self._base_url!r})" + def _auth_headers(self) -> Dict[str, str]: return {"Authorization": self._api_key} @@ -208,6 +211,104 @@ def _request(self, method: str, path: str, body: Any = None) -> Any: path=path, ) from exc + def _ai_config_headers(self) -> Dict[str, str]: + return {**self._auth_headers(), "LD-API-Version": "beta"} + + def get_model_configs(self, project_key: str) -> List[Dict[str, Any]]: + """Fetch all AI model configs for a project. + + :param project_key: LaunchDarkly project key. + :return: List of model config dicts (each has at minimum ``id`` and ``key``). + :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 + + def get_ai_config(self, project_key: str, config_key: str) -> Any: + """Fetch a single AI Config by key, including its variations. + + :param project_key: LaunchDarkly project key. + :param config_key: Key of the AI Config (aiConfigKey). + :return: Raw AI Config dict with a ``variations`` list. + :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 + + def create_ai_config_variation( + self, project_key: str, config_key: str, payload: Dict[str, Any] + ) -> Any: + """Create a new variation on an AI Config. + + :param project_key: LaunchDarkly project key. + :param config_key: Key of the AI Config. + :param payload: Variation payload (key, name, mode, instructions, model). + :return: Created AIConfigVariation dict. + :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 + def get_agent_optimization( self, project_key: str, optimization_key: str ) -> AgentOptimizationConfig: diff --git a/packages/optimization/tests/test_client.py b/packages/optimization/tests/test_client.py index 8cec72f..c88cef8 100644 --- a/packages/optimization/tests/test_client.py +++ b/packages/optimization/tests/test_client.py @@ -82,6 +82,7 @@ def _make_options( judges=None, max_attempts: int = 3, variable_choices=None, + **extra, ) -> OptimizationOptions: if handle_agent_call is None: handle_agent_call = AsyncMock(return_value=OptimizationResponse(output="The capital of France is Paris.")) @@ -103,6 +104,7 @@ def _make_options( handle_agent_call=handle_agent_call, handle_judge_call=handle_judge_call, judges=judges, + **extra, ) @@ -2165,7 +2167,7 @@ def test_on_turn_satisfies_criteria_requirement(self): def _make_gt_options(**overrides) -> GroundTruthOptimizationOptions: - defaults = dict( + defaults: Dict[str, Any] = dict( context_choices=[LD_CONTEXT], ground_truth_responses=[ GroundTruthSample(user_input="What is 2+2?", expected_response="4", variables={"lang": "English"}), @@ -2184,6 +2186,39 @@ def _make_gt_options(**overrides) -> GroundTruthOptimizationOptions: return GroundTruthOptimizationOptions(**defaults) +def _make_winning_context( + model: str = "gpt-4o", + instructions: str = "Be helpful.", + parameters: Dict[str, Any] | None = None, +) -> OptimizationContext: + """Return a minimal OptimizationContext representing a successful run.""" + return OptimizationContext( + scores={}, + completion_response="The answer is 4.", + current_instructions=instructions, + current_parameters=parameters or {}, + current_variables={}, + current_model=model, + iteration=1, + ) + + +def _make_api_client_for_commit( + existing_variation_keys: list | None = None, + model_configs: list | None = None, +) -> MagicMock: + """Return a mock LDApiClient pre-configured for _commit_variation calls.""" + mock = MagicMock() + existing = existing_variation_keys or [] + mock.get_ai_config.return_value = {"variations": [{"key": k} for k in existing]} + mock.get_model_configs.return_value = model_configs if model_configs is not None else [ + {"id": "gpt-4o", "key": "OpenAI.gpt-4o"}, + {"id": "gpt-4o-mini", "key": "OpenAI.gpt-4o-mini"}, + ] + mock.create_ai_config_variation.return_value = {"key": "new-variation"} + return mock + + class TestRunGroundTruthOptimization: def setup_method(self): self.mock_ldai = _make_ldai_client() @@ -2960,3 +2995,488 @@ async def test_no_duration_gate_in_gt_mode_when_no_latency_keywords(self): # Succeeds on first attempt even with slow duration (no latency keyword → no gate) assert isinstance(results, list) assert mock_execute.call_count == 2 + + +# --------------------------------------------------------------------------- +# _commit_variation +# --------------------------------------------------------------------------- + + +class TestCommitVariation: + def _make_client(self) -> OptimizationClient: + with patch.dict("os.environ", {"LAUNCHDARKLY_API_KEY": "test-api-key"}): + return OptimizationClient(_make_ldai_client()) + + # --- key generation --- + + def test_uses_output_key_as_variation_key(self): + client = self._make_client() + api_client = _make_api_client_for_commit() + + key = client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key="my-custom-key", api_client=api_client, + ) + + assert key == "my-custom-key" + payload = api_client.create_ai_config_variation.call_args[0][2] + assert payload["key"] == "my-custom-key" + assert payload["name"] == "my-custom-key" + + def test_generates_slug_when_output_key_is_none(self): + client = self._make_client() + api_client = _make_api_client_for_commit() + + with patch("ldai_optimization.client.generate_slug", return_value="fancy-panda"): + key = client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key=None, api_client=api_client, + ) + + assert key == "fancy-panda" + payload = api_client.create_ai_config_variation.call_args[0][2] + assert payload["key"] == "fancy-panda" + assert payload["name"] == "fancy-panda" + + # --- collision handling --- + + def test_appends_hex_suffix_on_key_collision(self): + client = self._make_client() + api_client = _make_api_client_for_commit(existing_variation_keys=["my-key"]) + + with patch("ldai_optimization.client.random.randint", return_value=0x1234): + key = client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key="my-key", api_client=api_client, + ) + + assert key == "my-key-1234" + payload = api_client.create_ai_config_variation.call_args[0][2] + assert payload["key"] == "my-key-1234" + + def test_no_suffix_when_key_does_not_collide(self): + client = self._make_client() + api_client = _make_api_client_for_commit(existing_variation_keys=["other-key"]) + + key = client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key="my-key", api_client=api_client, + ) + + assert key == "my-key" + + def test_proceeds_with_candidate_when_get_ai_config_raises(self): + client = self._make_client() + api_client = _make_api_client_for_commit() + api_client.get_ai_config.side_effect = Exception("network error") + + key = client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key="my-key", api_client=api_client, + ) + + assert key == "my-key" + api_client.create_ai_config_variation.assert_called_once() + + # --- payload shape --- + + def test_payload_mode_is_agent(self): + client = self._make_client() + api_client = _make_api_client_for_commit() + + client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key="k", api_client=api_client, + ) + + payload = api_client.create_ai_config_variation.call_args[0][2] + assert payload["mode"] == "agent" + + def test_payload_instructions_from_context(self): + client = self._make_client() + api_client = _make_api_client_for_commit() + ctx = _make_winning_context(instructions="You are a travel assistant.") + + client._commit_variation( + ctx, project_key="my-project", + ai_config_key="my-agent", output_key="k", api_client=api_client, + ) + + payload = api_client.create_ai_config_variation.call_args[0][2] + assert payload["instructions"] == "You are a travel assistant." + + def test_create_called_with_correct_project_and_config_key(self): + client = self._make_client() + api_client = _make_api_client_for_commit() + + client._commit_variation( + _make_winning_context(), project_key="proj-abc", + ai_config_key="agent-xyz", output_key="k", api_client=api_client, + ) + + args = api_client.create_ai_config_variation.call_args[0] + assert args[0] == "proj-abc" + assert args[1] == "agent-xyz" + + # --- modelConfigKey resolution --- + + def test_model_config_key_resolved_via_api_match_on_id(self): + client = self._make_client() + api_client = _make_api_client_for_commit(model_configs=[ + {"id": "gpt-4o", "key": "OpenAI.gpt-4o"}, + {"id": "claude-3", "key": "Anthropic.claude-3"}, + ]) + + client._commit_variation( + _make_winning_context(model="gpt-4o"), project_key="my-project", + ai_config_key="my-agent", output_key="k", api_client=api_client, + ) + + payload = api_client.create_ai_config_variation.call_args[0][2] + assert payload["modelConfigKey"] == "OpenAI.gpt-4o" + + def test_model_config_key_falls_back_to_model_name_when_no_id_match(self): + client = self._make_client() + api_client = _make_api_client_for_commit(model_configs=[ + {"id": "claude-3", "key": "Anthropic.claude-3"}, + ]) + + client._commit_variation( + _make_winning_context(model="gpt-4o"), project_key="my-project", + ai_config_key="my-agent", output_key="k", api_client=api_client, + ) + + payload = api_client.create_ai_config_variation.call_args[0][2] + assert payload["modelConfigKey"] == "gpt-4o" + + def test_model_config_key_falls_back_when_get_model_configs_raises(self): + client = self._make_client() + api_client = _make_api_client_for_commit() + api_client.get_model_configs.side_effect = Exception("network error") + + client._commit_variation( + _make_winning_context(model="gpt-4o"), project_key="my-project", + ai_config_key="my-agent", output_key="k", api_client=api_client, + ) + + payload = api_client.create_ai_config_variation.call_args[0][2] + assert payload["modelConfigKey"] == "gpt-4o" + + # --- retry logic --- + + def test_retries_on_transient_failure_and_succeeds(self): + client = self._make_client() + api_client = _make_api_client_for_commit() + api_client.create_ai_config_variation.side_effect = [ + Exception("transient"), + {"key": "my-key"}, + ] + + key = client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key="my-key", api_client=api_client, + ) + + assert key == "my-key" + assert api_client.create_ai_config_variation.call_count == 2 + + def test_raises_after_three_consecutive_failures(self): + client = self._make_client() + api_client = _make_api_client_for_commit() + api_client.create_ai_config_variation.side_effect = RuntimeError("permanent") + + with pytest.raises(RuntimeError, match="permanent"): + client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key="k", api_client=api_client, + ) + + assert api_client.create_ai_config_variation.call_count == 3 + + # --- LDApiClient construction --- + + def test_creates_api_client_from_stored_key_when_none_provided(self): + client = self._make_client() + + with patch("ldai_optimization.client.LDApiClient") as MockLDApiClient: + MockLDApiClient.return_value = _make_api_client_for_commit() + client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key="k", + ) + + MockLDApiClient.assert_called_once_with("test-api-key") + + def test_passes_base_url_when_creating_api_client(self): + client = self._make_client() + + with patch("ldai_optimization.client.LDApiClient") as MockLDApiClient: + MockLDApiClient.return_value = _make_api_client_for_commit() + client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key="k", + base_url="https://app.launchdarkly.us", + ) + + MockLDApiClient.assert_called_once_with( + "test-api-key", base_url="https://app.launchdarkly.us" + ) + + def test_reuses_provided_api_client_without_creating_new_one(self): + client = self._make_client() + api_client = _make_api_client_for_commit() + + with patch("ldai_optimization.client.LDApiClient") as MockLDApiClient: + client._commit_variation( + _make_winning_context(), project_key="my-project", + ai_config_key="my-agent", output_key="k", api_client=api_client, + ) + + MockLDApiClient.assert_not_called() + + +# --------------------------------------------------------------------------- +# auto_commit in optimize_from_options +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestAutoCommitInOptimizeFromOptions: + def _make_client_with_key(self) -> OptimizationClient: + with patch.dict("os.environ", {"LAUNCHDARKLY_API_KEY": "test-api-key"}): + return OptimizationClient(_make_ldai_client()) + + def _make_client_without_key(self) -> OptimizationClient: + client = OptimizationClient(_make_ldai_client()) + client._has_api_key = False + client._api_key = None + return client + + async def test_commit_called_on_success_when_auto_commit_true(self): + client = self._make_client_with_key() + options = _make_options(auto_commit=True, project_key="my-project") + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_options("test-agent", options) + + mock_commit.assert_called_once() + + async def test_commit_not_called_when_auto_commit_false(self): + client = self._make_client_with_key() + options = _make_options() # auto_commit defaults to False + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_options("test-agent", options) + + mock_commit.assert_not_called() + + async def test_commit_not_called_when_run_fails(self): + client = self._make_client_with_key() + options = _make_options( + auto_commit=True, + project_key="my-project", + handle_judge_call=AsyncMock(return_value=OptimizationResponse(output=JUDGE_FAIL_RESPONSE)), + max_attempts=1, + ) + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_options("test-agent", options) + + mock_commit.assert_not_called() + + async def test_raises_when_auto_commit_true_and_no_api_key(self): + client = self._make_client_without_key() + options = _make_options(auto_commit=True, project_key="my-project") + + with pytest.raises(ValueError, match="LAUNCHDARKLY_API_KEY"): + await client.optimize_from_options("test-agent", options) + + async def test_raises_when_auto_commit_true_and_no_project_key(self): + client = self._make_client_with_key() + options = _make_options(auto_commit=True, project_key=None) + + with pytest.raises(ValueError, match="project_key"): + await client.optimize_from_options("test-agent", options) + + async def test_output_key_forwarded_to_commit(self): + client = self._make_client_with_key() + options = _make_options( + auto_commit=True, project_key="my-project", output_key="my-variation" + ) + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_options("test-agent", options) + + assert mock_commit.call_args[1]["output_key"] == "my-variation" + + async def test_base_url_forwarded_to_commit(self): + client = self._make_client_with_key() + options = _make_options( + auto_commit=True, + project_key="my-project", + base_url="https://app.launchdarkly.us", + ) + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_options("test-agent", options) + + assert mock_commit.call_args[1]["base_url"] == "https://app.launchdarkly.us" + + async def test_agent_key_used_as_ai_config_key(self): + client = self._make_client_with_key() + options = _make_options(auto_commit=True, project_key="my-project") + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_options("test-agent", options) + + assert mock_commit.call_args[1]["ai_config_key"] == "test-agent" + + +# --------------------------------------------------------------------------- +# auto_commit in optimize_from_ground_truth_options +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestAutoCommitInOptimizeFromGroundTruthOptions: + def _make_client_with_key(self) -> OptimizationClient: + with patch.dict("os.environ", {"LAUNCHDARKLY_API_KEY": "test-api-key"}): + return OptimizationClient(_make_ldai_client()) + + def _make_client_without_key(self) -> OptimizationClient: + client = OptimizationClient(_make_ldai_client()) + client._has_api_key = False + client._api_key = None + return client + + async def test_commit_called_on_success_when_auto_commit_true(self): + client = self._make_client_with_key() + opts = _make_gt_options(auto_commit=True, project_key="my-project") + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_ground_truth_options("test-agent", opts) + + mock_commit.assert_called_once() + + async def test_commit_not_called_when_auto_commit_false(self): + client = self._make_client_with_key() + opts = _make_gt_options() # auto_commit defaults to False + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_ground_truth_options("test-agent", opts) + + mock_commit.assert_not_called() + + async def test_commit_not_called_when_run_fails(self): + client = self._make_client_with_key() + opts = _make_gt_options( + auto_commit=True, + project_key="my-project", + handle_judge_call=AsyncMock(return_value=OptimizationResponse(output=JUDGE_FAIL_RESPONSE)), + max_attempts=1, + ) + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_ground_truth_options("test-agent", opts) + + mock_commit.assert_not_called() + + async def test_raises_when_auto_commit_true_and_no_api_key(self): + client = self._make_client_without_key() + opts = _make_gt_options(auto_commit=True, project_key="my-project") + + with pytest.raises(ValueError, match="LAUNCHDARKLY_API_KEY"): + await client.optimize_from_ground_truth_options("test-agent", opts) + + async def test_raises_when_auto_commit_true_and_no_project_key(self): + client = self._make_client_with_key() + opts = _make_gt_options(auto_commit=True, project_key=None) + + with pytest.raises(ValueError, match="project_key"): + await client.optimize_from_ground_truth_options("test-agent", opts) + + async def test_output_key_forwarded_to_commit(self): + client = self._make_client_with_key() + opts = _make_gt_options( + auto_commit=True, project_key="my-project", output_key="my-variation" + ) + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_ground_truth_options("test-agent", opts) + + assert mock_commit.call_args[1]["output_key"] == "my-variation" + + async def test_base_url_forwarded_to_commit(self): + client = self._make_client_with_key() + opts = _make_gt_options( + auto_commit=True, + project_key="my-project", + base_url="https://app.launchdarkly.us", + ) + + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_ground_truth_options("test-agent", opts) + + assert mock_commit.call_args[1]["base_url"] == "https://app.launchdarkly.us" + + +# --------------------------------------------------------------------------- +# auto_commit in optimize_from_config +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestAutoCommitInOptimizeFromConfig: + def _make_client_with_key(self) -> OptimizationClient: + with patch.dict("os.environ", {"LAUNCHDARKLY_API_KEY": "test-api-key"}): + return OptimizationClient(_make_ldai_client()) + + async def test_commit_called_by_default(self): + """auto_commit=True is the default for optimize_from_config.""" + client = self._make_client_with_key() + mock_api = _make_mock_api_client() + mock_api.get_agent_optimization = MagicMock(return_value=dict(_API_CONFIG)) + + with patch("ldai_optimization.client.LDApiClient", return_value=mock_api): + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_config("my-opt", _make_from_config_options()) + + mock_commit.assert_called_once() + + async def test_commit_not_called_when_auto_commit_false(self): + client = self._make_client_with_key() + mock_api = _make_mock_api_client() + mock_api.get_agent_optimization = MagicMock(return_value=dict(_API_CONFIG)) + + with patch("ldai_optimization.client.LDApiClient", return_value=mock_api): + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_config( + "my-opt", _make_from_config_options(auto_commit=False) + ) + + mock_commit.assert_not_called() + + async def test_commit_receives_pre_built_api_client(self): + """The api_client created for fetching config is reused for _commit_variation.""" + client = self._make_client_with_key() + mock_api = _make_mock_api_client() + mock_api.get_agent_optimization = MagicMock(return_value=dict(_API_CONFIG)) + + with patch("ldai_optimization.client.LDApiClient", return_value=mock_api): + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_config("my-opt", _make_from_config_options()) + + assert mock_commit.call_args[1]["api_client"] is mock_api + + async def test_output_key_forwarded_to_commit(self): + client = self._make_client_with_key() + mock_api = _make_mock_api_client() + mock_api.get_agent_optimization = MagicMock(return_value=dict(_API_CONFIG)) + + with patch("ldai_optimization.client.LDApiClient", return_value=mock_api): + with patch.object(client, "_commit_variation") as mock_commit: + await client.optimize_from_config( + "my-opt", _make_from_config_options(output_key="my-variation") + ) + + assert mock_commit.call_args[1]["output_key"] == "my-variation"