diff options
Diffstat (limited to 'internal/controlplane/usage.go')
| -rw-r--r-- | internal/controlplane/usage.go | 128 |
1 files changed, 121 insertions, 7 deletions
diff --git a/internal/controlplane/usage.go b/internal/controlplane/usage.go index b436d8c..ebabaf6 100644 --- a/internal/controlplane/usage.go +++ b/internal/controlplane/usage.go @@ -14,7 +14,16 @@ import ( type UsageQuery struct { TenantID string ProjectID string + KeyID string Model string + Provider string + Protocol string + ErrorType string + Stream *bool + RequestID string + Status string + From time.Time + To time.Time Limit int } @@ -39,6 +48,11 @@ func (s *Store) RecordUsage(ctx context.Context, event domain.UsageEvent) error if err != nil { return fmt.Errorf("persist usage event: %w", err) } + if _, err := tx.Exec(ctx, `UPDATE api_keys + SET last_used_at = GREATEST(COALESCE(last_used_at, $2), $2) + WHERE id = $1`, event.KeyID, event.StartedAt); err != nil { + return fmt.Errorf("update API key last used time: %w", err) + } if command.RowsAffected() == 0 { return tx.Commit(ctx) } @@ -96,20 +110,55 @@ func (s *Store) ListUsage(ctx context.Context, query UsageQuery) ([]UsageRecord, limit = 200 } where := []string{"1=1"} - args := make([]any, 0, 5) + args := make([]any, 0, 13) index := 1 - for _, item := range []struct{ value, clause string }{{query.TenantID, "tenant_id=$"}, {query.ProjectID, "project_id=$"}, {query.Model, "public_model=$"}} { + 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) != "" { where = append(where, item.clause+fmt.Sprint(index)) args = append(args, item.value) index++ } } + for _, item := range []struct{ value, clause string }{{query.Protocol, "protocol=$"}, {query.ErrorType, "error_type=$"}} { + if strings.TrimSpace(item.value) != "" { + where = append(where, item.clause+fmt.Sprint(index)) + args = append(args, item.value) + index++ + } + } + if strings.TrimSpace(query.Provider) != "" { + where = append(where, "EXISTS (SELECT 1 FROM providers filter_provider WHERE filter_provider.id::text=usage_events.provider_id AND filter_provider.slug=$"+fmt.Sprint(index)+")") + args = append(args, query.Provider) + index++ + } + if query.Stream != nil { + where = append(where, "stream=$"+fmt.Sprint(index)) + args = append(args, *query.Stream) + index++ + } + if query.Status == "success" { + where = append(where, "success=TRUE") + } else if query.Status == "error" { + where = append(where, "success=FALSE") + } + if !query.From.IsZero() { + where = append(where, "started_at >= $"+fmt.Sprint(index)) + args = append(args, query.From) + index++ + } + if !query.To.IsZero() { + where = append(where, "started_at < $"+fmt.Sprint(index)) + args = append(args, query.To) + index++ + } args = append(args, limit) - rows, err := s.db.Query(ctx, `SELECT request_id, tenant_id::text, project_id::text, key_id::text, public_model, - COALESCE(provider_id,''), COALESCE(upstream_model,''), protocol, stream, status_code, success, error_type, + 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, - cache_read_input_tokens, cost_micros, charged_micros, uncollected_micros FROM usage_events WHERE `+ + 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...) if err != nil { return nil, fmt.Errorf("query usage events: %w", err) @@ -118,10 +167,10 @@ func (s *Store) ListUsage(ctx context.Context, query UsageQuery) ([]UsageRecord, result := make([]UsageRecord, 0) for rows.Next() { var item UsageRecord - if err := rows.Scan(&item.RequestID, &item.TenantID, &item.ProjectID, &item.KeyID, &item.PublicModel, &item.ProviderID, &item.UpstreamModel, + 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.CostMicros, &item.ChargedMicros, &item.UncollectedMicros); err != nil { + &item.CostMicros, &item.ChargedMicros, &item.UncollectedMicros, &item.UsageReported, &item.MeteringStatus); err != nil { return nil, fmt.Errorf("scan usage event: %w", err) } result = append(result, item) @@ -129,6 +178,71 @@ func (s *Store) ListUsage(ctx context.Context, query UsageQuery) ([]UsageRecord, return result, rows.Err() } +func (s *Store) UsageDaily(ctx context.Context, query UsageQuery) ([]UsageDailyPoint, error) { + where := []string{"1=1"} + args := make([]any, 0, 12) + 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) != "" { + where = append(where, item.clause+fmt.Sprint(index)) + args = append(args, item.value) + index++ + } + } + for _, item := range []struct{ value, clause string }{{query.Protocol, "protocol=$"}, {query.ErrorType, "error_type=$"}} { + if strings.TrimSpace(item.value) != "" { + where = append(where, item.clause+fmt.Sprint(index)) + args = append(args, item.value) + index++ + } + } + if strings.TrimSpace(query.Provider) != "" { + where = append(where, "EXISTS (SELECT 1 FROM providers filter_provider WHERE filter_provider.id::text=usage_events.provider_id AND filter_provider.slug=$"+fmt.Sprint(index)+")") + args = append(args, query.Provider) + index++ + } + if query.Stream != nil { + where = append(where, "stream=$"+fmt.Sprint(index)) + args = append(args, *query.Stream) + index++ + } + if query.Status == "success" { + where = append(where, "success=TRUE") + } else if query.Status == "error" { + where = append(where, "success=FALSE") + } + if !query.From.IsZero() { + where = append(where, "started_at >= $"+fmt.Sprint(index)) + args = append(args, query.From) + index++ + } + if !query.To.IsZero() { + where = append(where, "started_at < $"+fmt.Sprint(index)) + args = append(args, query.To) + index++ + } + rows, err := s.db.Query(ctx, `SELECT date_trunc('day', started_at AT TIME ZONE 'UTC'), count(*), + 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 + 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) + } + defer rows.Close() + result := make([]UsageDailyPoint, 0) + 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 { + return nil, fmt.Errorf("scan daily usage: %w", err) + } + result = append(result, item) + } + return result, rows.Err() +} + func (s *Store) UsageSummary(ctx context.Context, tenantID, projectID string) ([]UsageSummary, error) { where := []string{"1=1"} args := make([]any, 0, 2) |