opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add columns for cached tokens from aibridge (#23832)

Browse Source 
Two new columns added to aibridge_token_usages:
  - cache_read_input_tokens (BIGINT, default 0)
  - cache_write_input_tokens (BIGINT, default 0)

Migration backfills existing rows by extracting values from the metadata
JSONB column (cache_read_input, input_cached, prompt_cached for reads
(max value selected since only 1 should be set), cache_creation_input
for writes).

All references to data from metadata were updated to reference new
columns. No other changes then changing where data is extracted from.

Requires aibridge library version bump to include:
https://github.com/coder/aibridge/pull/229
Fixes: https://github.com/coder/aibridge/issues/150
This commit is contained in:
Paweł Banaszewski
2026-04-03 16:27:31 +02:00
committed by GitHub
parent da3c46b557
commit 8369fa88fd
28 changed files with 423 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/queries.sql.go
+48 -43
View File
@@ -148,21 +148,8 @@ token_aggregates AS (
SELECT
COALESCE(SUM(tu.input_tokens), 0) AS token_count_input,
COALESCE(SUM(tu.output_tokens), 0) AS token_count_output,
-- Cached tokens are stored in metadata JSON, extract if available.
-- Read tokens may be stored in:
-- - cache_read_input (Anthropic)
-- - prompt_cached (OpenAI)
COALESCE(SUM(
COALESCE((tu.metadata->>'cache_read_input')::bigint, 0) +
COALESCE((tu.metadata->>'prompt_cached')::bigint, 0)
), 0) AS token_count_cached_read,
-- Written tokens may be stored in:
-- - cache_creation_input (Anthropic)
-- Note that cache_ephemeral_5m_input and cache_ephemeral_1h_input on
-- Anthropic are included in the cache_creation_input field.
COALESCE(SUM(
COALESCE((tu.metadata->>'cache_creation_input')::bigint, 0)
), 0) AS token_count_cached_written,
COALESCE(SUM(tu.cache_read_input_tokens), 0) AS token_count_cached_read,
COALESCE(SUM(tu.cache_write_input_tokens), 0) AS token_count_cached_written,
COUNT(tu.id) AS token_usages_count
FROM
interceptions_in_range i
@@ -559,7 +546,7 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn
const getAIBridgeTokenUsagesByInterceptionID = `-- name: GetAIBridgeTokenUsagesByInterceptionID :many
SELECT
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens
FROM
aibridge_token_usages WHERE interception_id = $1::uuid
ORDER BY
@@ -584,6 +571,8 @@ func (q *sqlQuerier) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context,
&i.OutputTokens,
&i.Metadata,
&i.CreatedAt,
&i.CacheReadInputTokens,
&i.CacheWriteInputTokens,
); err != nil {
return nil, err
}
@@ -781,21 +770,23 @@ func (q *sqlQuerier) InsertAIBridgeModelThought(ctx context.Context, arg InsertA
const insertAIBridgeTokenUsage = `-- name: InsertAIBridgeTokenUsage :one
INSERT INTO aibridge_token_usages (
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
id, interception_id, provider_response_id, input_tokens, output_tokens, cache_read_input_tokens, cache_write_input_tokens, metadata, created_at
) VALUES (
$1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb), $7
$1, $2, $3, $4, $5, $6, $7, COALESCE($8::jsonb, '{}'::jsonb), $9
)
RETURNING id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
RETURNING id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens
`
type InsertAIBridgeTokenUsageParams struct {
ID uuid.UUID `db:"id" json:"id"`
InterceptionID uuid.UUID `db:"interception_id" json:"interception_id"`
ProviderResponseID string `db:"provider_response_id" json:"provider_response_id"`
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
ID uuid.UUID `db:"id" json:"id"`
InterceptionID uuid.UUID `db:"interception_id" json:"interception_id"`
ProviderResponseID string `db:"provider_response_id" json:"provider_response_id"`
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
CacheReadInputTokens int64 `db:"cache_read_input_tokens" json:"cache_read_input_tokens"`
CacheWriteInputTokens int64 `db:"cache_write_input_tokens" json:"cache_write_input_tokens"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
}
func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) (AIBridgeTokenUsage, error) {
@@ -805,6 +796,8 @@ func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIB
arg.ProviderResponseID,
arg.InputTokens,
arg.OutputTokens,
arg.CacheReadInputTokens,
arg.CacheWriteInputTokens,
arg.Metadata,
arg.CreatedAt,
)
@@ -817,6 +810,8 @@ func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIB
&i.OutputTokens,
&i.Metadata,
&i.CreatedAt,
&i.CacheReadInputTokens,
&i.CacheWriteInputTokens,
)
return i, err
}
@@ -1448,6 +1443,8 @@ SELECT
sp.threads,
COALESCE(st.input_tokens, 0)::bigint AS input_tokens,
COALESCE(st.output_tokens, 0)::bigint AS output_tokens,
COALESCE(st.cache_read_input_tokens, 0)::bigint AS cache_read_input_tokens,
COALESCE(st.cache_write_input_tokens, 0)::bigint AS cache_write_input_tokens,
COALESCE(slp.prompt, '') AS last_prompt
FROM
session_page sp
@@ -1469,7 +1466,9 @@ LEFT JOIN LATERAL (
-- Aggregate tokens only for this session's interceptions.
SELECT
COALESCE(SUM(tu.input_tokens), 0)::bigint AS input_tokens,
COALESCE(SUM(tu.output_tokens), 0)::bigint AS output_tokens
COALESCE(SUM(tu.output_tokens), 0)::bigint AS output_tokens,
COALESCE(SUM(tu.cache_read_input_tokens), 0)::bigint AS cache_read_input_tokens,
COALESCE(SUM(tu.cache_write_input_tokens), 0)::bigint AS cache_write_input_tokens
FROM aibridge_token_usages tu
WHERE tu.interception_id = ANY(sr.interception_ids)
) st ON true
@@ -1501,21 +1500,23 @@ type ListAIBridgeSessionsParams struct {
}
type ListAIBridgeSessionsRow struct {
SessionID string `db:"session_id" json:"session_id"`
UserID uuid.UUID `db:"user_id" json:"user_id"`
UserUsername string `db:"user_username" json:"user_username"`
UserName string `db:"user_name" json:"user_name"`
UserAvatarUrl string `db:"user_avatar_url" json:"user_avatar_url"`
Providers []string `db:"providers" json:"providers"`
Models []string `db:"models" json:"models"`
Client string `db:"client" json:"client"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
StartedAt time.Time `db:"started_at" json:"started_at"`
EndedAt time.Time `db:"ended_at" json:"ended_at"`
Threads int64 `db:"threads" json:"threads"`
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
LastPrompt string `db:"last_prompt" json:"last_prompt"`
SessionID string `db:"session_id" json:"session_id"`
UserID uuid.UUID `db:"user_id" json:"user_id"`
UserUsername string `db:"user_username" json:"user_username"`
UserName string `db:"user_name" json:"user_name"`
UserAvatarUrl string `db:"user_avatar_url" json:"user_avatar_url"`
Providers []string `db:"providers" json:"providers"`
Models []string `db:"models" json:"models"`
Client string `db:"client" json:"client"`
Metadata json.RawMessage `db:"metadata" json:"metadata"`
StartedAt time.Time `db:"started_at" json:"started_at"`
EndedAt time.Time `db:"ended_at" json:"ended_at"`
Threads int64 `db:"threads" json:"threads"`
InputTokens int64 `db:"input_tokens" json:"input_tokens"`
OutputTokens int64 `db:"output_tokens" json:"output_tokens"`
CacheReadInputTokens int64 `db:"cache_read_input_tokens" json:"cache_read_input_tokens"`
CacheWriteInputTokens int64 `db:"cache_write_input_tokens" json:"cache_write_input_tokens"`
LastPrompt string `db:"last_prompt" json:"last_prompt"`
}
// Returns paginated sessions with aggregated metadata, token counts, and
@@ -1560,6 +1561,8 @@ func (q *sqlQuerier) ListAIBridgeSessions(ctx context.Context, arg ListAIBridgeS
&i.Threads,
&i.InputTokens,
&i.OutputTokens,
&i.CacheReadInputTokens,
&i.CacheWriteInputTokens,
&i.LastPrompt,
); err != nil {
return nil, err
@@ -1577,7 +1580,7 @@ func (q *sqlQuerier) ListAIBridgeSessions(ctx context.Context, arg ListAIBridgeS
const listAIBridgeTokenUsagesByInterceptionIDs = `-- name: ListAIBridgeTokenUsagesByInterceptionIDs :many
SELECT
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at
id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens
FROM
aibridge_token_usages
WHERE
@@ -1604,6 +1607,8 @@ func (q *sqlQuerier) ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Contex
&i.OutputTokens,
&i.Metadata,
&i.CreatedAt,
&i.CacheReadInputTokens,
&i.CacheWriteInputTokens,
); err != nil {
return nil, err
}