import { DateTime } from "https://cdn.skypack.dev/luxon@2.3.2";
import { InStatement, sqlite } from "https://esm.town/v/std/sqlite?v=6";
const createStatsQuery = (hasUserID: boolean) =>
`WITH params AS (
SELECT
? AS interval_minutes,
? AS start_time,
? AS end_time${hasUserID ? ", ? AS user_id" : ""}
)
SELECT
strftime('%Y-%m-%d %H:%M:00', datetime(
(strftime('%s', timestamp) / (interval_minutes * 60)) * (interval_minutes * 60), 'unixepoch'
)) AS bucket,
SUM(tokens) AS tokens,
model,
COUNT(*) AS count
FROM
openai_usage,
params
WHERE
timestamp BETWEEN start_time AND end_time
${hasUserID ? "AND user_id = ?" : ""}
GROUP BY bucket, model
ORDER BY bucket`;
const STATS_QUERY = createStatsQuery(false);
const USER_SCOPED_STATS_QUERY = createStatsQuery(true);
const HOUR = 60;
type UsageRow = {
userId: string;
handle: string;
tier: string;
timestamp?: Date;
tokens: number;
model: string;
};
export class OpenAIUsage {
constructor() {}
async migrate() {
await sqlite.batch([`CREATE TABLE IF NOT EXISTS openai_usage (
id INTEGER PRIMARY KEY,
user_id TEXT NOT NULL,
handle TEXT NOT NULL,
tier TEXT NOT NULL,
timestamp TEXT NOT NULL DEFAULT (datetime('now')),
tokens INTEGER NOT NULL,
model TEXT NOT NULL
)`]);
}
async drop() {
await sqlite.batch([`DROP TABLE IF EXISTS openai_usage`]);
}
async writeUsage(ur: UsageRow) {
sqlite.execute({
sql: "INSERT INTO openai_usage (user_id, handle, tier, tokens, model) VALUES (?, ?, ?, ?, ?)",
args: [ur.userId, ur.handle, ur.tier, ur.tokens, ur.model],
});
}
async query(statement: InStatement): Promise<UsageRow[]> {
let rows = await sqlite.execute(statement);
return rows.rows.map(r => ({
userId: r[1],
handle: r[2],
tier: r[3],
timestamp: new Date(r[4]),
tokens: r[5],
model: r[6],
}));
}
async recentGpt4Usage(userId: string) {
const now = new Date();
now.setHours(now.getHours() - 24);
let resp = await sqlite.execute({
sql: `select count(*) from openai_usage where model LIKE 'gpt-4%' and model NOT LIKE '%mini%' and user_id = ? and timestamp > ?`,
args: [userId, now.toISOString()],
});
return resp.rows[0][0];
}
async usageStats(
{ user_id, start, end, intervalMinutes, timeZone }: {
user_id?: string;
start: string;
end: string;
intervalMinutes: number;
timeZone: string;
},
): Promise<{ bucket: string; count: number }[]> {
const startDate = DateTime.fromISO(start, { zone: "utc" }).setZone(timeZone);
const endDate = DateTime.fromISO(end, { zone: "utc" }).setZone(timeZone);
const startString = startDate.toFormat("yyyy-MM-dd HH:mm:ss");
const endString = endDate.toFormat("yyyy-MM-dd HH:mm:ss");
const args = [intervalMinutes, startString, endString];
user_id && args.push(user_id);
let resp = await sqlite.execute({
sql: user_id ? USER_SCOPED_STATS_QUERY : STATS_QUERY,