fix: address cursor feedback

Author: andrewklatzke

packages/optimization/src/ldai_optimizer/client.py

@@ -212,6 +212,7 @@ def __init__(self, ldClient: LDAIClient) -> None:
         self._initial_tool_keys: List[str] = []
         self._total_token_usage: int = 0
         self._model_configs: List[Dict[str, Any]] = []
+        self._last_batch_size: int = 1
 
         if os.environ.get("LAUNCHDARKLY_API_KEY"):
             self._has_api_key = True
@@ -1123,6 +1124,7 @@ async def _run_ground_truth_optimization(
         self._last_succeeded_context = None
         self._last_optimization_result_id = None
         self._total_token_usage = 0
+        self._last_batch_size = 1
         self._initialize_class_members_from_config(agent_config)
 
         # Seed from the first model choice on the first iteration
@@ -1331,6 +1333,9 @@ async def _run_ground_truth_optimization(
                 self._record_baseline_from_batch(attempt_results)
             self._history.extend(attempt_results)
             self._history = _trim_history(self._history, n)
+            # Track batch size so _all_judges_passing checks every sample in this
+            # attempt, not just the last one.
+            self._last_batch_size = n
 
             logger.info(
                 "[GT Attempt %d] -> %d/%d samples failed — generating new variation",
@@ -2140,28 +2145,35 @@ def _evaluate_cost(self, optimize_context: OptimizationContext) -> bool:
         return passed
 
     def _all_judges_passing(self) -> bool:
-        """Return True if every user-configured judge passed in the most recent history entry.
+        """Return True if every user-configured judge passed in every sample of the most recent batch.
 
-        Inspects the last context in ``_history`` and checks each score key that
-        corresponds to a judge defined in ``_options.judges`` (skipping synthetic gate
-        entries whose keys begin with ``_``). Returns False when history is empty or any
-        judge score does not meet its threshold.
+        In ground-truth mode the last ``_last_batch_size`` entries in ``_history``
+        correspond to the samples from the latest attempt. All of them must pass;
+        checking only the last entry would incorrectly return True when a middle sample
+        failed but the final sample passed.
+
+        In single-sample (non-GT) mode ``_last_batch_size`` is 1, so only the most
+        recent entry is inspected (original behaviour).
+
+        Synthetic gate entries (keys beginning with ``_``) are skipped.
+        Returns False when history is empty or any judge score does not meet its threshold.
 
         This is used to decide whether variation generation should preserve the current
-        behavior and only optimise for cost, rather than trying to improve quality further.
+        behaviour and only optimise for cost, rather than trying to improve quality further.
         """
         if not self._history or not self._options.judges:
             return False
-        recent = self._history[-1]
-        if not recent.scores:
-            return False
-        for key, judge in self._options.judges.items():
-            result = recent.scores.get(key)
-            if result is None:
-                return False
-            threshold = judge.threshold if judge.threshold is not None else 1.0
-            if not judge_passed(result.score, threshold, judge.is_inverted):
+        batch = self._history[-self._last_batch_size:]
+        for ctx in batch:
+            if not ctx.scores:
                 return False
+            for key, judge in self._options.judges.items():
+                result = ctx.scores.get(key)
+                if result is None:
+                    return False
+                threshold = judge.threshold if judge.threshold is not None else 1.0
+                if not judge_passed(result.score, threshold, judge.is_inverted):
+                    return False
         return True
 
     def _apply_duration_gate(
@@ -2201,15 +2213,12 @@ def _apply_duration_gate(
                 rationale = "Latency gate passed (no baseline)."
             score = 1.0
         else:
-            if self._baseline_duration_ms is not None and ctx.duration_ms is not None:
-                rationale = (
-                    f"Latency improvement gate failed: {ctx.duration_ms:.0f}ms did not improve "
-                    f"by {int((1 - _DURATION_TOLERANCE) * 100)}% vs baseline "
-                    f"{self._baseline_duration_ms:.0f}ms "
-                    f"(required < {self._baseline_duration_ms * _DURATION_TOLERANCE:.0f}ms)."
-                )
-            else:
-                rationale = "Latency gate failed (no baseline data)."
+            rationale = (
+                f"Latency improvement gate failed: {ctx.duration_ms:.0f}ms did not improve "
+                f"by {int((1 - _DURATION_TOLERANCE) * 100)}% vs baseline "
+                f"{self._baseline_duration_ms:.0f}ms "
+                f"(required < {self._baseline_duration_ms * _DURATION_TOLERANCE:.0f}ms)."
+            )
             score = 0.0
         ctx = dataclasses.replace(
             ctx,
@@ -2257,15 +2266,12 @@ def _apply_cost_gate(
                 rationale = "Cost gate passed (no baseline)."
             score = 1.0
         else:
-            if self._baseline_cost_usd is not None and ctx.estimated_cost_usd is not None:
-                rationale = (
-                    f"Cost improvement gate failed: {ctx.estimated_cost_usd:.6f} did not improve "
-                    f"by {int((1 - _COST_TOLERANCE) * 100)}% vs baseline "
-                    f"{self._baseline_cost_usd:.6f} "
-                    f"(required < {self._baseline_cost_usd * _COST_TOLERANCE:.6f})."
-                )
-            else:
-                rationale = "Cost gate failed (no baseline data)."
+            rationale = (
+                f"Cost improvement gate failed: {ctx.estimated_cost_usd:.6f} did not improve "
+                f"by {int((1 - _COST_TOLERANCE) * 100)}% vs baseline "
+                f"{self._baseline_cost_usd:.6f} "
+                f"(required < {self._baseline_cost_usd * _COST_TOLERANCE:.6f})."
+            )
             score = 0.0
         ctx = dataclasses.replace(
             ctx,
@@ -2600,6 +2606,7 @@ async def _run_optimization(
         self._last_succeeded_context = None
         self._last_optimization_result_id = None
         self._total_token_usage = 0
+        self._last_batch_size = 1
         self._initialize_class_members_from_config(agent_config)
 
         # If the LD flag doesn't carry a model name, seed from the first model choice

packages/optimization/src/ldai_optimizer/util.py

@@ -336,7 +336,8 @@ def estimate_cost(
 
     :param usage: Token usage from the agent call. When ``None``, returns ``None``.
     :param model_config: Model config dict from ``get_model_configs()``, or ``None``.
-    :return: Estimated cost in USD, or ``None`` if usage or pricing data is absent.
+    :return: Estimated cost in USD, or ``None`` if usage or pricing data is absent, or if
+        both ``usage.input`` and ``usage.output`` are ``None`` (no token counts available).
     """
     if usage is None:
         return None
@@ -348,8 +349,11 @@ def estimate_cost(
         return None
 
     cost = 0.0
+    computed = False
     if input_price is not None and usage.input is not None:
         cost += usage.input * input_price
+        computed = True
     if output_price is not None and usage.output is not None:
         cost += usage.output * output_price
-    return cost
+        computed = True
+    return cost if computed else None

packages/optimization/tests/test_client.py

@@ -5094,6 +5094,26 @@ def test_zero_usage_with_pricing_returns_zero(self):
         model_config = {"costPerInputToken": 0.001, "costPerOutputToken": 0.002}
         assert estimate_cost(usage, model_config) == pytest.approx(0.0)
 
+    def test_returns_none_when_both_token_counts_are_none(self):
+        # Pricing exists but both input and output are None — no token counts to
+        # compute from, so we must return None rather than 0.0 to avoid
+        # cost-gate treating unknown cost as zero cost.
+        usage = TokenUsage(total=None, input=None, output=None)
+        model_config = {"costPerInputToken": 0.001, "costPerOutputToken": 0.002}
+        assert estimate_cost(usage, model_config) is None
+
+    def test_returns_partial_cost_when_only_input_count_is_none(self):
+        # Only output count available — should still compute a partial cost.
+        usage = TokenUsage(total=40, input=None, output=40)
+        model_config = {"costPerInputToken": 0.001, "costPerOutputToken": 0.002}
+        assert estimate_cost(usage, model_config) == pytest.approx(40 * 0.002)
+
+    def test_returns_partial_cost_when_only_output_count_is_none(self):
+        # Only input count available — should still compute a partial cost.
+        usage = TokenUsage(total=60, input=60, output=None)
+        model_config = {"costPerInputToken": 0.001, "costPerOutputToken": 0.002}
+        assert estimate_cost(usage, model_config) == pytest.approx(60 * 0.001)
+
 
 # ---------------------------------------------------------------------------
 # _acceptance_criteria_implies_cost_optimization
@@ -5629,10 +5649,11 @@ def test_gate_scores_do_not_affect_result(self):
         assert self.client._all_judges_passing() is True
 
     def test_uses_most_recent_history_entry(self):
-        """Only the last history entry is inspected."""
+        """In non-GT mode (_last_batch_size=1) only the last history entry is inspected."""
         self.client._options = _make_options(judges={
             "accuracy": OptimizationJudge(threshold=0.8, acceptance_statement="accurate"),
         })
+        self.client._last_batch_size = 1
         self.client._history = [
             self._ctx_with_scores({"accuracy": JudgeResult(score=0.5, rationale="early fail")}, iteration=1),
             self._ctx_with_scores({"accuracy": JudgeResult(score=1.0, rationale="later pass")}, iteration=2),
@@ -5653,6 +5674,60 @@ def test_inverted_judge_fails_when_score_above_threshold(self):
         self.client._history = [self._ctx_with_scores({"toxicity": JudgeResult(score=0.5, rationale="toxic")})]
         assert self.client._all_judges_passing() is False
 
+    # --- GT batch tests ---
+
+    def test_gt_batch_last_sample_passes_but_earlier_fails_returns_false(self):
+        """Core GT bug: if any sample in the batch failed, must return False even if the last passed."""
+        self.client._options = _make_options(judges={
+            "accuracy": OptimizationJudge(threshold=0.8, acceptance_statement="accurate"),
+        })
+        self.client._last_batch_size = 3
+        self.client._history = [
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.3, rationale="fail")}, iteration=1),  # FAILS
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.9, rationale="ok")}, iteration=2),
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.95, rationale="ok")}, iteration=3),
+        ]
+        assert self.client._all_judges_passing() is False
+
+    def test_gt_batch_all_samples_pass_returns_true(self):
+        self.client._options = _make_options(judges={
+            "accuracy": OptimizationJudge(threshold=0.8, acceptance_statement="accurate"),
+        })
+        self.client._last_batch_size = 3
+        self.client._history = [
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.85, rationale="ok")}, iteration=1),
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.90, rationale="ok")}, iteration=2),
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.95, rationale="ok")}, iteration=3),
+        ]
+        assert self.client._all_judges_passing() is True
+
+    def test_gt_batch_middle_sample_fails_returns_false(self):
+        self.client._options = _make_options(judges={
+            "accuracy": OptimizationJudge(threshold=0.8, acceptance_statement="accurate"),
+        })
+        self.client._last_batch_size = 3
+        self.client._history = [
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.95, rationale="ok")}, iteration=1),
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.20, rationale="fail")}, iteration=2),  # FAILS
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.95, rationale="ok")}, iteration=3),
+        ]
+        assert self.client._all_judges_passing() is False
+
+    def test_gt_batch_size_respected_ignores_older_batches(self):
+        """Entries outside the current batch window should not influence the result."""
+        self.client._options = _make_options(judges={
+            "accuracy": OptimizationJudge(threshold=0.8, acceptance_statement="accurate"),
+        })
+        self.client._last_batch_size = 2
+        # 4 entries; batch covers last 2; first 2 are stale (from a previous attempt)
+        self.client._history = [
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.1, rationale="old fail")}, iteration=1),
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.1, rationale="old fail")}, iteration=2),
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.9, rationale="ok")}, iteration=3),
+            self._ctx_with_scores({"accuracy": JudgeResult(score=0.9, rationale="ok")}, iteration=4),
+        ]
+        assert self.client._all_judges_passing() is True
+
 
 class TestBuildNewVariationPromptCost:
     def _make_history(self) -> list: