mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: add columns for cached tokens from aibridge (#23832)
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:
@@ -1037,8 +1037,10 @@ func AIBridgeSession(row database.ListAIBridgeSessionsRow) codersdk.AIBridgeSess
|
||||
StartedAt: row.StartedAt,
|
||||
Threads: row.Threads,
|
||||
TokenUsageSummary: codersdk.AIBridgeSessionTokenUsageSummary{
|
||||
InputTokens: row.InputTokens,
|
||||
OutputTokens: row.OutputTokens,
|
||||
InputTokens: row.InputTokens,
|
||||
OutputTokens: row.OutputTokens,
|
||||
CacheReadInputTokens: row.CacheReadInputTokens,
|
||||
CacheWriteInputTokens: row.CacheWriteInputTokens,
|
||||
},
|
||||
}
|
||||
// Ensure non-nil slices for JSON serialization.
|
||||
@@ -1062,13 +1064,15 @@ func AIBridgeSession(row database.ListAIBridgeSessionsRow) codersdk.AIBridgeSess
|
||||
|
||||
func AIBridgeTokenUsage(usage database.AIBridgeTokenUsage) codersdk.AIBridgeTokenUsage {
|
||||
return codersdk.AIBridgeTokenUsage{
|
||||
ID: usage.ID,
|
||||
InterceptionID: usage.InterceptionID,
|
||||
ProviderResponseID: usage.ProviderResponseID,
|
||||
InputTokens: usage.InputTokens,
|
||||
OutputTokens: usage.OutputTokens,
|
||||
Metadata: jsonOrEmptyMap(usage.Metadata),
|
||||
CreatedAt: usage.CreatedAt,
|
||||
ID: usage.ID,
|
||||
InterceptionID: usage.InterceptionID,
|
||||
ProviderResponseID: usage.ProviderResponseID,
|
||||
InputTokens: usage.InputTokens,
|
||||
OutputTokens: usage.OutputTokens,
|
||||
CacheReadInputTokens: usage.CacheReadInputTokens,
|
||||
CacheWriteInputTokens: usage.CacheWriteInputTokens,
|
||||
Metadata: jsonOrEmptyMap(usage.Metadata),
|
||||
CreatedAt: usage.CreatedAt,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1179,9 +1183,11 @@ func AIBridgeSessionThreads(
|
||||
PageStartedAt: pageStartedAt,
|
||||
PageEndedAt: pageEndedAt,
|
||||
TokenUsageSummary: codersdk.AIBridgeSessionThreadsTokenUsage{
|
||||
InputTokens: session.InputTokens,
|
||||
OutputTokens: session.OutputTokens,
|
||||
Metadata: sessionTokenMeta,
|
||||
InputTokens: session.InputTokens,
|
||||
OutputTokens: session.OutputTokens,
|
||||
CacheReadInputTokens: session.CacheReadInputTokens,
|
||||
CacheWriteInputTokens: session.CacheWriteInputTokens,
|
||||
Metadata: sessionTokenMeta,
|
||||
},
|
||||
Threads: threads,
|
||||
}
|
||||
@@ -1314,17 +1320,19 @@ func buildAIBridgeThread(
|
||||
|
||||
// aggregateTokenUsage sums token usage rows and aggregates metadata.
|
||||
func aggregateTokenUsage(tokens []database.AIBridgeTokenUsage) codersdk.AIBridgeSessionThreadsTokenUsage {
|
||||
var inputTokens, outputTokens int64
|
||||
var inputTokens, outputTokens, cacheRead, cacheWrite int64
|
||||
for _, tu := range tokens {
|
||||
inputTokens += tu.InputTokens
|
||||
outputTokens += tu.OutputTokens
|
||||
// TODO: once https://github.com/coder/aibridge/issues/150 lands we
|
||||
// should aggregate the other token types.
|
||||
cacheRead += tu.CacheReadInputTokens
|
||||
cacheWrite += tu.CacheWriteInputTokens
|
||||
}
|
||||
return codersdk.AIBridgeSessionThreadsTokenUsage{
|
||||
InputTokens: inputTokens,
|
||||
OutputTokens: outputTokens,
|
||||
Metadata: aggregateTokenMetadata(tokens),
|
||||
InputTokens: inputTokens,
|
||||
OutputTokens: outputTokens,
|
||||
CacheReadInputTokens: cacheRead,
|
||||
CacheWriteInputTokens: cacheWrite,
|
||||
Metadata: aggregateTokenMetadata(tokens),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -259,11 +259,13 @@ func TestAIBridgeInterception(t *testing.T) {
|
||||
},
|
||||
tokenUsages: []database.AIBridgeTokenUsage{
|
||||
{
|
||||
ID: uuid.New(),
|
||||
InterceptionID: interceptionID,
|
||||
ProviderResponseID: "resp-123",
|
||||
InputTokens: 100,
|
||||
OutputTokens: 200,
|
||||
ID: uuid.New(),
|
||||
InterceptionID: interceptionID,
|
||||
ProviderResponseID: "resp-123",
|
||||
InputTokens: 100,
|
||||
OutputTokens: 200,
|
||||
CacheReadInputTokens: 50,
|
||||
CacheWriteInputTokens: 10,
|
||||
Metadata: pqtype.NullRawMessage{
|
||||
RawMessage: json.RawMessage(`{"cache":"hit"}`),
|
||||
Valid: true,
|
||||
@@ -413,6 +415,8 @@ func TestAIBridgeInterception(t *testing.T) {
|
||||
require.Equal(t, tu.ProviderResponseID, result.TokenUsages[i].ProviderResponseID)
|
||||
require.Equal(t, tu.InputTokens, result.TokenUsages[i].InputTokens)
|
||||
require.Equal(t, tu.OutputTokens, result.TokenUsages[i].OutputTokens)
|
||||
require.Equal(t, tu.CacheReadInputTokens, result.TokenUsages[i].CacheReadInputTokens)
|
||||
require.Equal(t, tu.CacheWriteInputTokens, result.TokenUsages[i].CacheWriteInputTokens)
|
||||
}
|
||||
|
||||
// Verify user prompts are converted correctly.
|
||||
|
||||
@@ -1613,13 +1613,15 @@ func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertA
|
||||
|
||||
func AIBridgeTokenUsage(t testing.TB, db database.Store, seed database.InsertAIBridgeTokenUsageParams) database.AIBridgeTokenUsage {
|
||||
usage, err := db.InsertAIBridgeTokenUsage(genCtx, database.InsertAIBridgeTokenUsageParams{
|
||||
ID: takeFirst(seed.ID, uuid.New()),
|
||||
InterceptionID: takeFirst(seed.InterceptionID, uuid.New()),
|
||||
ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"),
|
||||
InputTokens: takeFirst(seed.InputTokens, 100),
|
||||
OutputTokens: takeFirst(seed.OutputTokens, 100),
|
||||
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
|
||||
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
|
||||
ID: takeFirst(seed.ID, uuid.New()),
|
||||
InterceptionID: takeFirst(seed.InterceptionID, uuid.New()),
|
||||
ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"),
|
||||
InputTokens: takeFirst(seed.InputTokens, 100),
|
||||
OutputTokens: takeFirst(seed.OutputTokens, 100),
|
||||
CacheReadInputTokens: seed.CacheReadInputTokens,
|
||||
CacheWriteInputTokens: seed.CacheWriteInputTokens,
|
||||
Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")),
|
||||
CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()),
|
||||
})
|
||||
require.NoError(t, err, "insert aibridge token usage")
|
||||
return usage
|
||||
|
||||
Generated
+3
-1
@@ -1134,7 +1134,9 @@ CREATE TABLE aibridge_token_usages (
|
||||
input_tokens bigint NOT NULL,
|
||||
output_tokens bigint NOT NULL,
|
||||
metadata jsonb,
|
||||
created_at timestamp with time zone NOT NULL
|
||||
created_at timestamp with time zone NOT NULL,
|
||||
cache_read_input_tokens bigint DEFAULT 0 NOT NULL,
|
||||
cache_write_input_tokens bigint DEFAULT 0 NOT NULL
|
||||
);
|
||||
|
||||
COMMENT ON TABLE aibridge_token_usages IS 'Audit log of tokens used by intercepted requests in AI Bridge';
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
ALTER TABLE aibridge_token_usages
|
||||
DROP COLUMN cache_read_input_tokens,
|
||||
DROP COLUMN cache_write_input_tokens;
|
||||
@@ -0,0 +1,26 @@
|
||||
ALTER TABLE aibridge_token_usages
|
||||
ADD COLUMN cache_read_input_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
ADD COLUMN cache_write_input_tokens BIGINT NOT NULL DEFAULT 0;
|
||||
|
||||
-- Backfill from metadata JSONB. Old rows stored cache tokens under
|
||||
-- provider-specific keys; new rows use the dedicated columns above.
|
||||
UPDATE aibridge_token_usages
|
||||
SET
|
||||
|
||||
-- Cache-read metadata keys by provider:
|
||||
-- Anthropic (/v1/messages): "cache_read_input"
|
||||
-- OpenAI (/v1/responses): "input_cached"
|
||||
-- OpenAI (/v1/chat/completions): "prompt_cached"
|
||||
cache_read_input_tokens = GREATEST(
|
||||
COALESCE((metadata->>'cache_read_input')::bigint, 0),
|
||||
COALESCE((metadata->>'input_cached')::bigint, 0),
|
||||
COALESCE((metadata->>'prompt_cached')::bigint, 0)
|
||||
),
|
||||
|
||||
-- Cache-write metadata keys by provider:
|
||||
-- Anthropic (/v1/messages): "cache_creation_input"
|
||||
-- OpenAI does not report cache-write tokens.
|
||||
cache_write_input_tokens = COALESCE((metadata->>'cache_creation_input')::bigint, 0)
|
||||
WHERE metadata IS NOT NULL
|
||||
AND cache_read_input_tokens = 0
|
||||
AND cache_write_input_tokens = 0;
|
||||
@@ -1029,6 +1029,8 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeSessions(ctx context.Context, arg Lis
|
||||
&i.Threads,
|
||||
&i.InputTokens,
|
||||
&i.OutputTokens,
|
||||
&i.CacheReadInputTokens,
|
||||
&i.CacheWriteInputTokens,
|
||||
&i.LastPrompt,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -4055,11 +4055,13 @@ type AIBridgeTokenUsage struct {
|
||||
ID uuid.UUID `db:"id" json:"id"`
|
||||
InterceptionID uuid.UUID `db:"interception_id" json:"interception_id"`
|
||||
// The ID for the response in which the tokens were used, produced by the provider.
|
||||
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 pqtype.NullRawMessage `db:"metadata" json:"metadata"`
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
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 pqtype.NullRawMessage `db:"metadata" json:"metadata"`
|
||||
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
||||
CacheReadInputTokens int64 `db:"cache_read_input_tokens" json:"cache_read_input_tokens"`
|
||||
CacheWriteInputTokens int64 `db:"cache_write_input_tokens" json:"cache_write_input_tokens"`
|
||||
}
|
||||
|
||||
// Audit log of tool calls in intercepted requests in AI Bridge
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -31,9 +31,9 @@ WHERE aibridge_interceptions.id = (
|
||||
|
||||
-- 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 (
|
||||
@id, @interception_id, @provider_response_id, @input_tokens, @output_tokens, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at
|
||||
@id, @interception_id, @provider_response_id, @input_tokens, @output_tokens, @cache_read_input_tokens, @cache_write_input_tokens, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at
|
||||
)
|
||||
RETURNING *;
|
||||
|
||||
@@ -299,21 +299,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
|
||||
@@ -552,6 +539,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
|
||||
@@ -573,7 +562,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
|
||||
|
||||
Reference in New Issue
Block a user