compass-wallet-usage

Centralized wallet usage module (Python) for multi-agent Compass systems.

What this module provides

• Per-user, per-IST-day usage recording (total_tokens, request count, optional agent_id split)
• Daily quota enforcement from env (WALLET_DAILY_TOKEN_LIMIT)
• Pluggable storage backends:
  • InMemoryWalletStorage (dev/tests)
  • FirestoreWalletStorage (prod)
• ADK callback factory (record_tokens_after_model) for Gemini usage_metadata
• Built-in ADK user-id resolver (resolve_user_id) for common callback context shapes
• Optional FastAPI router exposing wallet status + record routes

Install (local repo)

pip install -e .

Quick start

from wallet import WalletService
from storage import InMemoryWalletStorage

wallet = WalletService(InMemoryWalletStorage())
wallet.record("user-123", 210, agent_id="mktg_banner_agent")
print(wallet.status("user-123"))

Quota

Set in host service env:

WALLET_DAILY_TOKEN_LIMIT=9000

Then call before model request:

wallet.enforce_daily_quota_if_configured(user_id)

ADK callback usage

from adk_callbacks import record_tokens_after_model

after_model_callback = record_tokens_after_model(
    wallet=wallet,
    agent_id="mktg_banner_agent",
)

record_tokens_after_model defaults to resolve_user_id, which checks:

• callback_context.state["user_id" | "uid" | "firebase_uid" | "user_email"]
• then callback_context.user_id / callback_context.uid
• fallback: "unknown-user"

Optional FastAPI routes

from fastapi import FastAPI
from api import build_wallet_router

app = FastAPI()
app.include_router(build_wallet_router(wallet))

Routes:

• GET /wallet/status?user_id=...
• POST /wallet/record

Notes

• This module is intentionally standalone and not integrated into other repos yet.
• Firestore backend requires google-cloud-firestore and ADC/service credentials.