summaryrefslogtreecommitdiff
path: root/internal/controlplane/usage.go
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--internal/controlplane/usage.go82
1 files changed, 67 insertions, 15 deletions
diff --git a/internal/controlplane/usage.go b/internal/controlplane/usage.go
index ebabaf6..05bc155 100644
--- a/internal/controlplane/usage.go
+++ b/internal/controlplane/usage.go
@@ -2,6 +2,9 @@ package controlplane
import (
"context"
+ "encoding/base64"
+ "encoding/json"
+ "errors"
"fmt"
"strings"
"time"
@@ -25,6 +28,14 @@ type UsageQuery struct {
From time.Time
To time.Time
Limit int
+ Cursor string
+}
+
+var ErrInvalidUsageCursor = errors.New("invalid usage cursor")
+
+type usageCursor struct {
+ StartedAt time.Time `json:"started_at"`
+ RequestID string `json:"request_id"`
}
func (s *Store) RecordUsage(ctx context.Context, event domain.UsageEvent) error {
@@ -37,13 +48,13 @@ func (s *Store) RecordUsage(ctx context.Context, event domain.UsageEvent) error
INSERT INTO usage_events (
request_id, tenant_id, project_id, key_id, public_model, provider_id, upstream_model,
protocol, stream, status_code, success, error_type, attempts, started_at, duration_ms,
- input_tokens, output_tokens, total_tokens, cache_creation_input_tokens, cache_read_input_tokens,
+ ttft_ms, input_tokens, output_tokens, total_tokens, cache_creation_input_tokens, cache_read_input_tokens,
cost_micros, charged_micros, uncollected_micros, usage_reported, metering_status)
- VALUES ($1,$2,$3,$4,$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,0,0,0,$21,$22)
+ VALUES ($1,$2,$3,$4,$5,NULLIF($6,''),NULLIF($7,''),$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,0,0,0,$22,$23)
ON CONFLICT (request_id) DO NOTHING`, event.RequestID, event.TenantID, event.ProjectID, event.KeyID,
event.PublicModel, event.ProviderID, event.UpstreamModel, string(event.Protocol), event.Stream,
event.StatusCode, event.Success, event.ErrorType, event.Attempts, event.StartedAt, event.DurationMS,
- event.Usage.InputTokens, event.Usage.OutputTokens, event.Usage.TotalTokens,
+ event.TTFTMS, event.Usage.InputTokens, event.Usage.OutputTokens, event.Usage.TotalTokens,
event.Usage.CacheCreationInputTokens, event.Usage.CacheReadInputTokens, event.UsageReported, usageMeteringStatus(event))
if err != nil {
return fmt.Errorf("persist usage event: %w", err)
@@ -104,13 +115,13 @@ func boolInt(value bool) int {
return 0
}
-func (s *Store) ListUsage(ctx context.Context, query UsageQuery) ([]UsageRecord, error) {
+func (s *Store) ListUsage(ctx context.Context, query UsageQuery) (UsagePage, error) {
limit := query.Limit
if limit < 1 || limit > 1000 {
limit = 200
}
where := []string{"1=1"}
- args := make([]any, 0, 13)
+ args := make([]any, 0, 16)
index := 1
for _, item := range []struct{ value, clause string }{{query.TenantID, "tenant_id=$"}, {query.ProjectID, "project_id=$"}, {query.KeyID, "key_id=$"}, {query.Model, "public_model=$"}, {query.RequestID, "request_id=$"}} {
if strings.TrimSpace(item.value) != "" {
@@ -151,31 +162,68 @@ func (s *Store) ListUsage(ctx context.Context, query UsageQuery) ([]UsageRecord,
args = append(args, query.To)
index++
}
- args = append(args, limit)
+ if strings.TrimSpace(query.Cursor) != "" {
+ cursor, err := decodeUsageCursor(query.Cursor)
+ if err != nil {
+ return UsagePage{}, err
+ }
+ where = append(where, "(started_at, request_id) < ($"+fmt.Sprint(index)+",$"+fmt.Sprint(index+1)+")")
+ args = append(args, cursor.StartedAt, cursor.RequestID)
+ index += 2
+ }
+ args = append(args, limit+1)
rows, err := s.db.Query(ctx, `SELECT request_id, tenant_id::text, project_id::text,
COALESCE((SELECT name FROM projects p WHERE p.id=usage_events.project_id),''), key_id::text,
COALESCE((SELECT name FROM api_keys k WHERE k.id=usage_events.key_id),''), public_model,
COALESCE(provider_id,''), COALESCE((SELECT name FROM providers p WHERE p.id::text=usage_events.provider_id),''),
COALESCE(upstream_model,''), protocol, stream, status_code, success, error_type,
- attempts, started_at, duration_ms, input_tokens, output_tokens, total_tokens, cache_creation_input_tokens,
+ attempts, started_at, duration_ms, ttft_ms, input_tokens, output_tokens, total_tokens, cache_creation_input_tokens,
cache_read_input_tokens, cost_micros, charged_micros, uncollected_micros, usage_reported, metering_status FROM usage_events WHERE `+
- strings.Join(where, " AND ")+` ORDER BY created_at DESC LIMIT $`+fmt.Sprint(index), args...)
+ strings.Join(where, " AND ")+` ORDER BY started_at DESC, request_id DESC LIMIT $`+fmt.Sprint(index), args...)
if err != nil {
- return nil, fmt.Errorf("query usage events: %w", err)
+ return UsagePage{}, fmt.Errorf("query usage events: %w", err)
}
defer rows.Close()
- result := make([]UsageRecord, 0)
+ result := make([]UsageRecord, 0, limit+1)
for rows.Next() {
var item UsageRecord
if err := rows.Scan(&item.RequestID, &item.TenantID, &item.ProjectID, &item.ProjectName, &item.KeyID, &item.KeyName, &item.PublicModel, &item.ProviderID, &item.ProviderName, &item.UpstreamModel,
&item.Protocol, &item.Stream, &item.StatusCode, &item.Success, &item.ErrorType, &item.Attempts, &item.StartedAt, &item.DurationMS,
- &item.InputTokens, &item.OutputTokens, &item.TotalTokens, &item.CacheCreationInputTokens, &item.CacheReadInputTokens,
+ &item.TTFTMS, &item.InputTokens, &item.OutputTokens, &item.TotalTokens, &item.CacheCreationInputTokens, &item.CacheReadInputTokens,
&item.CostMicros, &item.ChargedMicros, &item.UncollectedMicros, &item.UsageReported, &item.MeteringStatus); err != nil {
- return nil, fmt.Errorf("scan usage event: %w", err)
+ return UsagePage{}, fmt.Errorf("scan usage event: %w", err)
}
result = append(result, item)
}
- return result, rows.Err()
+ if err := rows.Err(); err != nil {
+ return UsagePage{}, err
+ }
+ page := UsagePage{Data: result}
+ if len(result) > limit {
+ page.Data = result[:limit]
+ page.NextCursor = encodeUsageCursor(page.Data[len(page.Data)-1])
+ }
+ return page, nil
+}
+
+func encodeUsageCursor(record UsageRecord) string {
+ payload, _ := json.Marshal(usageCursor{StartedAt: record.StartedAt.UTC(), RequestID: record.RequestID})
+ return base64.RawURLEncoding.EncodeToString(payload)
+}
+
+func decodeUsageCursor(raw string) (usageCursor, error) {
+ if len(raw) > 2048 {
+ return usageCursor{}, ErrInvalidUsageCursor
+ }
+ payload, err := base64.RawURLEncoding.DecodeString(raw)
+ if err != nil {
+ return usageCursor{}, ErrInvalidUsageCursor
+ }
+ var cursor usageCursor
+ if json.Unmarshal(payload, &cursor) != nil || cursor.StartedAt.IsZero() || strings.TrimSpace(cursor.RequestID) == "" {
+ return usageCursor{}, ErrInvalidUsageCursor
+ }
+ return cursor, nil
}
func (s *Store) UsageDaily(ctx context.Context, query UsageQuery) ([]UsageDailyPoint, error) {
@@ -225,7 +273,10 @@ func (s *Store) UsageDaily(ctx context.Context, query UsageQuery) ([]UsageDailyP
count(*) FILTER (WHERE success), COALESCE(sum(input_tokens),0), COALESCE(sum(output_tokens),0),
COALESCE(sum(total_tokens),0), COALESCE(sum(charged_micros),0), COALESCE(sum(uncollected_micros),0),
COALESCE(round(avg(duration_ms)),0)::bigint,
- COALESCE(round(percentile_cont(0.95) WITHIN GROUP (ORDER BY duration_ms)),0)::bigint
+ COALESCE(round(percentile_cont(0.50) WITHIN GROUP (ORDER BY duration_ms)),0)::bigint,
+ COALESCE(round(percentile_cont(0.95) WITHIN GROUP (ORDER BY duration_ms)),0)::bigint,
+ COALESCE(round(percentile_cont(0.50) WITHIN GROUP (ORDER BY ttft_ms) FILTER (WHERE ttft_ms > 0)),0)::bigint,
+ COALESCE(round(percentile_cont(0.95) WITHIN GROUP (ORDER BY ttft_ms) FILTER (WHERE ttft_ms > 0)),0)::bigint
FROM usage_events WHERE `+strings.Join(where, " AND ")+` GROUP BY 1 ORDER BY 1`, args...)
if err != nil {
return nil, fmt.Errorf("query daily usage: %w", err)
@@ -235,7 +286,8 @@ func (s *Store) UsageDaily(ctx context.Context, query UsageQuery) ([]UsageDailyP
for rows.Next() {
var item UsageDailyPoint
if err := rows.Scan(&item.Day, &item.RequestCount, &item.SuccessfulRequests, &item.InputTokens, &item.OutputTokens,
- &item.TotalTokens, &item.ChargedMicros, &item.UncollectedMicros, &item.AverageDurationMS, &item.P95DurationMS); err != nil {
+ &item.TotalTokens, &item.ChargedMicros, &item.UncollectedMicros, &item.AverageDurationMS, &item.P50DurationMS,
+ &item.P95DurationMS, &item.P50TTFTMS, &item.P95TTFTMS); err != nil {
return nil, fmt.Errorf("scan daily usage: %w", err)
}
result = append(result, item)
generated by cgit v1.2.3 (git 2.46.0) at 2026-10-03 13:53:26 +0000