Merge be876a6c22 into d30e3e052d

headroom/proxy/output_savings.py

@@ -147,12 +147,23 @@ class BaselineModel:
             self.strata.setdefault(key, _Accum()).merge(acc)
         self.glob.merge(other.glob)
 
-    def lookup(self, key: str) -> tuple[float, float, int]:
+    def lookup(self, key: str, *, fall_back_to_global: bool = False) -> tuple[float, float, int]:
         """Return ``(mean, var, n)`` for *key* with hierarchical back-off.
 
-        Falls back by trimming trailing (least-specific) stratum fields, then
-        to the global mean. Back-off keeps the estimate defined for strata the
-        baseline never saw, at the cost of specificity.
+        Falls back by trimming trailing (least-specific) stratum fields, so a
+        stratum the baseline never saw is still scored against its nearest
+        observed neighbours. Returns ``(0.0, 0.0, 0)`` when even that finds
+        nothing -- callers already treat ``n == 0`` as "no evidence".
+
+        ``fall_back_to_global`` restores the old last resort of the
+        all-requests mean. It is off by default because that mean is not a
+        control for anything: the baseline is seeded once, from whatever the
+        user ran *before* installing, so every model family they adopt later
+        resolves to it. On a real ledger that meant 48% of requests -- sonnet
+        and fable turns whose replies average 43-770 tokens -- being scored
+        against one opus-derived mean of 1,083, which alone produced 74% of the
+        reported savings. Scoring a short no-tool ask against a long
+        tool-calling turn is not a synthetic control, it is a unit conversion.
         """
         acc = self.strata.get(key)
         if acc is not None and acc.n > 0:
@@ -160,11 +171,16 @@ class BaselineModel:
         parts = key.split("|")
         while len(parts) > 1:
             parts = parts[:-1]
-            prefix = "|".join(parts)
+            prefix = "|".join(parts) + "|"
+            neighbours = _Accum()
             for k, a in self.strata.items():
-                if k.startswith(prefix + "|") and a.n > 0:
-                    return a.mean, a.var, a.n
-        return self.glob.mean, self.glob.var, self.glob.n
+                if a.n > 0 and k.startswith(prefix):
+                    neighbours.merge(a)
+            if neighbours.n > 0:
+                return neighbours.mean, neighbours.var, neighbours.n
+        if fall_back_to_global:
+            return self.glob.mean, self.glob.var, self.glob.n
+        return 0.0, 0.0, 0
 
     def to_dict(self) -> dict[str, Any]:
         return {

tests/test_output_savings.py

@@ -178,14 +178,38 @@ class TestBaselineModel:
         assert mean == 500.0
         assert n == 1
 
-    def test_lookup_falls_back_to_global(self):
+    def test_an_unobserved_family_is_not_scored_against_the_global_mean(self):
+        # The baseline is seeded once, from what the user ran before installing.
+        # Every family adopted later lands here, and the all-requests mean is
+        # not a control for any of them.
         m = BaselineModel()
         m.observe("opus|a|s|tools", 100)
         m.observe("sonnet|b|m|notools", 300)
-        mean, _, n = m.lookup("gpt|totally|xl|tools")
+        assert m.lookup("gpt|totally|xl|tools") == (0.0, 0.0, 0)
+
+    def test_the_global_mean_remains_available_on_request(self):
+        m = BaselineModel()
+        m.observe("opus|a|s|tools", 100)
+        m.observe("sonnet|b|m|notools", 300)
+        mean, _, n = m.lookup("gpt|totally|xl|tools", fall_back_to_global=True)
         assert mean == 200.0  # global mean of 100 and 300
         assert n == 2
 
+    def test_prefix_backoff_merges_every_neighbour_not_the_first_one_hashed(self):
+        # Taking the first matching stratum in dict order made the answer
+        # depend on insertion order, so a round-tripped ledger could score the
+        # same request differently.
+        m = BaselineModel()
+        m.observe("opus|ask|l|tools", 1000)
+        m.observe("opus|ask|l|notools", 100)
+        mean, _, n = m.lookup("opus|ask|l|other")
+        assert (mean, n) == (550.0, 2)
+
+        reversed_order = BaselineModel()
+        reversed_order.observe("opus|ask|l|notools", 100)
+        reversed_order.observe("opus|ask|l|tools", 1000)
+        assert reversed_order.lookup("opus|ask|l|other") == m.lookup("opus|ask|l|other")
+
     def test_roundtrip_serialization(self):
         m = BaselineModel()
         for v in (10, 20, 30):
@@ -277,6 +301,28 @@ class TestEstimateFromBaseline:
 # ---------------------------------------------------------------------------
 
 
+class TestEstimateExcludesUnobservedStrata:
+    def test_requests_without_baseline_evidence_are_left_out(self):
+        ledger = SavingsLedger()
+        for _ in range(10):
+            ledger.baseline.observe("opus|ask|l|tools", 1000)
+            ledger.record("treatment", "opus|ask|l|tools", 800)
+        # A family the baseline never saw, with far shorter replies. Scoring it
+        # against the global mean would credit ~950 saved tokens per request.
+        for _ in range(40):
+            ledger.record("treatment", "sonnet|new_user_ask|m|notools", 50)
+
+        est = ledger.estimate_from_baseline()
+        assert est.n_requests == 10, "only the observed stratum is scored"
+        assert abs(est.tokens_saved - 2000) < 1e-6  # 10 * (1000 - 800)
+        assert abs(est.pct - 20.0) < 1e-6
+
+    def test_per_request_savings_are_zero_without_evidence(self):
+        ledger = SavingsLedger()
+        ledger.baseline.observe("opus|ask|l|tools", 1000)
+        assert ledger.baseline.lookup("sonnet|new_user_ask|m|notools") == (0.0, 0.0, 0)
+
+
 class TestEstimateFromHoldout:
     def test_none_without_control_data(self):
         ledger = SavingsLedger()