fix: remove unnecessary token path

Author: andrewklatzke

packages/optimization/src/ldai_optimizer/client.py

@@ -791,34 +791,21 @@ async def _evaluate_acceptance_judge(
                 else None
             )
             if current_cost is not None:
-                has_pricing = (
-                    _find_model_config(self._current_model or "", self._model_configs) or {}
-                ).get("costPerInputToken") is not None
-                if has_pricing:
-                    cost_str = f"${current_cost:.6f}"
-                else:
-                    cost_str = f"{int(current_cost)} tokens"
                 instructions += (
                     f"\n\nThe acceptance criteria for this judge includes a cost/token-usage goal. "
                 )
                 if agent_usage is not None:
                     instructions += (
                         f"The agent's response used {agent_usage.input} input tokens "
                         f"and {agent_usage.output} output tokens "
-                        f"(estimated cost: {cost_str}). "
+                        f"(estimated cost: ${current_cost:.6f}). "
                     )
                 if baseline_cost is not None:
                     delta = current_cost - baseline_cost
                     direction = "less" if delta < 0 else "more"
-                    if has_pricing:
-                        baseline_str = f"${baseline_cost:.6f}"
-                        delta_str = f"${abs(delta):.6f}"
-                    else:
-                        baseline_str = f"{int(baseline_cost)} tokens"
-                        delta_str = f"{int(abs(delta))} tokens"
                     instructions += (
-                        f"The baseline cost (first iteration) was {baseline_str}. "
-                        f"This response cost {delta_str} {direction} than the baseline. "
+                        f"The baseline cost (first iteration) was ${baseline_cost:.6f}. "
+                        f"This response cost ${abs(delta):.6f} {direction} than the baseline. "
                     )
                 instructions += (
                     "In your rationale, state the token usage and cost, and any change from baseline. "

packages/optimization/src/ldai_optimizer/prompts.py

@@ -18,7 +18,7 @@
 
 _COST_KEYWORDS = re.compile(
     r"\b(cheap|cheaper|cheapest|costs?|costly|expensive|budget|affordable|"
-    r"tokens?|spend|spending|economical|cost-effective|frugal|"
+    r"spend|spending|economical|cost-effective|frugal|"
     r"price|pricing|bill|billing)\b",
     re.IGNORECASE,
 )

packages/optimization/src/ldai_optimizer/util.py

@@ -322,37 +322,34 @@ def estimate_cost(
     usage: Optional["TokenUsage"],
     model_config: Optional[Dict[str, Any]],
 ) -> Optional[float]:
-    """Estimate the monetary cost of a single agent call.
+    """Estimate the monetary cost of a single agent call in USD.
 
-    Uses ``costPerInputToken`` and ``costPerOutputToken`` from the model config
-    when available.  If the model config has no pricing fields, falls back to
-    returning the raw total token count as a dimensionless proxy so the cost
-    gate can still operate comparatively.  Returns ``None`` only when ``usage``
-    itself is ``None``.
+    Uses ``costPerInputToken`` and ``costPerOutputToken`` from the model config.
+    Returns ``None`` when either ``usage`` is ``None`` or no pricing fields are
+    present on the model config — ensuring the return value is always in USD or
+    absent, never a raw token count. This prevents unit-mismatch bugs when
+    comparing costs across iterations where the model (and its pricing
+    availability) may differ.
 
     ``costPerCachedInputToken`` is intentionally ignored — the estimate uses
-    input/output tokens only, which is sufficient for relative comparison
-    across optimization iterations.
+    input/output tokens only.
 
     :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 raw total token count as proxy, or ``None``.
+    :return: Estimated cost in USD, or ``None`` if usage or pricing data is absent.
     """
     if usage is None:
         return None
 
     input_price = model_config.get("costPerInputToken") if model_config else None
     output_price = model_config.get("costPerOutputToken") if model_config else None
 
-    if input_price is not None or output_price is not None:
-        cost = 0.0
-        if input_price is not None and usage.input is not None:
-            cost += usage.input * input_price
-        if output_price is not None and usage.output is not None:
-            cost += usage.output * output_price
-        return cost
+    if input_price is None and output_price is None:
+        return None
 
-    logger.debug(
-        "No pricing data on model config for cost estimation; falling back to total token count"
-    )
-    return float(usage.total or 0)
+    cost = 0.0
+    if input_price is not None and usage.input is not None:
+        cost += usage.input * input_price
+    if output_price is not None and usage.output is not None:
+        cost += usage.output * output_price
+    return cost

packages/optimization/tests/test_client.py

@@ -535,7 +535,7 @@ async def test_duration_context_added_to_instructions_when_latency_keyword_prese
         )
         _, config, _, _ = self.handle_judge_call.call_args.args
         assert "1500ms" in config.instructions
-        assert "mention the duration" in config.instructions
+        assert "state the duration" in config.instructions
 
     async def test_duration_context_includes_baseline_comparison_when_history_present(self):
         """When history[0] has a duration, the judge instructions include a baseline comparison."""
@@ -1842,11 +1842,11 @@ async def test_apply_variation_response_calls_restore_and_logs_warning(self):
 
         with patch("ldai_optimizer.client.logger") as mock_logger:
             await client._generate_new_variation(iteration=1, variables={})
-            warning_calls = [
-                call for call in mock_logger.warning.call_args_list
+            debug_calls = [
+                call for call in mock_logger.debug.call_args_list
                 if "user-123" in str(call) or "business" in str(call)
             ]
-            assert len(warning_calls) >= 1
+            assert len(debug_calls) >= 1
 
         assert "{{user_id}}" in client._current_instructions
         assert "user-123" not in client._current_instructions
@@ -4727,15 +4727,13 @@ def test_uses_only_output_price_when_input_absent(self):
         cost = estimate_cost(usage, model_config)
         assert cost == pytest.approx(40 * 0.002)
 
-    def test_falls_back_to_total_token_count_when_no_pricing(self):
+    def test_returns_none_when_no_pricing_in_config(self):
         usage = self._usage(total=100)
-        cost = estimate_cost(usage, {})
-        assert cost == 100.0
+        assert estimate_cost(usage, {}) is None
 
-    def test_falls_back_to_total_token_count_when_model_config_none(self):
+    def test_returns_none_when_model_config_none(self):
         usage = self._usage(total=250)
-        cost = estimate_cost(usage, None)
-        assert cost == 250.0
+        assert estimate_cost(usage, None) is None
 
     def test_ignores_cached_input_token_price(self):
         usage = self._usage(total=100, inp=60, out=40)
@@ -4783,8 +4781,8 @@ def test_detects_costs_plural(self):
     def test_detects_budget(self):
         assert _acceptance_criteria_implies_cost_optimization(self._judge("Stay within budget."))
 
-    def test_detects_tokens(self):
-        assert _acceptance_criteria_implies_cost_optimization(self._judge("Use fewer tokens."))
+    def test_does_not_detect_token_to_avoid_false_positives(self):
+        assert not _acceptance_criteria_implies_cost_optimization(self._judge("Generate a valid authentication token."))
 
     def test_detects_billing(self):
         assert _acceptance_criteria_implies_cost_optimization(self._judge("Minimize billing."))
@@ -4803,7 +4801,7 @@ def test_no_match_on_unrelated_statement(self):
     def test_multiple_judges_one_matches(self):
         judges = {
             "j1": OptimizationJudge(threshold=0.9, acceptance_statement="Be accurate."),
-            "j2": OptimizationJudge(threshold=0.9, acceptance_statement="Use fewer tokens."),
+            "j2": OptimizationJudge(threshold=0.9, acceptance_statement="Keep costs low."),
         }
         assert _acceptance_criteria_implies_cost_optimization(judges)
 
@@ -4861,11 +4859,11 @@ def test_skips_gracefully_when_candidate_cost_none(self):
         ctx = self._ctx(None)  # type: ignore[arg-type]
         assert self.client._evaluate_cost(ctx) is True
 
-    def test_works_with_token_count_proxy(self):
-        # When no pricing data, cost is raw token count — gate still compares numerically
-        self._seed_history(1000.0)
-        assert self.client._evaluate_cost(self._ctx(750.0)) is True
-        assert self.client._evaluate_cost(self._ctx(900.0)) is False
+    def test_skips_gracefully_when_units_differ_across_model_switch(self):
+        # If baseline was captured with pricing (USD) but candidate has no pricing,
+        # candidate cost is None and the gate skips rather than comparing incompatible units.
+        self._seed_history(0.010)
+        assert self.client._evaluate_cost(self._ctx(None)) is True
 
 
 # ---------------------------------------------------------------------------
@@ -4986,10 +4984,18 @@ def setup_method(self):
     def _cost_judge(self) -> OptimizationJudge:
         return OptimizationJudge(
             threshold=0.9,
-            acceptance_statement="Use fewer tokens and keep costs low.",
+            acceptance_statement="Keep costs low and stay within budget.",
         )
 
+    def _set_pricing(self):
+        """Give the client a model config with pricing so estimate_cost returns USD."""
+        self.client._current_model = "gpt-4o"
+        self.client._model_configs = [
+            {"id": "gpt-4o", "costPerInputToken": 0.000005, "costPerOutputToken": 0.000015}
+        ]
+
     async def test_cost_context_injected_into_instructions(self):
+        self._set_pricing()
         usage = TokenUsage(total=100, input=60, output=40)
         captured: list = []
 
@@ -5040,6 +5046,7 @@ async def _capture_judge_call(judge_key, judge_config, ctx, is_judge):
         assert "cost/token-usage goal" not in instructions
 
     async def test_baseline_cost_shown_when_history_present(self):
+        self._set_pricing()
         usage = TokenUsage(total=100, input=60, output=40)
         captured: list = []