diff --git a/packages/optimization/src/ldai_optimizer/client.py b/packages/optimization/src/ldai_optimizer/client.py index db232b1..ff10d8a 100644 --- a/packages/optimization/src/ldai_optimizer/client.py +++ b/packages/optimization/src/ldai_optimizer/client.py @@ -15,6 +15,7 @@ about the caller's broader runtime environment beyond the key itself. """ +import asyncio import dataclasses import json import logging @@ -2118,6 +2119,15 @@ def _persist_and_forward( api_client.patch_agent_optimization_result( project_key, optimization_key, result_id, patch ) + # When the winning result is marked successful, make sure + # _last_optimization_result_id tracks it. Without this the + # auto-commit PATCH (which attaches createdVariationKey) is + # sent to the last-posted Phase 2 record rather than to the + # winner, giving a non-winning RUNNING record the latest + # updatedAt and causing the backend to report the run as still + # running even after the optimization completes. + if status == "success": + self._last_optimization_result_id = result_id # Reset tracking state after terminal events so the next main-loop # attempt starts fresh. @@ -2673,16 +2683,20 @@ async def _run_cost_latency_phase( frozen_user_input = winning_ctx.user_input # Build a deterministic, deduplicated list of models to evaluate: - # start with the Phase 1 winner's model, then add each model_choice - # that hasn't been seen yet. This guarantees every user-selected model - # is tried exactly once, in a predictable order. + # always start from model_choices, skipping the Phase 1 winner so it + # doesn't appear as an extra input inside the quality iteration in the + # UI. Fall back to the Phase 1 winner only when no distinct choices + # are provided. phase1_model = winning_ctx.current_model or "" seen_models: set = {phase1_model} - ordered_models: List[str] = [phase1_model] + ordered_models: List[str] = [] for m in self._options.model_choices or []: if m not in seen_models: seen_models.add(m) ordered_models.append(m) + # Fall back to Phase 1 model only if no distinct alternatives exist. + if not ordered_models: + ordered_models.append(phase1_model) # Ensure at least 2 iterations while len(ordered_models) < 2: ordered_models.append(ordered_models[-1]) @@ -2714,12 +2728,29 @@ async def _run_cost_latency_phase( variables=frozen_variables, user_input=frozen_user_input, ) + + gate_placeholders: Dict[str, JudgeResult] = {} + if self._options.latency_optimization: + gate_placeholders["_latency_gate"] = JudgeResult( + score=0.0, rationale="evaluating" + ) + if self._options.token_optimization: + gate_placeholders["_cost_gate"] = JudgeResult( + score=0.0, rationale="evaluating" + ) + if gate_placeholders: + ctx = dataclasses.replace( + ctx, scores={**ctx.scores, **gate_placeholders} + ) self._safe_status_update("generating", ctx, iteration) try: - ctx = await self._execute_agent_turn(ctx, iteration) - except Exception: + ctx = await asyncio.wait_for( + self._execute_agent_turn(ctx, iteration), + timeout=120, + ) + except (Exception, asyncio.TimeoutError): logger.warning( - "[Phase 2 Iter %d] -> Agent call failed (model=%s); " + "[Phase 2 Iter %d] -> Agent call failed or timed out (model=%s); " "skipping this model and trying the next", iteration, self._current_model, @@ -2758,16 +2789,27 @@ async def _run_cost_latency_phase( if i < max_iters - 1: self._safe_status_update("turn completed", ctx, iteration) - # Send terminal FAILED status for each non-winning model attempt. - # We use _safe_status_update directly rather than _handle_failure so that - # exploratory Phase 2 misses don't corrupt _last_run_succeeded, - # _last_succeeded_context, or trigger on_failing_result — those are - # run-level signals that should only fire if the whole optimization fails. - for failed_ctx in non_candidates: - self._safe_status_update("failure", failed_ctx, failed_ctx.iteration) + # Phase 2 is complete. The `status` field on each result carries the + # *run-level* outcome (PASSED / FAILED / RUNNING), not the quality of + # that individual result — the scores already encode individual quality. + # The backend derives the visible run status from the highest-iteration + # result, so every Phase 2 result must end with status=PASSED after a + # successful run; otherwise the highest-numbered result keeps the run in + # RUNNING indefinitely. We use _safe_status_update directly (not + # _handle_failure / _handle_success) for non-winners so that + # _last_run_succeeded, _last_succeeded_context, and on_failing_result are + # not corrupted — those are reserved for run-level outcomes. + for non_candidate_ctx in non_candidates: + self._safe_status_update("success", non_candidate_ctx, non_candidate_ctx.iteration) if candidates: best = self._pick_best_candidate(candidates) + # Non-best candidates: mark PASSED (run succeeded) before the winner + # so _handle_success remains the very last update, preserving the + # correct _last_succeeded_context for the caller. + for other in candidates: + if other.iteration != best.iteration: + self._safe_status_update("success", other, other.iteration) # Suppress on_passing_result here — the caller fires it once with the # true final winner after Phase 2 returns, so it is never double-fired. self._handle_success(best, best.iteration, suppress_user_callbacks=True)