fix: don't only evaluate final input in GT results

Author: andrewklatzke

packages/optimization/src/ldai_optimizer/client.py

@@ -224,17 +224,51 @@ def _initialize_class_members_from_config(
         self._baseline_cost_usd: Optional[float] = None
 
     def _record_baseline(self, ctx: OptimizationContext) -> None:
-        """Capture duration/cost baseline from the first iteration appended to history.
+        """Capture duration/cost baseline from a single context.
 
-        Called once per run (subsequent calls are no-ops once both values are set).
-        Storing these explicitly lets ``_trim_history`` use a simple tail slice without
-        needing to preserve ``history[0]`` as an anchor.
+        Used by the standard (non-GT) optimization loop where each iteration
+        produces one result. Called once per run (subsequent calls are no-ops
+        once both values are set). Storing these explicitly lets
+        ``_trim_history`` use a simple tail slice without needing to preserve
+        ``history[0]`` as an anchor.
         """
         if self._baseline_duration_ms is None and ctx.duration_ms is not None:
             self._baseline_duration_ms = ctx.duration_ms
         if self._baseline_cost_usd is None and ctx.estimated_cost_usd is not None:
             self._baseline_cost_usd = ctx.estimated_cost_usd
 
+    def _record_baseline_from_batch(self, attempt_results: List[OptimizationContext]) -> None:
+        """Capture duration/cost baseline as the average across a GT batch.
+
+        Used by the GT optimization loop. The first attempt's N samples form
+        the baseline; averaging them gives a more stable reference than a
+        single sample and ensures comparisons in subsequent attempts reflect
+        the typical performance of the original configuration rather than an
+        outlier measurement.
+
+        Called once per run (subsequent calls are no-ops once both values are
+        set).
+
+        :param attempt_results: All completed sample contexts from the first
+            GT attempt.
+        """
+        if not attempt_results:
+            return
+        if self._baseline_duration_ms is None:
+            durations = [
+                ctx.duration_ms for ctx in attempt_results if ctx.duration_ms is not None
+            ]
+            if durations:
+                self._baseline_duration_ms = sum(durations) / len(durations)
+        if self._baseline_cost_usd is None:
+            costs = [
+                ctx.estimated_cost_usd
+                for ctx in attempt_results
+                if ctx.estimated_cost_usd is not None
+            ]
+            if costs:
+                self._baseline_cost_usd = sum(costs) / len(costs)
+
     def _build_agent_config_for_context(
         self, ctx: OptimizationContext, skip_interpolation: bool = False
     ) -> AIAgentConfig:
@@ -1272,7 +1306,7 @@ async def _run_ground_truth_optimization(
             # from all of the previous samples, then trim to one full attempt's worth so
             # judge prompts don't grow unboundedly across many failed attempts.
             if attempt_results:
-                self._record_baseline(attempt_results[0])
+                self._record_baseline_from_batch(attempt_results)
             self._history.extend(attempt_results)
             self._history = _trim_history(self._history, n)
 

packages/optimization/tests/test_client.py

@@ -1,5 +1,6 @@
 """Tests for OptimizationClient."""
 
+import dataclasses
 import json
 from typing import Any, Dict
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -5176,6 +5177,52 @@ def test_skips_gracefully_when_units_differ_across_model_switch(self):
         assert self.client._evaluate_cost(self._ctx(None)) is True
 
 
+# ---------------------------------------------------------------------------
+# _record_baseline_from_batch
+# ---------------------------------------------------------------------------
+
+
+class TestRecordBaselineFromBatch:
+    def setup_method(self):
+        self.client = _make_client()
+        self.client._initialize_class_members_from_config(_make_agent_config())
+
+    def _ctx(self, duration_ms=None, cost=None):
+        ctx = self.client._create_optimization_context(iteration=1, variables={})
+        return dataclasses.replace(ctx, duration_ms=duration_ms, estimated_cost_usd=cost)
+
+    def test_averages_duration_across_batch(self):
+        results = [self._ctx(duration_ms=1000), self._ctx(duration_ms=2000), self._ctx(duration_ms=3000)]
+        self.client._record_baseline_from_batch(results)
+        assert self.client._baseline_duration_ms == 2000.0
+
+    def test_averages_cost_across_batch(self):
+        results = [self._ctx(cost=0.01), self._ctx(cost=0.02), self._ctx(cost=0.03)]
+        self.client._record_baseline_from_batch(results)
+        assert abs(self.client._baseline_cost_usd - 0.02) < 1e-9
+
+    def test_skips_none_values_in_average(self):
+        results = [self._ctx(duration_ms=1000), self._ctx(duration_ms=None), self._ctx(duration_ms=3000)]
+        self.client._record_baseline_from_batch(results)
+        assert self.client._baseline_duration_ms == 2000.0
+
+    def test_noop_when_already_set(self):
+        self.client._baseline_duration_ms = 999.0
+        results = [self._ctx(duration_ms=1000), self._ctx(duration_ms=2000)]
+        self.client._record_baseline_from_batch(results)
+        assert self.client._baseline_duration_ms == 999.0
+
+    def test_noop_on_empty_list(self):
+        self.client._record_baseline_from_batch([])
+        assert self.client._baseline_duration_ms is None
+        assert self.client._baseline_cost_usd is None
+
+    def test_noop_when_all_values_none(self):
+        results = [self._ctx(duration_ms=None), self._ctx(duration_ms=None)]
+        self.client._record_baseline_from_batch(results)
+        assert self.client._baseline_duration_ms is None
+
+
 # ---------------------------------------------------------------------------
 # _apply_duration_gate
 # ---------------------------------------------------------------------------