diff --git a/packages/tracker/meshnet_tracker/cli.py b/packages/tracker/meshnet_tracker/cli.py index 5a5d728..5b08f0f 100644 --- a/packages/tracker/meshnet_tracker/cli.py +++ b/packages/tracker/meshnet_tracker/cli.py @@ -8,6 +8,8 @@ from .accounts import DEFAULT_ACCOUNTS_DB_PATH from .billing import DEFAULT_BILLING_DB_PATH from .server import TrackerServer, derive_relay_url_from_public_tracker_url +DEFAULT_REGISTRY_DB_PATH = "meshnet_registry.sqlite3" + def main() -> None: common = argparse.ArgumentParser(add_help=False) @@ -49,20 +51,34 @@ def main() -> None: f"(default: {DEFAULT_BILLING_DB_PATH}; ADR-0015)" ), ) - common.add_argument( - "--no-billing", - action="store_true", - help="Disable the USDT billing ledger", - ) - common.add_argument( - "--max-charge-per-request", - type=float, - default=None, - help=( - "Reject chat completion requests whose requested token limit would " - "cost more than this many USDT" - ), - ) + common.add_argument( + "--no-billing", + action="store_true", + help="Disable the USDT billing ledger", + ) + common.add_argument( + "--max-charge-per-request", + type=float, + default=None, + help=( + "Reject chat completion requests whose prompt plus requested completion " + "token bound would cost more than this many USDT" + ), + ) + common.add_argument( + "--registry-db", + default=DEFAULT_REGISTRY_DB_PATH, + metavar="PATH", + help=( + "SQLite database path for persisted strike/ban/reputation registry " + f"state (default: {DEFAULT_REGISTRY_DB_PATH})" + ), + ) + common.add_argument( + "--no-registry-contracts", + action="store_true", + help="Disable the local contract registry used for strike/ban/reputation enforcement", + ) common.add_argument( "--accounts-db", default=DEFAULT_ACCOUNTS_DB_PATH, @@ -149,6 +165,11 @@ def main() -> None: treasury = SolanaCustodialTreasury( args.solana_rpc_url, args.usdt_mint, args.treasury_keypair, ) + contracts = None + if not args.no_registry_contracts: + from meshnet_contracts import LocalSolanaContracts # type: ignore[import-not-found] + + contracts = LocalSolanaContracts(registry_db=args.registry_db) server = TrackerServer( host=args.host, port=args.port, @@ -158,9 +179,10 @@ def main() -> None: stats_db=getattr(args, "stats_db", None), relay_url=relay_url, enable_billing=not args.no_billing, - billing_db=None if args.no_billing else args.billing_db, - max_charge_per_request=args.max_charge_per_request, - accounts_db=None if args.no_accounts else args.accounts_db, + billing_db=None if args.no_billing else args.billing_db, + max_charge_per_request=args.max_charge_per_request, + contracts=contracts, + accounts_db=None if args.no_accounts else args.accounts_db, treasury=treasury, settle_period=args.settle_period, payout_threshold=args.payout_threshold, diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index eea2247..5c29f85 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -1293,15 +1293,43 @@ def _request_total_token_upper_bound(body: dict) -> int | None: def _billable_non_stream_tokens(payload: dict, request_body: dict) -> int: reported = _usage_total_tokens(payload) - if reported is None: - return 0 - billable = max(0, reported) upper_bound = _request_total_token_upper_bound(request_body) + if reported is None: + completion_estimate = _observed_non_stream_completion_tokens(payload) + if completion_estimate > 0: + billable = completion_estimate + (_estimate_prompt_tokens(request_body) or 0) + elif upper_bound is not None: + billable = upper_bound + else: + return 0 + return min(billable, upper_bound) if upper_bound is not None else billable + billable = max(0, reported) if upper_bound is not None: billable = min(billable, upper_bound) return billable +def _observed_non_stream_completion_tokens(payload: dict) -> int: + choices = payload.get("choices") + if not isinstance(choices, list): + return 0 + observed = 0 + for choice in choices: + if not isinstance(choice, dict): + continue + message = choice.get("message") + estimated: int | None = None + if isinstance(message, dict): + estimated = _estimate_text_tokens(message.get("content")) + if estimated is None: + estimated = _estimate_text_tokens(choice.get("text")) + if estimated is not None: + observed += estimated + elif choice: + observed += 1 + return observed + + def _observed_stream_tokens(payload: dict) -> int: choices = payload.get("choices") if not isinstance(choices, list): @@ -1881,7 +1909,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "code": "missing_token_limit", }}) return - estimated_charge = server.billing.price_for(model) * token_limit / 1000.0 + total_token_bound = _request_total_token_upper_bound(body) or token_limit + estimated_charge = server.billing.price_for(model) * total_token_bound / 1000.0 if estimated_charge > server.max_charge_per_request: self._send_json(402, {"error": { "message": ( diff --git a/tests/test_billing_ledger.py b/tests/test_billing_ledger.py index 9dcc5d7..26cb543 100644 --- a/tests/test_billing_ledger.py +++ b/tests/test_billing_ledger.py @@ -21,6 +21,7 @@ from meshnet_tracker.server import ( _billable_non_stream_tokens, _billable_stream_tokens, _estimate_prompt_tokens, + _observed_non_stream_completion_tokens, _observed_stream_tokens, ) @@ -99,6 +100,14 @@ def test_non_stream_billable_tokens_cap_usage_by_request_bound(): assert _billable_non_stream_tokens(payload, request) == 5 +def test_non_stream_billable_tokens_fallback_when_usage_missing(): + payload = {"choices": [{"message": {"role": "assistant", "content": "ok now"}}]} + request = {"messages": [{"role": "user", "content": "hello tracker"}], "max_tokens": 10} + + assert _observed_non_stream_completion_tokens(payload) == 2 + assert _billable_non_stream_tokens(payload, request) == 4 + + def test_stream_billable_tokens_allow_usage_to_lower_observed_count_only(): observed_payload = { "choices": [{ @@ -557,7 +566,12 @@ def test_proxy_chat_rejects_request_above_spend_cap_before_routing(): r.read() with pytest.raises(urllib.error.HTTPError) as exc_info: - _chat(tracker_url, api_key="capped-client", max_tokens=1000) + _chat( + tracker_url, + api_key="capped-client", + messages=[{"role": "user", "content": " ".join(["prompt"] * 1000)}], + max_tokens=1, + ) body = json.loads(exc_info.value.read()) assert exc_info.value.code == 402 assert body["error"]["code"] == "spend_cap_exceeded"