Add /usage endpoint with token tracking, live balances, and cost calculations
This commit is contained in:
William Stuckey
@@ -1,5 +1,5 @@
|
||||
import { createServer } from "node:http";
|
||||
import { readFileSync, existsSync } from "node:fs";
|
||||
import { readFileSync, existsSync, writeFileSync, renameSync, mkdirSync } from "node:fs";
|
||||
import { join, dirname } from "node:path";
|
||||
import { fileURLToPath } from "node:url";
|
||||
|
||||
@@ -8,6 +8,7 @@ const __dirname = dirname(fileURLToPath(import.meta.url));
|
||||
// --- Configuration -----------------------------------------------------------
|
||||
|
||||
const CONFIG_PATH = join(__dirname, "config.json");
|
||||
const USAGE_DATA_PATH = join(__dirname, "usage-data.json");
|
||||
|
||||
function loadConfig() {
|
||||
if (!existsSync(CONFIG_PATH)) {
|
||||
@@ -22,6 +23,13 @@ const config = loadConfig();
|
||||
const PORT = config.port || 3100;
|
||||
const AUTH_TOKEN = config.authToken;
|
||||
|
||||
// Default pricing if not in config
|
||||
const DEFAULT_PRICING = {
|
||||
anthropic: { inputPerMTok: 15, outputPerMTok: 75 }, // Opus rates
|
||||
deepseek: { inputPerMTok: 0.27, outputPerMTok: 1.10 },
|
||||
zai: { type: "flat", monthlyCostUsd: 10 },
|
||||
};
|
||||
|
||||
// Provider definitions: maps provider name -> { baseUrl, apiKey, authStyle }
|
||||
// authStyle: "anthropic" sends x-api-key header, "bearer" sends Authorization: Bearer
|
||||
const PROVIDERS = {
|
||||
@@ -69,6 +77,248 @@ const MODEL_ROUTES = {
|
||||
// Default route for unknown models
|
||||
const DEFAULT_ROUTE = { provider: "zai", rewriteModel: "glm-4.7" };
|
||||
|
||||
// --- Usage Tracking -----------------------------------------------------------
|
||||
|
||||
function getCurrentPeriod() {
|
||||
const now = new Date();
|
||||
return `${now.getFullYear()}-${String(now.getMonth() + 1).padStart(2, "0")}`;
|
||||
}
|
||||
|
||||
function loadUsageData() {
|
||||
if (!existsSync(USAGE_DATA_PATH)) {
|
||||
return {
|
||||
version: 1,
|
||||
lastUpdated: new Date().toISOString(),
|
||||
models: {},
|
||||
};
|
||||
}
|
||||
try {
|
||||
return JSON.parse(readFileSync(USAGE_DATA_PATH, "utf-8"));
|
||||
} catch (err) {
|
||||
console.error(`[usage] Failed to load usage data: ${err.message}`);
|
||||
return {
|
||||
version: 1,
|
||||
lastUpdated: new Date().toISOString(),
|
||||
models: {},
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
function saveUsageData(data) {
|
||||
try {
|
||||
const tempPath = USAGE_DATA_PATH + ".tmp";
|
||||
writeFileSync(tempPath, JSON.stringify(data, null, 2), "utf-8");
|
||||
renameSync(tempPath, USAGE_DATA_PATH);
|
||||
} catch (err) {
|
||||
console.error(`[usage] Failed to save usage data: ${err.message}`);
|
||||
}
|
||||
}
|
||||
|
||||
function recordUsage(requestedModel, provider, actualModel, inputTokens, outputTokens) {
|
||||
const usageData = loadUsageData();
|
||||
const currentPeriod = getCurrentPeriod();
|
||||
|
||||
if (!usageData.models[requestedModel]) {
|
||||
usageData.models[requestedModel] = {
|
||||
provider,
|
||||
actualModel,
|
||||
total: { inputTokens: 0, outputTokens: 0, requests: 0 },
|
||||
monthly: { period: currentPeriod, inputTokens: 0, outputTokens: 0, requests: 0 },
|
||||
};
|
||||
}
|
||||
|
||||
const modelData = usageData.models[requestedModel];
|
||||
|
||||
// Check if monthly period has rolled over
|
||||
if (modelData.monthly.period !== currentPeriod) {
|
||||
modelData.monthly = { period: currentPeriod, inputTokens: 0, outputTokens: 0, requests: 0 };
|
||||
}
|
||||
|
||||
// Update totals
|
||||
modelData.total.inputTokens += inputTokens || 0;
|
||||
modelData.total.outputTokens += outputTokens || 0;
|
||||
modelData.total.requests += 1;
|
||||
|
||||
// Update monthly
|
||||
modelData.monthly.inputTokens += inputTokens || 0;
|
||||
modelData.monthly.outputTokens += outputTokens || 0;
|
||||
modelData.monthly.requests += 1;
|
||||
|
||||
usageData.lastUpdated = new Date().toISOString();
|
||||
saveUsageData(usageData);
|
||||
}
|
||||
|
||||
// --- Provider Balance Fetching ----------------------------------------------
|
||||
|
||||
async function fetchDeepSeekBalance(apiKey) {
|
||||
try {
|
||||
const res = await fetch("https://api.deepseek.com/user/balance", {
|
||||
headers: { "Authorization": `Bearer ${apiKey}` },
|
||||
signal: AbortSignal.timeout(5000),
|
||||
});
|
||||
if (res.ok) {
|
||||
const data = await res.json();
|
||||
if (data.balance_infos && data.balance_infos.length > 0) {
|
||||
const balance = data.balance_infos.find((b) => b.currency === "USD") || data.balance_infos[0];
|
||||
return {
|
||||
totalBalance: balance.total_balance || balance.total_granted_balance,
|
||||
currency: balance.currency || "USD",
|
||||
};
|
||||
}
|
||||
}
|
||||
return null;
|
||||
} catch (err) {
|
||||
console.error(`[usage] Failed to fetch DeepSeek balance: ${err.message}`);
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
async function fetchAnthropicBalance(apiKey) {
|
||||
// Anthropic doesn't have a public balance API for regular API keys
|
||||
// Admin API exists but requires a different key type and org access
|
||||
return null;
|
||||
}
|
||||
|
||||
async function fetchZaiBalance(apiKey) {
|
||||
// Z.ai doesn't have a documented balance API; flat-rate subscription
|
||||
return null;
|
||||
}
|
||||
|
||||
async function fetchProviderBalances() {
|
||||
const balances = {};
|
||||
const promises = [];
|
||||
|
||||
if (PROVIDERS.deepseek?.apiKey) {
|
||||
promises.push(
|
||||
fetchDeepSeekBalance(PROVIDERS.deepseek.apiKey).then((b) => {
|
||||
if (b) balances.deepseek = b;
|
||||
})
|
||||
);
|
||||
}
|
||||
if (PROVIDERS.anthropic?.apiKey) {
|
||||
promises.push(
|
||||
fetchAnthropicBalance(PROVIDERS.anthropic.apiKey).then((b) => {
|
||||
if (b) balances.anthropic = b;
|
||||
})
|
||||
);
|
||||
}
|
||||
if (PROVIDERS.zai?.apiKey) {
|
||||
promises.push(
|
||||
fetchZaiBalance(PROVIDERS.zai.apiKey).then((b) => {
|
||||
if (b) balances.zai = b;
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
await Promise.allSettled(promises);
|
||||
return balances;
|
||||
}
|
||||
|
||||
// --- Cost Calculation --------------------------------------------------------
|
||||
|
||||
function calculateCost(provider, inputTokens, outputTokens) {
|
||||
const pricing = config.pricing?.[provider] || DEFAULT_PRICING[provider];
|
||||
if (!pricing || pricing.type === "flat") {
|
||||
return null;
|
||||
}
|
||||
const inputCost = ((inputTokens || 0) / 1_000_000) * pricing.inputPerMTok;
|
||||
const outputCost = ((outputTokens || 0) / 1_000_000) * pricing.outputPerMTok;
|
||||
return inputCost + outputCost;
|
||||
}
|
||||
|
||||
// --- Usage Data Aggregation -------------------------------------------------
|
||||
|
||||
function buildUsageResponse() {
|
||||
const usageData = loadUsageData();
|
||||
const balances = fetchProviderBalances(); // Start fetching but don't await
|
||||
const currentPeriod = getCurrentPeriod();
|
||||
|
||||
const providers = {
|
||||
anthropic: { balance: null, models: {} },
|
||||
deepseek: { balance: null, models: {} },
|
||||
zai: { balance: null, models: {} },
|
||||
};
|
||||
|
||||
let totalRequests = 0;
|
||||
let monthlyRequests = 0;
|
||||
let totalCost = 0;
|
||||
let monthlyCost = 0;
|
||||
let flatRateModels = 0;
|
||||
|
||||
// Process model data
|
||||
for (const [requestedModel, modelData] of Object.entries(usageData.models || {})) {
|
||||
const providerName = modelData.provider;
|
||||
if (!providers[providerName]) continue;
|
||||
|
||||
const { total, monthly, actualModel } = modelData;
|
||||
|
||||
providers[providerName].models[requestedModel] = {
|
||||
routesTo: actualModel || requestedModel,
|
||||
total: {
|
||||
inputTokens: total.inputTokens,
|
||||
outputTokens: total.outputTokens,
|
||||
requests: total.requests,
|
||||
costUsd: calculateCost(providerName, total.inputTokens, total.outputTokens),
|
||||
},
|
||||
monthly: {
|
||||
period: monthly.period,
|
||||
inputTokens: monthly.inputTokens,
|
||||
outputTokens: monthly.outputTokens,
|
||||
requests: monthly.requests,
|
||||
costUsd: calculateCost(providerName, monthly.inputTokens, monthly.outputTokens),
|
||||
},
|
||||
};
|
||||
|
||||
totalRequests += total.requests;
|
||||
monthlyRequests += monthly.requests;
|
||||
|
||||
const totalCostVal = calculateCost(providerName, total.inputTokens, total.outputTokens);
|
||||
const monthlyCostVal = calculateCost(providerName, monthly.inputTokens, monthly.outputTokens);
|
||||
|
||||
if (totalCostVal !== null) totalCost += totalCostVal;
|
||||
else flatRateModels++;
|
||||
|
||||
if (monthlyCostVal !== null) monthlyCost += monthlyCostVal;
|
||||
}
|
||||
|
||||
// Calculate pricing per provider
|
||||
for (const [providerName, providerData] of Object.entries(providers)) {
|
||||
const pricing = config.pricing?.[providerName] || DEFAULT_PRICING[providerName];
|
||||
if (pricing?.type === "flat") {
|
||||
providerData.pricing = "flat_rate";
|
||||
providerData.monthlyCostUsd = pricing.monthlyCostUsd;
|
||||
} else {
|
||||
providerData.pricing = {
|
||||
inputPerMTok: pricing?.inputPerMTok,
|
||||
outputPerMTok: pricing?.outputPerMTok,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// Build note
|
||||
const notes = [];
|
||||
if (flatRateModels > 0) {
|
||||
notes.push(`Cost excludes ${flatRateModels} flat-rate model(s). Z.ai costs $10/mo regardless of usage.`);
|
||||
}
|
||||
if (!balances.anthropic) {
|
||||
notes.push("Anthropic has no public balance API — check console.anthropic.com");
|
||||
}
|
||||
if (!balances.zai) {
|
||||
notes.push("Z.ai is flat-rate; no balance API available");
|
||||
}
|
||||
|
||||
return {
|
||||
usageData,
|
||||
providers,
|
||||
totalRequests,
|
||||
monthlyRequests,
|
||||
totalCost,
|
||||
monthlyCost,
|
||||
currentPeriod,
|
||||
notes,
|
||||
};
|
||||
}
|
||||
|
||||
// --- Helpers -----------------------------------------------------------------
|
||||
|
||||
function resolveRoute(model) {
|
||||
@@ -109,6 +359,44 @@ function buildUpstreamHeaders(provider, originalHeaders) {
|
||||
|
||||
// --- Request Handler ---------------------------------------------------------
|
||||
|
||||
// Extract usage from SSE stream data
|
||||
function extractUsageFromSSE(text, provider) {
|
||||
let inputTokens = 0;
|
||||
let outputTokens = 0;
|
||||
|
||||
// Split by "data:" lines (SSE format)
|
||||
const lines = text.split("\n");
|
||||
for (const line of lines) {
|
||||
const trimmed = line.trim();
|
||||
if (trimmed.startsWith("event: message_start")) {
|
||||
// Find the next data line
|
||||
const dataIndex = lines.indexOf(line) + 1;
|
||||
if (dataIndex < lines.length && lines[dataIndex].trim().startsWith("data:")) {
|
||||
try {
|
||||
const dataStr = lines[dataIndex].trim().slice(5).trim();
|
||||
const data = JSON.parse(dataStr);
|
||||
if (data.message?.usage?.input_tokens) {
|
||||
inputTokens = data.message.usage.input_tokens;
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
} else if (trimmed.startsWith("event: message_delta")) {
|
||||
const dataIndex = lines.indexOf(line) + 1;
|
||||
if (dataIndex < lines.length && lines[dataIndex].trim().startsWith("data:")) {
|
||||
try {
|
||||
const dataStr = lines[dataIndex].trim().slice(5).trim();
|
||||
const data = JSON.parse(dataStr);
|
||||
if (data.usage?.output_tokens) {
|
||||
outputTokens = data.usage.output_tokens;
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return { inputTokens, outputTokens };
|
||||
}
|
||||
|
||||
async function handleMessages(req, res) {
|
||||
// Read body
|
||||
const chunks = [];
|
||||
@@ -135,6 +423,7 @@ async function handleMessages(req, res) {
|
||||
}
|
||||
|
||||
// Rewrite model name if needed
|
||||
const actualModel = route.rewriteModel || body.model;
|
||||
if (route.rewriteModel) {
|
||||
body.model = route.rewriteModel;
|
||||
}
|
||||
@@ -158,7 +447,7 @@ async function handleMessages(req, res) {
|
||||
const statusCode = upstreamRes.status;
|
||||
|
||||
if (isStreaming && statusCode === 200) {
|
||||
// SSE streaming: pipe through directly
|
||||
// SSE streaming: pipe through directly while tracking usage
|
||||
res.writeHead(statusCode, {
|
||||
"content-type": "text/event-stream",
|
||||
"cache-control": "no-cache",
|
||||
@@ -168,14 +457,22 @@ async function handleMessages(req, res) {
|
||||
|
||||
const reader = upstreamRes.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let fullText = "";
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
const text = decoder.decode(value, { stream: true });
|
||||
fullText += text;
|
||||
res.write(text);
|
||||
}
|
||||
|
||||
// Extract usage from the accumulated SSE text
|
||||
const { inputTokens, outputTokens } = extractUsageFromSSE(fullText, route.provider);
|
||||
if (inputTokens > 0 || outputTokens > 0) {
|
||||
recordUsage(requestedModel, route.provider, actualModel, inputTokens, outputTokens);
|
||||
}
|
||||
} catch (streamErr) {
|
||||
console.error(`[proxy] Stream error from ${route.provider}:`, streamErr.message);
|
||||
} finally {
|
||||
@@ -185,6 +482,23 @@ async function handleMessages(req, res) {
|
||||
// Non-streaming or error: forward the full response
|
||||
const responseBody = await upstreamRes.text();
|
||||
const responseHeaders = { "content-type": upstreamRes.headers.get("content-type") || "application/json" };
|
||||
|
||||
// Try to extract usage from non-streaming response
|
||||
if (statusCode === 200 && responseHeaders["content-type"]?.includes("application/json")) {
|
||||
try {
|
||||
const json = JSON.parse(responseBody);
|
||||
if (json.usage?.input_tokens || json.usage?.output_tokens) {
|
||||
recordUsage(
|
||||
requestedModel,
|
||||
route.provider,
|
||||
actualModel,
|
||||
json.usage.input_tokens,
|
||||
json.usage.output_tokens
|
||||
);
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
|
||||
res.writeHead(statusCode, responseHeaders);
|
||||
res.end(responseBody);
|
||||
}
|
||||
@@ -234,6 +548,50 @@ const server = createServer(async (req, res) => {
|
||||
return;
|
||||
}
|
||||
|
||||
// Usage endpoint (no auth required - no sensitive data)
|
||||
if (req.method === "GET" && urlPath === "/usage") {
|
||||
const { usageData, providers, totalRequests, monthlyRequests, totalCost, monthlyCost, currentPeriod, notes } =
|
||||
buildUsageResponse();
|
||||
|
||||
// Fetch live balances in parallel
|
||||
const balances = await fetchProviderBalances();
|
||||
|
||||
// Apply balances to provider data
|
||||
for (const [providerName, balance] of Object.entries(balances)) {
|
||||
if (providers[providerName]) {
|
||||
providers[providerName].balance = balance;
|
||||
}
|
||||
}
|
||||
|
||||
// Add balance notes
|
||||
if (!balances.anthropic) {
|
||||
providers.anthropic.balance = null;
|
||||
providers.anthropic.pricingNote = "No balance API available — check console.anthropic.com";
|
||||
}
|
||||
if (!balances.zai) {
|
||||
providers.zai.balance = null;
|
||||
providers.zai.pricingNote = "Flat-rate subscription, no balance API available";
|
||||
}
|
||||
|
||||
res.writeHead(200, { "content-type": "application/json" });
|
||||
res.end(
|
||||
JSON.stringify({
|
||||
service: "ein-ai-proxy",
|
||||
generatedAt: new Date().toISOString(),
|
||||
currentPeriod,
|
||||
providers,
|
||||
totals: {
|
||||
totalRequests,
|
||||
monthlyRequests,
|
||||
totalCostUsd: Number(totalCost.toFixed(4)),
|
||||
monthlyCostUsd: Number(monthlyCost.toFixed(4)),
|
||||
note: notes.join(" ") || null,
|
||||
},
|
||||
})
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
// Auth check (everything below requires auth)
|
||||
if (AUTH_TOKEN) {
|
||||
const authHeader = req.headers["authorization"] || "";
|
||||
|
||||
Reference in New Issue
Block a user