89 lines
2.9 KiB
TypeScript
89 lines
2.9 KiB
TypeScript
/**
|
|
* qwencloud-usage: Session-cost tracker for QwenCloud models.
|
|
*
|
|
* QwenCloud Token Plan is a prepaid credits subscription without a public
|
|
* usage/balance API. This extension accumulates cost from pi's built-in
|
|
* Usage tracking (every assistant message carries token counts + cost) and
|
|
* displays a running session total in the footer.
|
|
*
|
|
* Set QWENCLOUD_MONTHLY_BUDGET (dollars) to enable a coloured
|
|
* budget-usage percentage. For example:
|
|
* QWENCLOUD_MONTHLY_BUDGET=6 # Lite plan ($6/mo)
|
|
* QWENCLOUD_MONTHLY_BUDGET=18 # Standard ($18/mo)
|
|
* QWENCLOUD_MONTHLY_BUDGET=68 # Pro ($68/mo)
|
|
*
|
|
* The extension sets the "qwencloud-usage" footer-status key and emits
|
|
* "usage:update" so prompt-frame (or any other listener) can refresh.
|
|
*/
|
|
|
|
const STATUS_KEY = "qwencloud-usage";
|
|
const PROVIDER = "qw";
|
|
|
|
function budgetPct(cost: number, budget: number): number | null {
|
|
if (!Number.isFinite(budget) || budget <= 0) return null;
|
|
return Math.round((cost / budget) * 100);
|
|
}
|
|
|
|
function costColor(pct: number | null): "success" | "warning" | "error" {
|
|
if (pct === null) return "success";
|
|
if (pct >= 90) return "error";
|
|
if (pct >= 70) return "warning";
|
|
return "success";
|
|
}
|
|
|
|
export default function (pi: any) {
|
|
const budget = parseFloat(process.env.QWENCLOUD_MONTHLY_BUDGET || "0");
|
|
|
|
let sessionCost = 0;
|
|
let sessionInput = 0;
|
|
let sessionOutput = 0;
|
|
|
|
pi.on("message_end", (event: any, ctx: any) => {
|
|
const msg = event.message;
|
|
if (msg?.model?.provider !== PROVIDER) return;
|
|
if (!msg.usage) return;
|
|
|
|
sessionCost += msg.usage.cost?.total ?? 0;
|
|
sessionInput += msg.usage.input ?? 0;
|
|
sessionOutput += msg.usage.output ?? 0;
|
|
|
|
if (!ctx?.hasUI) return;
|
|
|
|
const theme = ctx.ui.theme;
|
|
const pct = budgetPct(sessionCost, budget);
|
|
const color = costColor(pct);
|
|
const totalTokens = sessionInput + sessionOutput;
|
|
|
|
// Build a compact status string: qw · $0.42 · 7% · 12.3k tok
|
|
const parts: string[] = [];
|
|
parts.push(theme.fg("dim", "qw"));
|
|
parts.push(theme.fg(color, `$${sessionCost.toFixed(2)}`));
|
|
if (pct !== null) {
|
|
parts.push(theme.fg(color, `${pct}%`));
|
|
}
|
|
parts.push(theme.fg("dim", `${formatTokens(totalTokens)} tok`));
|
|
|
|
ctx.ui.setStatus(STATUS_KEY, parts.join(" " + theme.fg("muted", "·") + " "));
|
|
pi.events.emit("usage:update");
|
|
});
|
|
|
|
// Reset counters at the start of every new conversation.
|
|
pi.on("session_start", () => {
|
|
sessionCost = 0;
|
|
sessionInput = 0;
|
|
sessionOutput = 0;
|
|
});
|
|
|
|
// Clean up when the session ends so the status doesn't linger.
|
|
pi.on("session_shutdown", (_event: any, ctx: any) => {
|
|
if (ctx?.hasUI) ctx.ui.setStatus(STATUS_KEY, undefined);
|
|
});
|
|
}
|
|
|
|
/** Human-readable token count (e.g. "12.3k", "1.2M"). */
|
|
function formatTokens(n: number): string {
|
|
if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`;
|
|
if (n >= 1_000) return `${(n / 1_000).toFixed(1)}k`;
|
|
return String(n);
|
|
}
|