221 lines
6.9 KiB
Go
221 lines
6.9 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"git.gobha.me/xcaliber/chat-switchboard/models"
|
|
"git.gobha.me/xcaliber/chat-switchboard/store"
|
|
)
|
|
|
|
type UsageStore struct{}
|
|
|
|
func NewUsageStore() *UsageStore { return &UsageStore{} }
|
|
|
|
// Log inserts a usage entry.
|
|
func (s *UsageStore) Log(ctx context.Context, entry *models.UsageEntry) error {
|
|
_, err := DB.ExecContext(ctx, `
|
|
INSERT INTO usage_log (
|
|
channel_id, user_id, provider_config_id, provider_scope,
|
|
model_id, role,
|
|
prompt_tokens, completion_tokens,
|
|
cache_creation_tokens, cache_read_tokens,
|
|
cost_input, cost_output
|
|
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
|
|
`,
|
|
entry.ChannelID, entry.UserID, entry.ProviderConfigID, entry.ProviderScope,
|
|
entry.ModelID, entry.Role,
|
|
entry.PromptTokens, entry.CompletionTokens,
|
|
entry.CacheCreationTokens, entry.CacheReadTokens,
|
|
entry.CostInput, entry.CostOutput,
|
|
)
|
|
return err
|
|
}
|
|
|
|
// QueryByUser returns aggregated usage for a specific user.
|
|
func (s *UsageStore) QueryByUser(ctx context.Context, userID string, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
|
|
where := []string{"user_id = $1"}
|
|
args := []interface{}{userID}
|
|
return s.query(ctx, where, args, opts)
|
|
}
|
|
|
|
// QueryByTeam returns aggregated usage for all members of a team (admin view).
|
|
func (s *UsageStore) QueryByTeam(ctx context.Context, teamID string, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
|
|
where := []string{"user_id IN (SELECT user_id FROM team_members WHERE team_id = $1)"}
|
|
args := []interface{}{teamID}
|
|
return s.query(ctx, where, args, opts)
|
|
}
|
|
|
|
// QueryByTeamProviders returns usage against providers owned by this team.
|
|
// Used by team admins to see how their team's API keys are being consumed.
|
|
func (s *UsageStore) QueryByTeamProviders(ctx context.Context, teamID string, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
|
|
where := []string{"provider_config_id IN (SELECT id FROM provider_configs WHERE scope = 'team' AND owner_id = $1)"}
|
|
args := []interface{}{teamID}
|
|
return s.query(ctx, where, args, opts)
|
|
}
|
|
|
|
// QueryByModel returns aggregated usage grouped by model.
|
|
func (s *UsageStore) QueryByModel(ctx context.Context, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
|
|
return s.query(ctx, nil, nil, opts)
|
|
}
|
|
|
|
// GetTotals returns aggregate totals (admin view — excludes BYOK by default).
|
|
func (s *UsageStore) GetTotals(ctx context.Context, opts store.UsageQueryOptions) (*models.UsageTotals, error) {
|
|
where, args := s.buildFilters(nil, nil, opts)
|
|
|
|
q := fmt.Sprintf(`
|
|
SELECT
|
|
COUNT(*),
|
|
COALESCE(SUM(prompt_tokens), 0),
|
|
COALESCE(SUM(completion_tokens), 0),
|
|
COALESCE(SUM(COALESCE(cost_input, 0) + COALESCE(cost_output, 0)), 0)
|
|
FROM usage_log
|
|
WHERE %s
|
|
`, strings.Join(where, " AND "))
|
|
|
|
var t models.UsageTotals
|
|
err := DB.QueryRowContext(ctx, q, args...).Scan(
|
|
&t.Requests, &t.InputTokens, &t.OutputTokens, &t.TotalCost,
|
|
)
|
|
return &t, err
|
|
}
|
|
|
|
// GetTeamProviderTotals returns aggregate totals for providers owned by a team.
|
|
func (s *UsageStore) GetTeamProviderTotals(ctx context.Context, teamID string, opts store.UsageQueryOptions) (*models.UsageTotals, error) {
|
|
baseWhere := []string{"provider_config_id IN (SELECT id FROM provider_configs WHERE scope = 'team' AND owner_id = $1)"}
|
|
baseArgs := []interface{}{teamID}
|
|
where, args := s.buildFilters(baseWhere, baseArgs, opts)
|
|
|
|
q := fmt.Sprintf(`
|
|
SELECT
|
|
COUNT(*),
|
|
COALESCE(SUM(prompt_tokens), 0),
|
|
COALESCE(SUM(completion_tokens), 0),
|
|
COALESCE(SUM(COALESCE(cost_input, 0) + COALESCE(cost_output, 0)), 0)
|
|
FROM usage_log
|
|
WHERE %s
|
|
`, strings.Join(where, " AND "))
|
|
|
|
var t models.UsageTotals
|
|
err := DB.QueryRowContext(ctx, q, args...).Scan(
|
|
&t.Requests, &t.InputTokens, &t.OutputTokens, &t.TotalCost,
|
|
)
|
|
return &t, err
|
|
}
|
|
|
|
// GetPersonalTotals returns aggregate totals for a specific user (all scopes).
|
|
func (s *UsageStore) GetPersonalTotals(ctx context.Context, userID string, opts store.UsageQueryOptions) (*models.UsageTotals, error) {
|
|
baseWhere := []string{"user_id = $1"}
|
|
baseArgs := []interface{}{userID}
|
|
// Don't exclude BYOK for personal view
|
|
opts.ExcludeBYOK = false
|
|
where, args := s.buildFilters(baseWhere, baseArgs, opts)
|
|
|
|
q := fmt.Sprintf(`
|
|
SELECT
|
|
COUNT(*),
|
|
COALESCE(SUM(prompt_tokens), 0),
|
|
COALESCE(SUM(completion_tokens), 0),
|
|
COALESCE(SUM(COALESCE(cost_input, 0) + COALESCE(cost_output, 0)), 0)
|
|
FROM usage_log
|
|
WHERE %s
|
|
`, strings.Join(where, " AND "))
|
|
|
|
var t models.UsageTotals
|
|
err := DB.QueryRowContext(ctx, q, args...).Scan(
|
|
&t.Requests, &t.InputTokens, &t.OutputTokens, &t.TotalCost,
|
|
)
|
|
return &t, err
|
|
}
|
|
|
|
// ── Internal ───────────────────────────────
|
|
|
|
func (s *UsageStore) buildFilters(baseWhere []string, baseArgs []interface{}, opts store.UsageQueryOptions) ([]string, []interface{}) {
|
|
where := append([]string{}, baseWhere...)
|
|
args := append([]interface{}{}, baseArgs...)
|
|
idx := len(args) + 1
|
|
|
|
if opts.ExcludeBYOK {
|
|
where = append(where, fmt.Sprintf("provider_scope != $%d", idx))
|
|
args = append(args, "personal")
|
|
idx++
|
|
}
|
|
if opts.Since != nil {
|
|
where = append(where, fmt.Sprintf("created_at >= $%d", idx))
|
|
args = append(args, *opts.Since)
|
|
idx++
|
|
}
|
|
if opts.Until != nil {
|
|
where = append(where, fmt.Sprintf("created_at < $%d", idx))
|
|
args = append(args, *opts.Until)
|
|
idx++
|
|
}
|
|
|
|
if len(where) == 0 {
|
|
where = append(where, "1=1")
|
|
}
|
|
return where, args
|
|
}
|
|
|
|
func (s *UsageStore) query(ctx context.Context, baseWhere []string, baseArgs []interface{}, opts store.UsageQueryOptions) ([]models.UsageAggregate, error) {
|
|
where, args := s.buildFilters(baseWhere, baseArgs, opts)
|
|
|
|
// Determine GROUP BY expression and label
|
|
groupExpr := "model_id"
|
|
labelExpr := "model_id"
|
|
switch opts.GroupBy {
|
|
case "day":
|
|
groupExpr = "DATE(created_at)"
|
|
labelExpr = "DATE(created_at)::text"
|
|
case "user":
|
|
groupExpr = "user_id"
|
|
labelExpr = "user_id::text"
|
|
case "provider":
|
|
groupExpr = "provider_config_id"
|
|
labelExpr = "provider_config_id::text"
|
|
case "model":
|
|
groupExpr = "model_id"
|
|
labelExpr = "model_id"
|
|
default:
|
|
groupExpr = "model_id"
|
|
labelExpr = "model_id"
|
|
}
|
|
|
|
limit := opts.Limit
|
|
if limit <= 0 {
|
|
limit = 100
|
|
}
|
|
|
|
q := fmt.Sprintf(`
|
|
SELECT
|
|
%s AS group_key,
|
|
%s AS label,
|
|
COUNT(*) AS requests,
|
|
COALESCE(SUM(prompt_tokens), 0) AS input_tokens,
|
|
COALESCE(SUM(completion_tokens), 0) AS output_tokens,
|
|
COALESCE(SUM(COALESCE(cost_input, 0) + COALESCE(cost_output, 0)), 0) AS total_cost
|
|
FROM usage_log
|
|
WHERE %s
|
|
GROUP BY %s
|
|
ORDER BY total_cost DESC
|
|
LIMIT %d
|
|
`, groupExpr, labelExpr, strings.Join(where, " AND "), groupExpr, limit)
|
|
|
|
rows, err := DB.QueryContext(ctx, q, args...)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var results []models.UsageAggregate
|
|
for rows.Next() {
|
|
var a models.UsageAggregate
|
|
if err := rows.Scan(&a.GroupKey, &a.Label, &a.Requests, &a.InputTokens, &a.OutputTokens, &a.TotalCost); err != nil {
|
|
return nil, err
|
|
}
|
|
results = append(results, a)
|
|
}
|
|
return results, rows.Err()
|
|
}
|