dual billing; tracker to node model sharing

packages/tracker/meshnet_tracker/billing.py

@@ -24,6 +24,18 @@ DEFAULT_BILLING_DB_PATH = "billing.sqlite"
 NODE_REVENUE_SHARE = 0.90  # nodes 90% / protocol cut 10% (ADR-0015)
 
 
+def _normalize_rates(value: "float | tuple[float, float] | dict") -> tuple[float, float]:
+    """Coerce a price spec into an (input_per_1k, output_per_1k) pair."""
+    if isinstance(value, dict):
+        base = value.get("price")
+        inp = value.get("input", base)
+        out = value.get("output", base)
+        return (float(inp), float(out))
+    if isinstance(value, (tuple, list)):
+        return (float(value[0]), float(value[1]))
+    return (float(value), float(value))
+
+
 class BillingLedger:
     """Thread-safe USDT ledger with SQLite persistence and event replication."""
 
@@ -33,13 +45,17 @@ class BillingLedger:
         self,
         db_path: str | None = None,
         *,
-        prices: dict[str, float] | None = None,
+        prices: dict[str, float | tuple[float, float]] | None = None,
         default_price_per_1k: float = DEFAULT_PRICE_PER_1K_TOKENS,
         starting_credit: float = DEFAULT_STARTING_CREDIT,
         node_share: float = NODE_REVENUE_SHARE,
     ) -> None:
         self._db_path = db_path
-        self._prices = dict(prices) if prices else {}
+        # US-045: per-model (input_per_1k, output_per_1k). A bare float in
+        # ``prices`` sets both rates (legacy single-rate models).
+        self._prices: dict[str, tuple[float, float]] = {
+            model: _normalize_rates(value) for model, value in (prices or {}).items()
+        }
         self._default_price_per_1k = default_price_per_1k
         self._starting_credit = starting_credit
         self._node_share = node_share
@@ -62,11 +78,23 @@ class BillingLedger:
     # ---- pricing ----
 
     def price_for(self, model: str) -> float:
-        return self._prices.get(model, self._default_price_per_1k)
+        """Blended (average) per-1k rate — kept for estimators and history logs."""
+        rates = self._prices.get(model)
+        if rates is None:
+            return self._default_price_per_1k
+        return (rates[0] + rates[1]) / 2.0
+
+    def prices_for(self, model: str) -> tuple[float, float]:
+        """(input_per_1k, output_per_1k) for a model (US-045)."""
+        return self._prices.get(model, (self._default_price_per_1k, self._default_price_per_1k))
 
     def set_price(self, model: str, price_per_1k: float) -> None:
+        """Legacy single-rate setter — applies the same rate to input and output."""
+        self.set_prices(model, price_per_1k, price_per_1k)
+
+    def set_prices(self, model: str, input_per_1k: float, output_per_1k: float) -> None:
         with self._lock:
-            self._prices[model] = price_per_1k
+            self._prices[model] = (float(input_per_1k), float(output_per_1k))
 
     # ---- local operations (create + apply + log an event) ----
 
@@ -130,9 +158,17 @@ class BillingLedger:
         model: str,
         total_tokens: int,
         node_work: list[tuple[str | None, int]],
+        *,
+        input_tokens: int | None = None,
+        output_tokens: int | None = None,
     ) -> dict:
         """Debit the client and split the fee 90/10.
 
+        With ``input_tokens``/``output_tokens`` (US-045) the cost is
+        ``input·in_rate + output·out_rate``; without them, legacy behavior —
+        ``total_tokens`` at the blended rate. Replayed/gossiped events apply
+        their recorded ``cost``, so old events are unaffected either way.
+
         ``node_work`` is ``[(wallet_address | None, work_units), ...]`` for the
         nodes that served the request. Work units of nodes without a wallet
         accrue to the protocol cut — there is nowhere to pay them out.
@@ -140,7 +176,14 @@ class BillingLedger:
         request is then rejected by ``has_funds`` (post-pay drift, standard
         metered-billing behavior).
         """
-        cost = self.price_for(model) * max(0, total_tokens) / 1000.0
+        if input_tokens is not None or output_tokens is not None:
+            in_rate, out_rate = self.prices_for(model)
+            in_tokens = max(0, input_tokens or 0)
+            out_tokens = max(0, output_tokens or 0)
+            cost = (in_tokens * in_rate + out_tokens * out_rate) / 1000.0
+            total_tokens = in_tokens + out_tokens
+        else:
+            cost = self.price_for(model) * max(0, total_tokens) / 1000.0
         total_work = sum(max(0, w) for _, w in node_work)
         node_pool = cost * self._node_share
         shares: dict[str, float] = {}
@@ -165,6 +208,9 @@ class BillingLedger:
                 "api_key": api_key,
                 "model": model,
                 "total_tokens": total_tokens,
+                **({"input_tokens": max(0, input_tokens or 0),
+                    "output_tokens": max(0, output_tokens or 0)}
+                   if (input_tokens is not None or output_tokens is not None) else {}),
                 "cost": cost,
                 "shares": shares,
                 "protocol_amount": protocol_amount,