feat: token usage state (#153)

* feat: token usage state

* fix: typo
This commit is contained in:
Acbox Liu
2026-03-01 02:19:07 +08:00
committed by GitHub
parent 443ede30b4
commit 0cdf822603
33 changed files with 1916 additions and 35 deletions
+2 -1
View File
@@ -9,9 +9,10 @@ SET status = $2,
result_text = $3,
error_message = $4,
usage = $5,
model_id = $6,
completed_at = now()
WHERE id = $1
RETURNING id, bot_id, status, result_text, error_message, usage, started_at, completed_at;
RETURNING id, bot_id, status, result_text, error_message, usage, model_id, started_at, completed_at;
-- name: ListHeartbeatLogsByBot :many
SELECT id, bot_id, status, result_text, error_message, usage, started_at, completed_at
+4 -2
View File
@@ -10,7 +10,8 @@ INSERT INTO bot_history_messages (
role,
content,
metadata,
usage
usage,
model_id
)
VALUES (
sqlc.arg(bot_id),
@@ -23,7 +24,8 @@ VALUES (
sqlc.arg(role),
sqlc.arg(content),
sqlc.arg(metadata),
sqlc.arg(usage)
sqlc.arg(usage),
sqlc.narg(model_id)::uuid
)
RETURNING
id,
+69
View File
@@ -0,0 +1,69 @@
-- name: GetMessageTokenUsageByDay :many
SELECT
date_trunc('day', m.created_at)::date AS day,
COALESCE(SUM((m.usage->>'inputTokens')::bigint), 0)::bigint AS input_tokens,
COALESCE(SUM((m.usage->>'outputTokens')::bigint), 0)::bigint AS output_tokens,
COALESCE(SUM((m.usage->'inputTokenDetails'->>'cacheReadTokens')::bigint), 0)::bigint AS cache_read_tokens,
COALESCE(SUM((m.usage->'inputTokenDetails'->>'cacheWriteTokens')::bigint), 0)::bigint AS cache_write_tokens,
COALESCE(SUM((m.usage->'outputTokenDetails'->>'reasoningTokens')::bigint), 0)::bigint AS reasoning_tokens
FROM bot_history_messages m
WHERE m.bot_id = sqlc.arg(bot_id)
AND m.usage IS NOT NULL
AND m.created_at >= sqlc.arg(from_time)
AND m.created_at < sqlc.arg(to_time)
AND (sqlc.narg(model_id)::uuid IS NULL OR m.model_id = sqlc.narg(model_id)::uuid)
GROUP BY day
ORDER BY day;
-- name: GetHeartbeatTokenUsageByDay :many
SELECT
date_trunc('day', h.started_at)::date AS day,
COALESCE(SUM((h.usage->>'inputTokens')::bigint), 0)::bigint AS input_tokens,
COALESCE(SUM((h.usage->>'outputTokens')::bigint), 0)::bigint AS output_tokens,
COALESCE(SUM((h.usage->'inputTokenDetails'->>'cacheReadTokens')::bigint), 0)::bigint AS cache_read_tokens,
COALESCE(SUM((h.usage->'inputTokenDetails'->>'cacheWriteTokens')::bigint), 0)::bigint AS cache_write_tokens,
COALESCE(SUM((h.usage->'outputTokenDetails'->>'reasoningTokens')::bigint), 0)::bigint AS reasoning_tokens
FROM bot_heartbeat_logs h
WHERE h.bot_id = sqlc.arg(bot_id)
AND h.usage IS NOT NULL
AND h.started_at >= sqlc.arg(from_time)
AND h.started_at < sqlc.arg(to_time)
AND (sqlc.narg(model_id)::uuid IS NULL OR h.model_id = sqlc.narg(model_id)::uuid)
GROUP BY day
ORDER BY day;
-- name: GetMessageTokenUsageByModel :many
SELECT
m.model_id,
COALESCE(mo.model_id, 'unknown') AS model_slug,
COALESCE(mo.name, 'Unknown') AS model_name,
COALESCE(lp.name, 'Unknown') AS provider_name,
COALESCE(SUM((m.usage->>'inputTokens')::bigint), 0)::bigint AS input_tokens,
COALESCE(SUM((m.usage->>'outputTokens')::bigint), 0)::bigint AS output_tokens
FROM bot_history_messages m
LEFT JOIN models mo ON mo.id = m.model_id
LEFT JOIN llm_providers lp ON lp.id = mo.llm_provider_id
WHERE m.bot_id = sqlc.arg(bot_id)
AND m.usage IS NOT NULL
AND m.created_at >= sqlc.arg(from_time)
AND m.created_at < sqlc.arg(to_time)
GROUP BY m.model_id, mo.model_id, mo.name, lp.name
ORDER BY input_tokens DESC;
-- name: GetHeartbeatTokenUsageByModel :many
SELECT
h.model_id,
COALESCE(mo.model_id, 'unknown') AS model_slug,
COALESCE(mo.name, 'Unknown') AS model_name,
COALESCE(lp.name, 'Unknown') AS provider_name,
COALESCE(SUM((h.usage->>'inputTokens')::bigint), 0)::bigint AS input_tokens,
COALESCE(SUM((h.usage->>'outputTokens')::bigint), 0)::bigint AS output_tokens
FROM bot_heartbeat_logs h
LEFT JOIN models mo ON mo.id = h.model_id
LEFT JOIN llm_providers lp ON lp.id = mo.llm_provider_id
WHERE h.bot_id = sqlc.arg(bot_id)
AND h.usage IS NOT NULL
AND h.started_at >= sqlc.arg(from_time)
AND h.started_at < sqlc.arg(to_time)
GROUP BY h.model_id, mo.model_id, mo.name, lp.name
ORDER BY input_tokens DESC;