diff --git a/CHANGELOG.md b/CHANGELOG.md index 15b7d1d..353c519 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,10 @@ All notable changes to light-tokenmeter are listed here. The format follows ## [Unreleased] ### Added +- `cache_read_tokens` and `cache_write_tokens` on usage events (`usage.v1` fields 17 and 18) and + summed in summary and by-meta rows: the input tokens read from and written to the vendor's + prompt cache, which vendors bill apart from plain input. They are part of `tokens_input`. + Stored in Postgres and ClickHouse; events from older callers carry 0. - Authentication: bearer tokens with `write` and `read` scopes, configured as SHA-256 hashes in `TOKENMETER_TOKENS` or `TOKENMETER_TOKENS_FILE`, several at once for rotation. Refusals are `401`/`403` over HTTP and `UNAUTHENTICATED`/`PERMISSION_DENIED` over gRPC. Without tokens the diff --git a/README.md b/README.md index 9e348c8..a31b03d 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,7 @@ An event describes one model call: | `vendor_http_status` | int | no | HTTP status returned by the vendor. | | `meta` | map of string | no | Your own tags, e.g. `{"job_id": "…"}`: non-empty keys, at most 16 KiB of keys and values per event. | | `cost_nanousd` | int ≥0 | no | Call cost in nano-USD (1e-9 USD), priced by the caller; stored as-is and never recomputed. | +| `cache_read_tokens`, `cache_write_tokens` | int ≥0 | no | Input tokens read from and written to the vendor's prompt cache. They are part of `tokens_input`, not added to it, so together they must not exceed it; an event that breaks this is refused. Anthropic reports them outside its `input_tokens`: send the sum of all three as `tokens_input`. | Batches (`POST /v1/usage/events:batch`, or `RecordUsageEvents` over gRPC) take up to 500 events and report how many were created and how many were duplicates: @@ -183,7 +184,8 @@ curl -s "localhost:8080/v1/usage/summary?$Q&model=gpt-4o&group_by_user_id=true" ``` Each row carries `tokens_input`, `tokens_output`, `tokens_total`, `request_count`, -`sum_duration_ms` and `cost_nanousd` (summed spend; divide by 1e9 for USD), so you can +`sum_duration_ms`, `cost_nanousd` (summed spend; divide by 1e9 for USD) and +`cache_read_tokens` / `cache_write_tokens` (the cached part of `tokens_input`), so you can compute tokens per request, cost per request, average latency and error share yourself. The signals that usually separate greedy from worthy: - high `tokens_total / request_count` on a trivial `action_id` (bloated prompts), diff --git a/gen/usage/v1/usage.pb.go b/gen/usage/v1/usage.pb.go index cd6d296..0c63154 100644 --- a/gen/usage/v1/usage.pb.go +++ b/gen/usage/v1/usage.pb.go @@ -113,9 +113,15 @@ type UsageEvent struct { Meta map[string]string `protobuf:"bytes,15,rep,name=meta,proto3" json:"meta,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` // Cost of the call in nano-USD (1e-9 USD), computed by the caller from its // price catalog at settle time; stored as an immutable fact (optional, >= 0) - CostNanousd int64 `protobuf:"varint,16,opt,name=cost_nanousd,json=costNanousd,proto3" json:"cost_nanousd,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + CostNanousd int64 `protobuf:"varint,16,opt,name=cost_nanousd,json=costNanousd,proto3" json:"cost_nanousd,omitempty"` + // Input tokens served from the vendor's prompt cache (optional, >= 0). Part of + // tokens_input, not added to it; usually billed at a fraction of the input rate. + CacheReadTokens int64 `protobuf:"varint,17,opt,name=cache_read_tokens,json=cacheReadTokens,proto3" json:"cache_read_tokens,omitempty"` + // Input tokens written to the vendor's prompt cache (optional, >= 0). Part of + // tokens_input; cache_read_tokens + cache_write_tokens must not exceed it. + CacheWriteTokens int64 `protobuf:"varint,18,opt,name=cache_write_tokens,json=cacheWriteTokens,proto3" json:"cache_write_tokens,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *UsageEvent) Reset() { @@ -260,6 +266,20 @@ func (x *UsageEvent) GetCostNanousd() int64 { return 0 } +func (x *UsageEvent) GetCacheReadTokens() int64 { + if x != nil { + return x.CacheReadTokens + } + return 0 +} + +func (x *UsageEvent) GetCacheWriteTokens() int64 { + if x != nil { + return x.CacheWriteTokens + } + return 0 +} + // RecordUsageEventRequest is the request for recording a single usage event type RecordUsageEventRequest struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -624,9 +644,12 @@ type UsageSummaryRow struct { RequestCount int64 `protobuf:"varint,13,opt,name=request_count,json=requestCount,proto3" json:"request_count,omitempty"` SumDurationMs int64 `protobuf:"varint,14,opt,name=sum_duration_ms,json=sumDurationMs,proto3" json:"sum_duration_ms,omitempty"` // Summed cost in nano-USD (1e-9 USD) - CostNanousd int64 `protobuf:"varint,15,opt,name=cost_nanousd,json=costNanousd,proto3" json:"cost_nanousd,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + CostNanousd int64 `protobuf:"varint,15,opt,name=cost_nanousd,json=costNanousd,proto3" json:"cost_nanousd,omitempty"` + // Summed parts of tokens_input read from and written to the prompt cache + CacheReadTokens int64 `protobuf:"varint,16,opt,name=cache_read_tokens,json=cacheReadTokens,proto3" json:"cache_read_tokens,omitempty"` + CacheWriteTokens int64 `protobuf:"varint,17,opt,name=cache_write_tokens,json=cacheWriteTokens,proto3" json:"cache_write_tokens,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *UsageSummaryRow) Reset() { @@ -750,6 +773,20 @@ func (x *UsageSummaryRow) GetCostNanousd() int64 { return 0 } +func (x *UsageSummaryRow) GetCacheReadTokens() int64 { + if x != nil { + return x.CacheReadTokens + } + return 0 +} + +func (x *UsageSummaryRow) GetCacheWriteTokens() int64 { + if x != nil { + return x.CacheWriteTokens + } + return 0 +} + // GetUsageSummaryResponse is the response for getting usage summary type GetUsageSummaryResponse struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -890,9 +927,12 @@ type UsageByMetaRow struct { TokensTotal int64 `protobuf:"varint,6,opt,name=tokens_total,json=tokensTotal,proto3" json:"tokens_total,omitempty"` RequestCount int64 `protobuf:"varint,7,opt,name=request_count,json=requestCount,proto3" json:"request_count,omitempty"` // Summed cost in nano-USD (1e-9 USD) - CostNanousd int64 `protobuf:"varint,8,opt,name=cost_nanousd,json=costNanousd,proto3" json:"cost_nanousd,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + CostNanousd int64 `protobuf:"varint,8,opt,name=cost_nanousd,json=costNanousd,proto3" json:"cost_nanousd,omitempty"` + // Summed parts of tokens_input read from and written to the prompt cache + CacheReadTokens int64 `protobuf:"varint,9,opt,name=cache_read_tokens,json=cacheReadTokens,proto3" json:"cache_read_tokens,omitempty"` + CacheWriteTokens int64 `protobuf:"varint,10,opt,name=cache_write_tokens,json=cacheWriteTokens,proto3" json:"cache_write_tokens,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *UsageByMetaRow) Reset() { @@ -981,6 +1021,20 @@ func (x *UsageByMetaRow) GetCostNanousd() int64 { return 0 } +func (x *UsageByMetaRow) GetCacheReadTokens() int64 { + if x != nil { + return x.CacheReadTokens + } + return 0 +} + +func (x *UsageByMetaRow) GetCacheWriteTokens() int64 { + if x != nil { + return x.CacheWriteTokens + } + return 0 +} + // GetUsageByMetaResponse is the response for meta-scoped aggregation type GetUsageByMetaResponse struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -1038,7 +1092,7 @@ var File_usage_v1_usage_proto protoreflect.FileDescriptor const file_usage_v1_usage_proto_rawDesc = "" + "\n" + - "\x14usage/v1/usage.proto\x12\busage.v1\x1a\x1fgoogle/protobuf/timestamp.proto\"\xfc\x04\n" + + "\x14usage/v1/usage.proto\x12\busage.v1\x1a\x1fgoogle/protobuf/timestamp.proto\"\xd6\x05\n" + "\n" + "UsageEvent\x12\x19\n" + "\bevent_id\x18\x01 \x01(\tR\aeventId\x12;\n" + @@ -1061,7 +1115,9 @@ const file_usage_v1_usage_proto_rawDesc = "" + "errorClass\x12,\n" + "\x12vendor_http_status\x18\x0e \x01(\x05R\x10vendorHttpStatus\x122\n" + "\x04meta\x18\x0f \x03(\v2\x1e.usage.v1.UsageEvent.MetaEntryR\x04meta\x12!\n" + - "\fcost_nanousd\x18\x10 \x01(\x03R\vcostNanousd\x1a7\n" + + "\fcost_nanousd\x18\x10 \x01(\x03R\vcostNanousd\x12*\n" + + "\x11cache_read_tokens\x18\x11 \x01(\x03R\x0fcacheReadTokens\x12,\n" + + "\x12cache_write_tokens\x18\x12 \x01(\x03R\x10cacheWriteTokens\x1a7\n" + "\tMetaEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"E\n" + @@ -1090,7 +1146,7 @@ const file_usage_v1_usage_proto_rawDesc = "" + "\x0fgroup_by_vendor\x18\f \x01(\bR\rgroupByVendor\x12$\n" + "\x0egroup_by_model\x18\r \x01(\bR\fgroupByModel\x12+\n" + "\x12group_by_action_id\x18\x0e \x01(\bR\x0fgroupByActionId\x12&\n" + - "\x0fgroup_by_status\x18\x0f \x01(\bR\rgroupByStatus\"\xcd\x03\n" + + "\x0fgroup_by_status\x18\x0f \x01(\bR\rgroupByStatus\"\xa7\x04\n" + "\x0fUsageSummaryRow\x122\n" + "\x06bucket\x18\x01 \x01(\v2\x1a.google.protobuf.TimestampR\x06bucket\x12\x1d\n" + "\n" + @@ -1106,7 +1162,9 @@ const file_usage_v1_usage_proto_rawDesc = "" + "\ftokens_total\x18\f \x01(\x03R\vtokensTotal\x12#\n" + "\rrequest_count\x18\r \x01(\x03R\frequestCount\x12&\n" + "\x0fsum_duration_ms\x18\x0e \x01(\x03R\rsumDurationMs\x12!\n" + - "\fcost_nanousd\x18\x0f \x01(\x03R\vcostNanousd\"f\n" + + "\fcost_nanousd\x18\x0f \x01(\x03R\vcostNanousd\x12*\n" + + "\x11cache_read_tokens\x18\x10 \x01(\x03R\x0fcacheReadTokens\x12,\n" + + "\x12cache_write_tokens\x18\x11 \x01(\x03R\x10cacheWriteTokens\"f\n" + "\x17GetUsageSummaryResponse\x12-\n" + "\x04rows\x18\x01 \x03(\v2\x19.usage.v1.UsageSummaryRowR\x04rows\x12\x1c\n" + "\ttruncated\x18\x02 \x01(\bR\ttruncated\"\x9b\x01\n" + @@ -1116,7 +1174,7 @@ const file_usage_v1_usage_proto_rawDesc = "" + "meta_value\x18\x02 \x01(\tR\tmetaValue\x12)\n" + "\x11group_by_meta_key\x18\x03 \x01(\tR\x0egroupByMetaKey\x12\x1d\n" + "\n" + - "project_id\x18\x04 \x01(\tR\tprojectId\"\x92\x02\n" + + "project_id\x18\x04 \x01(\tR\tprojectId\"\xec\x02\n" + "\x0eUsageByMetaRow\x12\x1f\n" + "\vgroup_value\x18\x01 \x01(\tR\n" + "groupValue\x12\x16\n" + @@ -1126,7 +1184,10 @@ const file_usage_v1_usage_proto_rawDesc = "" + "\rtokens_output\x18\x05 \x01(\x03R\ftokensOutput\x12!\n" + "\ftokens_total\x18\x06 \x01(\x03R\vtokensTotal\x12#\n" + "\rrequest_count\x18\a \x01(\x03R\frequestCount\x12!\n" + - "\fcost_nanousd\x18\b \x01(\x03R\vcostNanousd\"d\n" + + "\fcost_nanousd\x18\b \x01(\x03R\vcostNanousd\x12*\n" + + "\x11cache_read_tokens\x18\t \x01(\x03R\x0fcacheReadTokens\x12,\n" + + "\x12cache_write_tokens\x18\n" + + " \x01(\x03R\x10cacheWriteTokens\"d\n" + "\x16GetUsageByMetaResponse\x12,\n" + "\x04rows\x18\x01 \x03(\v2\x18.usage.v1.UsageByMetaRowR\x04rows\x12\x1c\n" + "\ttruncated\x18\x02 \x01(\bR\ttruncated*p\n" + diff --git a/internal/database/migrations/008_cache_tokens.sql b/internal/database/migrations/008_cache_tokens.sql new file mode 100644 index 0000000..f476faf --- /dev/null +++ b/internal/database/migrations/008_cache_tokens.sql @@ -0,0 +1,15 @@ +-- +goose Up +-- +goose StatementBegin +-- Parts of tokens_input, kept apart because vendors bill cached input at a +-- fraction of the input rate (and cache writes above it). +ALTER TABLE usage_events + ADD COLUMN IF NOT EXISTS cache_read_tokens BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS cache_write_tokens BIGINT NOT NULL DEFAULT 0; +COMMENT ON COLUMN usage_events.cache_read_tokens IS 'Input tokens read from the prompt cache (part of tokens_input)'; +COMMENT ON COLUMN usage_events.cache_write_tokens IS 'Input tokens written to the prompt cache (part of tokens_input)'; +-- +goose StatementEnd + +-- +goose Down +-- +goose StatementBegin +ALTER TABLE usage_events DROP COLUMN IF EXISTS cache_read_tokens, DROP COLUMN IF EXISTS cache_write_tokens; +-- +goose StatementEnd diff --git a/internal/domain/event.go b/internal/domain/event.go index 19777e7..34d2c4d 100644 --- a/internal/domain/event.go +++ b/internal/domain/event.go @@ -69,6 +69,10 @@ type UsageEvent struct { VendorHTTPStatus int32 Meta map[string]string CostNanoUSD int64 + // Prompt-cache tokens, both part of TokensInput: read from the vendor's + // cache, and written to it, each priced differently from plain input. + CacheReadTokens int64 + CacheWriteTokens int64 } // Validation errors @@ -86,6 +90,7 @@ var ( ErrVendorTooLong = errors.New("vendor exceeds 64 characters") ErrModelTooLong = errors.New("model exceeds 128 characters") ErrTokensMismatch = errors.New("tokens_total must equal tokens_input + tokens_output") + ErrCacheTokens = errors.New("cache_read_tokens and cache_write_tokens must be >= 0 and together not exceed tokens_input") ErrNegativeCost = errors.New("cost_nanousd must not be negative") ErrOccurredAtRange = errors.New("occurred_at must be between 1900-01-01 and 2262-04-10") ErrErrorClassTooLong = errors.New("error_class exceeds 256 characters") @@ -150,6 +155,12 @@ func (e *UsageEvent) Validate() error { if e.CostNanoUSD < 0 { return ErrNegativeCost } + // Checked only when sent, so events without them validate as before; the + // subtraction cannot overflow once 0 <= read <= input. + if (e.CacheReadTokens != 0 || e.CacheWriteTokens != 0) && (e.CacheReadTokens < 0 || e.CacheWriteTokens < 0 || + e.CacheReadTokens > e.TokensInput || e.CacheWriteTokens > e.TokensInput-e.CacheReadTokens) { + return ErrCacheTokens + } if len(e.ErrorClass) > MaxStringLength { return ErrErrorClassTooLong } diff --git a/internal/domain/meta_usage.go b/internal/domain/meta_usage.go index 0addb2a..2037717 100644 --- a/internal/domain/meta_usage.go +++ b/internal/domain/meta_usage.go @@ -25,4 +25,7 @@ type MetaUsageRow struct { TokensTotal int64 RequestCount int64 CostNanoUSD int64 + // Parts of TokensInput served from and written to the prompt cache. + CacheReadTokens int64 + CacheWriteTokens int64 } diff --git a/internal/domain/summary.go b/internal/domain/summary.go index 3d0ef06..379f112 100644 --- a/internal/domain/summary.go +++ b/internal/domain/summary.go @@ -43,4 +43,7 @@ type SummaryRow struct { RequestCount int64 SumDurationMs int64 CostNanoUSD int64 + // Parts of TokensInput served from and written to the prompt cache. + CacheReadTokens int64 + CacheWriteTokens int64 } diff --git a/internal/repository/outbox.go b/internal/repository/outbox.go index 8402021..44212fc 100644 --- a/internal/repository/outbox.go +++ b/internal/repository/outbox.go @@ -179,7 +179,8 @@ func (r *Repository) OutboxBacklog(ctx context.Context) (int64, time.Duration, e const eventColumns = `event_id, created_at, occurred_at, user_id, project_id, action_id, vendor, model, status, duration_ms, tokens_input, tokens_output, tokens_total, - COALESCE(error_class, ''), COALESCE(vendor_http_status, 0), meta, cost_nanousd` + COALESCE(error_class, ''), COALESCE(vendor_http_status, 0), meta, cost_nanousd, + cache_read_tokens, cache_write_tokens` func (r *Repository) eventsByID(ctx context.Context, ids []uuid.UUID) ([]*domain.UsageEvent, error) { rows, err := r.pool.Query(ctx, `SELECT `+eventColumns+` FROM usage_events WHERE event_id = ANY($1)`, ids) @@ -200,7 +201,8 @@ func scanEvents(rows pgx.Rows) ([]*domain.UsageEvent, error) { ) if err := rows.Scan(&e.EventID, &e.CreatedAt, &e.OccurredAt, &e.UserID, &e.ProjectID, &e.ActionID, &e.Vendor, &e.Model, &status, &e.DurationMs, &e.TokensInput, &e.TokensOutput, &e.TokensTotal, - &e.ErrorClass, &e.VendorHTTPStatus, &meta, &e.CostNanoUSD); err != nil { + &e.ErrorClass, &e.VendorHTTPStatus, &meta, &e.CostNanoUSD, + &e.CacheReadTokens, &e.CacheWriteTokens); err != nil { return nil, fmt.Errorf("scan outbox event: %w", err) } e.Status = domain.Status(status) diff --git a/internal/repository/outbox_test.go b/internal/repository/outbox_test.go index aa5cfb1..90942f0 100644 --- a/internal/repository/outbox_test.go +++ b/internal/repository/outbox_test.go @@ -19,7 +19,7 @@ func outboxEvent(meta map[string]string) *domain.UsageEvent { EventID: uuid.New(), OccurredAt: &now, UserID: "u", ProjectID: "p", ActionID: "a", Vendor: "openai", Model: "m", Status: domain.StatusError, DurationMs: 12, TokensInput: 3, TokensOutput: 4, TokensTotal: 7, ErrorClass: "rate_limited", VendorHTTPStatus: 429, - Meta: meta, CostNanoUSD: 99, + Meta: meta, CostNanoUSD: 99, CacheReadTokens: 2, CacheWriteTokens: 1, } } @@ -87,6 +87,8 @@ func TestOutboxClaimAckAndTakeOver(t *testing.T) { require.Equal(t, first.ErrorClass, stored.ErrorClass) require.Equal(t, first.VendorHTTPStatus, stored.VendorHTTPStatus) require.Equal(t, first.CostNanoUSD, stored.CostNanoUSD) + require.Equal(t, first.CacheReadTokens, stored.CacheReadTokens) + require.Equal(t, first.CacheWriteTokens, stored.CacheWriteTokens) require.False(t, stored.CreatedAt.IsZero()) require.NoError(t, repo.AckOutbox(ctx, b1.ID)) diff --git a/internal/repository/postgres_test.go b/internal/repository/postgres_test.go index d97ef70..8c813fd 100644 --- a/internal/repository/postgres_test.go +++ b/internal/repository/postgres_test.go @@ -198,18 +198,20 @@ func TestCostAggregation(t *testing.T) { newEvent := func(model string, cost int64) *domain.UsageEvent { return &domain.UsageEvent{ - EventID: uuid.New(), - UserID: "u1", - ProjectID: "p1", - ActionID: "a1", - Vendor: "openai", - Model: model, - Status: domain.StatusSuccess, - TokensInput: 10, - TokensOutput: 20, - TokensTotal: 30, - Meta: map[string]string{"job_id": "job-cost"}, - CostNanoUSD: cost, + EventID: uuid.New(), + UserID: "u1", + ProjectID: "p1", + ActionID: "a1", + Vendor: "openai", + Model: model, + Status: domain.StatusSuccess, + TokensInput: 10, + TokensOutput: 20, + TokensTotal: 30, + Meta: map[string]string{"job_id": "job-cost"}, + CostNanoUSD: cost, + CacheReadTokens: 6, + CacheWriteTokens: 2, } } @@ -229,19 +231,28 @@ func TestCostAggregation(t *testing.T) { }) require.NoError(t, err) byModel := map[string]int64{} + var cacheRead, cacheWrite int64 for _, r := range summary { byModel[r.Model] += r.CostNanoUSD + cacheRead += r.CacheReadTokens + cacheWrite += r.CacheWriteTokens } + assert.Equal(t, int64(18), cacheRead, "InsertOne and InsertBatch both store cache tokens") + assert.Equal(t, int64(6), cacheWrite) assert.Equal(t, int64(3_000_000), byModel["gpt-4o"]) assert.Equal(t, int64(150_000), byModel["gpt-4o-mini"]) metaRows, _, err := repo.QueryUsageByMeta(ctx, domain.MetaUsageParams{MetaKey: "job_id", MetaValue: "job-cost"}) require.NoError(t, err) - var total int64 + var total, metaCacheRead, metaCacheWrite int64 for _, r := range metaRows { total += r.CostNanoUSD + metaCacheRead += r.CacheReadTokens + metaCacheWrite += r.CacheWriteTokens } assert.Equal(t, int64(3_150_000), total) + assert.Equal(t, int64(18), metaCacheRead) + assert.Equal(t, int64(6), metaCacheWrite) } // A retried batch inserts nothing new, and InsertBatch says so by ID, so the diff --git a/internal/repository/repository.go b/internal/repository/repository.go index dd99986..31011f8 100644 --- a/internal/repository/repository.go +++ b/internal/repository/repository.go @@ -44,12 +44,14 @@ func (r *Repository) InsertOne(ctx context.Context, event *domain.UsageEvent) (b event_id, occurred_at, user_id, project_id, action_id, vendor, model, status, duration_ms, tokens_input, tokens_output, tokens_total, - error_class, vendor_http_status, meta, cost_nanousd + error_class, vendor_http_status, meta, cost_nanousd, + cache_read_tokens, cache_write_tokens ) VALUES ( $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, - $13, $14, $15, $16 + $13, $14, $15, $16, + $17, $18 ) ON CONFLICT (event_id) DO NOTHING ` @@ -75,6 +77,8 @@ func (r *Repository) InsertOne(ctx context.Context, event *domain.UsageEvent) (b nullableInt32(event.VendorHTTPStatus), metaJSON, event.CostNanoUSD, + event.CacheReadTokens, + event.CacheWriteTokens, ) if err != nil { return false, fmt.Errorf("insert event: %w", err) @@ -136,6 +140,8 @@ func (r *Repository) InsertBatch(ctx context.Context, events []*domain.UsageEven nullableInt32(event.VendorHTTPStatus), metaJSON, event.CostNanoUSD, + event.CacheReadTokens, + event.CacheWriteTokens, }) } @@ -148,6 +154,7 @@ func (r *Repository) InsertBatch(ctx context.Context, events []*domain.UsageEven "vendor", "model", "status", "duration_ms", "tokens_input", "tokens_output", "tokens_total", "error_class", "vendor_http_status", "meta", "cost_nanousd", + "cache_read_tokens", "cache_write_tokens", }, pgx.CopyFromRows(rows), ) @@ -233,6 +240,8 @@ func (r *Repository) QuerySummary(ctx context.Context, params domain.SummaryPara "COUNT(*) AS request_count", "SUM(duration_ms) AS sum_duration_ms", "SUM(cost_nanousd) AS cost_nanousd", + "SUM(cache_read_tokens) AS cache_read_tokens", + "SUM(cache_write_tokens) AS cache_write_tokens", ) // Build WHERE clause @@ -332,6 +341,8 @@ func (r *Repository) QuerySummary(ctx context.Context, params domain.SummaryPara &row.RequestCount, &row.SumDurationMs, &row.CostNanoUSD, + &row.CacheReadTokens, + &row.CacheWriteTokens, ) if err := rows.Scan(scanDest...); err != nil { @@ -389,7 +400,9 @@ func (r *Repository) QueryUsageByMeta(ctx context.Context, params domain.MetaUsa SUM(tokens_output) AS tokens_output, SUM(tokens_total) AS tokens_total, COUNT(*) AS request_count, - SUM(cost_nanousd) AS cost_nanousd + SUM(cost_nanousd) AS cost_nanousd, + SUM(cache_read_tokens) AS cache_read_tokens, + SUM(cache_write_tokens) AS cache_write_tokens FROM usage_events %s GROUP BY group_value, vendor, model @@ -409,7 +422,7 @@ func (r *Repository) QueryUsageByMeta(ctx context.Context, params domain.MetaUsa if err := rows.Scan( &row.GroupValue, &row.Vendor, &row.Model, &row.TokensInput, &row.TokensOutput, &row.TokensTotal, &row.RequestCount, - &row.CostNanoUSD, + &row.CostNanoUSD, &row.CacheReadTokens, &row.CacheWriteTokens, ); err != nil { return nil, false, fmt.Errorf("scan meta usage row: %w", err) } diff --git a/internal/server/grpc.go b/internal/server/grpc.go index 346a47a..322405b 100644 --- a/internal/server/grpc.go +++ b/internal/server/grpc.go @@ -173,19 +173,21 @@ func (s *GRPCServer) GetUsageSummary(ctx context.Context, req *pb.GetUsageSummar pbRows := make([]*pb.UsageSummaryRow, 0, len(rows)) for _, row := range rows { pbRows = append(pbRows, &pb.UsageSummaryRow{ - Bucket: timestamppb.New(row.Bucket), - ProjectId: row.ProjectID, - UserId: row.UserID, - Vendor: row.Vendor, - Model: row.Model, - ActionId: row.ActionID, - Status: domainStatusToProto(row.Status), - TokensInput: row.TokensInput, - TokensOutput: row.TokensOutput, - TokensTotal: row.TokensTotal, - RequestCount: row.RequestCount, - SumDurationMs: row.SumDurationMs, - CostNanousd: row.CostNanoUSD, + Bucket: timestamppb.New(row.Bucket), + ProjectId: row.ProjectID, + UserId: row.UserID, + Vendor: row.Vendor, + Model: row.Model, + ActionId: row.ActionID, + Status: domainStatusToProto(row.Status), + TokensInput: row.TokensInput, + TokensOutput: row.TokensOutput, + TokensTotal: row.TokensTotal, + RequestCount: row.RequestCount, + SumDurationMs: row.SumDurationMs, + CostNanousd: row.CostNanoUSD, + CacheReadTokens: row.CacheReadTokens, + CacheWriteTokens: row.CacheWriteTokens, }) } @@ -217,14 +219,16 @@ func (s *GRPCServer) GetUsageByMeta(ctx context.Context, req *pb.GetUsageByMetaR pbRows := make([]*pb.UsageByMetaRow, 0, len(rows)) for _, row := range rows { pbRows = append(pbRows, &pb.UsageByMetaRow{ - GroupValue: row.GroupValue, - Vendor: row.Vendor, - Model: row.Model, - TokensInput: row.TokensInput, - TokensOutput: row.TokensOutput, - TokensTotal: row.TokensTotal, - RequestCount: row.RequestCount, - CostNanousd: row.CostNanoUSD, + GroupValue: row.GroupValue, + Vendor: row.Vendor, + Model: row.Model, + TokensInput: row.TokensInput, + TokensOutput: row.TokensOutput, + TokensTotal: row.TokensTotal, + RequestCount: row.RequestCount, + CostNanousd: row.CostNanoUSD, + CacheReadTokens: row.CacheReadTokens, + CacheWriteTokens: row.CacheWriteTokens, }) } return &pb.GetUsageByMetaResponse{ @@ -250,6 +254,8 @@ func protoToEvent(pbEvent *pb.UsageEvent) (*domain.UsageEvent, error) { VendorHTTPStatus: pbEvent.VendorHttpStatus, Meta: pbEvent.Meta, CostNanoUSD: pbEvent.CostNanousd, + CacheReadTokens: pbEvent.CacheReadTokens, + CacheWriteTokens: pbEvent.CacheWriteTokens, } // Parse event_id diff --git a/internal/server/http.go b/internal/server/http.go index 77e958c..9f7959a 100644 --- a/internal/server/http.go +++ b/internal/server/http.go @@ -137,6 +137,8 @@ type recordEventRequest struct { VendorHTTPStatus int32 `json:"vendor_http_status,omitempty"` Meta map[string]string `json:"meta,omitempty"` CostNanoUSD int64 `json:"cost_nanousd,omitempty"` + CacheReadTokens int64 `json:"cache_read_tokens,omitempty"` + CacheWriteTokens int64 `json:"cache_write_tokens,omitempty"` } type recordEventResponse struct { @@ -183,6 +185,9 @@ type summaryRow struct { RequestCount int64 `json:"request_count"` SumDurationMs int64 `json:"sum_duration_ms"` CostNanoUSD int64 `json:"cost_nanousd"` + // Parts of tokens_input read from and written to the prompt cache. + CacheReadTokens int64 `json:"cache_read_tokens"` + CacheWriteTokens int64 `json:"cache_write_tokens"` } type getSummaryResponse struct { @@ -324,19 +329,21 @@ func (s *HTTPServer) handleGetSummary(w http.ResponseWriter, r *http.Request) { status = row.Status.String() } respRows = append(respRows, summaryRow{ - Bucket: row.Bucket.Format(time.RFC3339), - ProjectID: row.ProjectID, - UserID: row.UserID, - Vendor: row.Vendor, - Model: row.Model, - ActionID: row.ActionID, - Status: status, - TokensInput: row.TokensInput, - TokensOutput: row.TokensOutput, - TokensTotal: row.TokensTotal, - RequestCount: row.RequestCount, - SumDurationMs: row.SumDurationMs, - CostNanoUSD: row.CostNanoUSD, + Bucket: row.Bucket.Format(time.RFC3339), + ProjectID: row.ProjectID, + UserID: row.UserID, + Vendor: row.Vendor, + Model: row.Model, + ActionID: row.ActionID, + Status: status, + TokensInput: row.TokensInput, + TokensOutput: row.TokensOutput, + TokensTotal: row.TokensTotal, + RequestCount: row.RequestCount, + SumDurationMs: row.SumDurationMs, + CostNanoUSD: row.CostNanoUSD, + CacheReadTokens: row.CacheReadTokens, + CacheWriteTokens: row.CacheWriteTokens, }) } @@ -363,6 +370,8 @@ func httpRequestToEvent(req *recordEventRequest) (*domain.UsageEvent, error) { Meta: req.Meta, OccurredAt: req.OccurredAt, CostNanoUSD: req.CostNanoUSD, + CacheReadTokens: req.CacheReadTokens, + CacheWriteTokens: req.CacheWriteTokens, } if err := event.EventID.UnmarshalText([]byte(req.EventID)); err != nil { diff --git a/internal/server/server_test.go b/internal/server/server_test.go index 0756de6..29356d2 100644 --- a/internal/server/server_test.go +++ b/internal/server/server_test.go @@ -152,6 +152,26 @@ func TestProtoToEvent_Cost(t *testing.T) { assert.ErrorContains(t, err, "cost_nanousd must not be negative") } +func TestProtoToEvent_CacheTokens(t *testing.T) { + event := func(read, write int64) *pb.UsageEvent { + return &pb.UsageEvent{ + EventId: uuid.NewString(), UserId: "u1", ProjectId: "p1", ActionId: "a1", + Vendor: "anthropic", Model: "claude", Status: pb.Status_STATUS_SUCCESS, + TokensInput: 100, TokensOutput: 5, CacheReadTokens: read, CacheWriteTokens: write, + } + } + ev, err := protoToEvent(event(70, 30)) + require.NoError(t, err) + assert.Equal(t, int64(70), ev.CacheReadTokens) + assert.Equal(t, int64(30), ev.CacheWriteTokens) + assert.Equal(t, int64(105), ev.TokensTotal, "cache tokens are part of tokens_input, not added") + + for _, bad := range [][2]int64{{71, 30}, {-1, 0}, {0, -1}, {101, 0}, {0, 101}} { + _, err = protoToEvent(event(bad[0], bad[1])) + assert.ErrorContains(t, err, "cache_read_tokens", "read=%d write=%d", bad[0], bad[1]) + } +} + func testVerifier(t *testing.T) *auth.Verifier { t.Helper() sum := func(s string) string { h := sha256.Sum256([]byte(s)); return hex.EncodeToString(h[:]) } @@ -341,3 +361,17 @@ func TestSummaryOmitsUngroupedStatus(t *testing.T) { assert.NotContains(t, get("&group_by_model=true"), `"status"`) assert.Contains(t, get("&group_by_status=true"), `"status":"success"`) } + +func TestHTTP_CacheTokens(t *testing.T) { + var req recordEventRequest + require.NoError(t, json.Unmarshal([]byte(`{"event_id":"`+uuid.NewString()+`","user_id":"u","project_id":"p", + "action_id":"a","vendor":"anthropic","model":"claude","status":"success", + "tokens_input":100,"tokens_output":5,"cache_read_tokens":60,"cache_write_tokens":40}`), &req)) + ev, err := httpRequestToEvent(&req) + require.NoError(t, err) + assert.Equal(t, [2]int64{60, 40}, [2]int64{ev.CacheReadTokens, ev.CacheWriteTokens}) + + out, err := json.Marshal(summaryRow{CacheReadTokens: 7, CacheWriteTokens: 3}) + require.NoError(t, err) + assert.Contains(t, string(out), `"cache_read_tokens":7,"cache_write_tokens":3`) +} diff --git a/internal/storage/clickhouse/client.go b/internal/storage/clickhouse/client.go index db749ee..4198ec3 100644 --- a/internal/storage/clickhouse/client.go +++ b/internal/storage/clickhouse/client.go @@ -198,6 +198,8 @@ func (c *Client) QuerySummary(ctx context.Context, params domain.SummaryParams) "count() AS request_count", "sum(duration_ms) AS sum_duration_ms", "sum(cost_nanousd) AS cost_nanousd", + "sum(cache_read_tokens) AS cache_read_tokens", + "sum(cache_write_tokens) AS cache_write_tokens", ) var whereClauses []string @@ -289,6 +291,8 @@ func (c *Client) QuerySummary(ctx context.Context, params domain.SummaryParams) &reqCount, &row.SumDurationMs, &row.CostNanoUSD, + &row.CacheReadTokens, + &row.CacheWriteTokens, ) if err := rows.Scan(scanTargets...); err != nil { @@ -343,7 +347,9 @@ func (c *Client) QueryUsageByMeta(ctx context.Context, params domain.MetaUsagePa sum(tokens_output) AS tokens_output, sum(tokens_total) AS tokens_total, count() AS request_count, - sum(cost_nanousd) AS cost_nanousd + sum(cost_nanousd) AS cost_nanousd, + sum(cache_read_tokens) AS cache_read_tokens, + sum(cache_write_tokens) AS cache_write_tokens FROM usage_events WHERE %s GROUP BY group_value, vendor, model @@ -375,6 +381,8 @@ func (c *Client) QueryUsageByMeta(ctx context.Context, params domain.MetaUsagePa &row.TokensTotal, &reqCount, &row.CostNanoUSD, + &row.CacheReadTokens, + &row.CacheWriteTokens, ); err != nil { return nil, false, fmt.Errorf("clickhouse scan meta usage row: %w", err) } diff --git a/internal/storage/clickhouse/client_test.go b/internal/storage/clickhouse/client_test.go index 8265a79..d1201ee 100644 --- a/internal/storage/clickhouse/client_test.go +++ b/internal/storage/clickhouse/client_test.go @@ -99,18 +99,20 @@ func TestClickHouse_CostAggregation(t *testing.T) { newEvent := func(model string, cost int64) *domain.UsageEvent { return &domain.UsageEvent{ - EventID: uuid.New(), - UserID: "u1", - ProjectID: "p1", - ActionID: "a1", - Vendor: "openai", - Model: model, - Status: domain.StatusSuccess, - TokensInput: 10, - TokensOutput: 20, - TokensTotal: 30, - Meta: map[string]string{"job_id": "job-cost"}, - CostNanoUSD: cost, + EventID: uuid.New(), + UserID: "u1", + ProjectID: "p1", + ActionID: "a1", + Vendor: "openai", + Model: model, + Status: domain.StatusSuccess, + TokensInput: 10, + TokensOutput: 20, + TokensTotal: 30, + Meta: map[string]string{"job_id": "job-cost"}, + CostNanoUSD: cost, + CacheReadTokens: 6, + CacheWriteTokens: 1, } } r := NewReplicator(nil, client.Conn(), 100, time.Hour) @@ -128,19 +130,26 @@ func TestClickHouse_CostAggregation(t *testing.T) { }) require.NoError(t, err) byModel := map[string]int64{} + var cacheRead, cacheWrite int64 for _, r := range summary { byModel[r.Model] += r.CostNanoUSD + cacheRead += r.CacheReadTokens + cacheWrite += r.CacheWriteTokens } assert.Equal(t, int64(3_000_000), byModel["gpt-4o"]) assert.Equal(t, int64(150_000), byModel["gpt-4o-mini"]) + assert.Equal(t, [2]int64{18, 3}, [2]int64{cacheRead, cacheWrite}) metaRows, _, err := client.QueryUsageByMeta(ctx, domain.MetaUsageParams{MetaKey: "job_id", MetaValue: "job-cost"}) require.NoError(t, err) - var total int64 + var total, metaRead, metaWrite int64 for _, r := range metaRows { total += r.CostNanoUSD + metaRead += r.CacheReadTokens + metaWrite += r.CacheWriteTokens } assert.Equal(t, int64(3_150_000), total) + assert.Equal(t, [2]int64{18, 3}, [2]int64{metaRead, metaWrite}) } // A least-privilege user without CREATE DATABASE must still connect to the diff --git a/internal/storage/clickhouse/migrations/004_add_cache_tokens.sql b/internal/storage/clickhouse/migrations/004_add_cache_tokens.sql new file mode 100644 index 0000000..62b8754 --- /dev/null +++ b/internal/storage/clickhouse/migrations/004_add_cache_tokens.sql @@ -0,0 +1 @@ +ALTER TABLE usage_events ADD COLUMN IF NOT EXISTS cache_read_tokens Int64 DEFAULT 0, ADD COLUMN IF NOT EXISTS cache_write_tokens Int64 DEFAULT 0 diff --git a/internal/storage/clickhouse/replicator.go b/internal/storage/clickhouse/replicator.go index 4bcae5e..6c8df12 100644 --- a/internal/storage/clickhouse/replicator.go +++ b/internal/storage/clickhouse/replicator.go @@ -217,7 +217,8 @@ func (r *Replicator) insert(ctx context.Context, batch uuid.UUID, events []*doma user_id, project_id, action_id, vendor, model, status, duration_ms, tokens_input, tokens_output, tokens_total, - error_class, vendor_http_status, meta, cost_nanousd + error_class, vendor_http_status, meta, cost_nanousd, + cache_read_tokens, cache_write_tokens ) `) if err != nil { @@ -265,6 +266,8 @@ func (r *Replicator) insert(ctx context.Context, batch uuid.UUID, events []*doma vendorStatus, meta, ev.CostNanoUSD, + ev.CacheReadTokens, + ev.CacheWriteTokens, ) if err != nil { _ = b.Abort() diff --git a/internal/storage/clickhouse/replicator_test.go b/internal/storage/clickhouse/replicator_test.go index 6d0ab54..42127cf 100644 --- a/internal/storage/clickhouse/replicator_test.go +++ b/internal/storage/clickhouse/replicator_test.go @@ -82,7 +82,7 @@ func (b *poisonBatch) Send() error { return b.Batch.Send() } -type totals struct{ requests, tokens, cost int64 } +type totals struct{ requests, tokens, cost, cacheRead, cacheWrite int64 } func pgTotals(t *testing.T, repo *repository.Repository, from, to time.Time) totals { rows, _, err := repo.QuerySummary(context.Background(), domain.SummaryParams{From: from, To: to}) @@ -92,6 +92,8 @@ func pgTotals(t *testing.T, repo *repository.Repository, from, to time.Time) tot out.requests += r.RequestCount out.tokens += r.TokensTotal out.cost += r.CostNanoUSD + out.cacheRead += r.CacheReadTokens + out.cacheWrite += r.CacheWriteTokens } return out } @@ -104,6 +106,8 @@ func chTotals(t *testing.T, c *Client, from, to time.Time) totals { out.requests += r.RequestCount out.tokens += r.TokensTotal out.cost += r.CostNanoUSD + out.cacheRead += r.CacheReadTokens + out.cacheWrite += r.CacheWriteTokens } return out } @@ -131,6 +135,7 @@ func (s *stack) record(t *testing.T, n int, model string) { EventID: uuid.New(), OccurredAt: &at, UserID: "u", ProjectID: "p", ActionID: "a", Vendor: "openai", Model: model, Status: domain.StatusSuccess, TokensInput: int64(i % 7), TokensOutput: 3, TokensTotal: int64(i%7) + 3, CostNanoUSD: int64(i), + CacheReadTokens: int64(i % 7 / 2), CacheWriteTokens: int64(i % 7 % 2), }) } _, err := s.repo.InsertBatch(context.Background(), batch) diff --git a/proto/usage/v1/usage.proto b/proto/usage/v1/usage.proto index 35f96cf..02b2888 100644 --- a/proto/usage/v1/usage.proto +++ b/proto/usage/v1/usage.proto @@ -65,6 +65,14 @@ message UsageEvent { // Cost of the call in nano-USD (1e-9 USD), computed by the caller from its // price catalog at settle time; stored as an immutable fact (optional, >= 0) int64 cost_nanousd = 16; + + // Input tokens served from the vendor's prompt cache (optional, >= 0). Part of + // tokens_input, not added to it; usually billed at a fraction of the input rate. + int64 cache_read_tokens = 17; + + // Input tokens written to the vendor's prompt cache (optional, >= 0). Part of + // tokens_input; cache_read_tokens + cache_write_tokens must not exceed it. + int64 cache_write_tokens = 18; } // RecordUsageEventRequest is the request for recording a single usage event @@ -139,6 +147,9 @@ message UsageSummaryRow { int64 sum_duration_ms = 14; // Summed cost in nano-USD (1e-9 USD) int64 cost_nanousd = 15; + // Summed parts of tokens_input read from and written to the prompt cache + int64 cache_read_tokens = 16; + int64 cache_write_tokens = 17; } // GetUsageSummaryResponse is the response for getting usage summary @@ -178,6 +189,9 @@ message UsageByMetaRow { int64 request_count = 7; // Summed cost in nano-USD (1e-9 USD) int64 cost_nanousd = 8; + // Summed parts of tokens_input read from and written to the prompt cache + int64 cache_read_tokens = 9; + int64 cache_write_tokens = 10; } // GetUsageByMetaResponse is the response for meta-scoped aggregation