Changeset 0.10.0 (#56)
This commit is contained in:
@@ -30,6 +30,8 @@ type Stores struct {
|
||||
Audit AuditStore
|
||||
Notes NoteStore
|
||||
GlobalConfig GlobalConfigStore
|
||||
Usage UsageStore
|
||||
Pricing PricingStore
|
||||
}
|
||||
|
||||
// =========================================
|
||||
@@ -276,6 +278,41 @@ type GlobalConfigStore interface {
|
||||
GetAll(ctx context.Context) (map[string]models.JSONMap, error)
|
||||
}
|
||||
|
||||
// =========================================
|
||||
// USAGE STORE
|
||||
// =========================================
|
||||
|
||||
type UsageStore interface {
|
||||
Log(ctx context.Context, entry *models.UsageEntry) error
|
||||
QueryByUser(ctx context.Context, userID string, opts UsageQueryOptions) ([]models.UsageAggregate, error)
|
||||
QueryByTeam(ctx context.Context, teamID string, opts UsageQueryOptions) ([]models.UsageAggregate, error)
|
||||
QueryByTeamProviders(ctx context.Context, teamID string, opts UsageQueryOptions) ([]models.UsageAggregate, error)
|
||||
QueryByModel(ctx context.Context, opts UsageQueryOptions) ([]models.UsageAggregate, error)
|
||||
GetTotals(ctx context.Context, opts UsageQueryOptions) (*models.UsageTotals, error)
|
||||
GetTeamProviderTotals(ctx context.Context, teamID string, opts UsageQueryOptions) (*models.UsageTotals, error)
|
||||
GetPersonalTotals(ctx context.Context, userID string, opts UsageQueryOptions) (*models.UsageTotals, error)
|
||||
}
|
||||
|
||||
type UsageQueryOptions struct {
|
||||
Since *time.Time
|
||||
Until *time.Time
|
||||
GroupBy string // "day", "model", "user", "provider"
|
||||
ExcludeBYOK bool // true for admin views — filters provider_scope != 'personal'
|
||||
Limit int
|
||||
}
|
||||
|
||||
// =========================================
|
||||
// PRICING STORE
|
||||
// =========================================
|
||||
|
||||
type PricingStore interface {
|
||||
GetForModel(ctx context.Context, providerConfigID, modelID string) (*models.PricingEntry, error)
|
||||
Upsert(ctx context.Context, entry *models.PricingEntry) error
|
||||
UpsertFromCatalog(ctx context.Context, providerConfigID, modelID string, pricing *models.ModelPricing) error
|
||||
List(ctx context.Context) ([]models.PricingEntry, error)
|
||||
Delete(ctx context.Context, providerConfigID, modelID string) error
|
||||
}
|
||||
|
||||
// =========================================
|
||||
// SHARED TYPES
|
||||
// =========================================
|
||||
|
||||
131
server/store/postgres/pricing.go
Normal file
131
server/store/postgres/pricing.go
Normal file
@@ -0,0 +1,131 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
|
||||
"git.gobha.me/xcaliber/chat-switchboard/models"
|
||||
)
|
||||
|
||||
type PricingStore struct{}
|
||||
|
||||
func NewPricingStore() *PricingStore { return &PricingStore{} }
|
||||
|
||||
// GetForModel returns the pricing entry for a specific provider+model pair.
|
||||
func (s *PricingStore) GetForModel(ctx context.Context, providerConfigID, modelID string) (*models.PricingEntry, error) {
|
||||
var p models.PricingEntry
|
||||
err := DB.QueryRowContext(ctx, `
|
||||
SELECT id, provider_config_id, model_id,
|
||||
input_per_m, output_per_m, cache_create_per_m, cache_read_per_m,
|
||||
currency, source, updated_at, updated_by
|
||||
FROM model_pricing
|
||||
WHERE provider_config_id = $1 AND model_id = $2
|
||||
`, providerConfigID, modelID).Scan(
|
||||
&p.ID, &p.ProviderConfigID, &p.ModelID,
|
||||
&p.InputPerM, &p.OutputPerM, &p.CacheCreatePerM, &p.CacheReadPerM,
|
||||
&p.Currency, &p.Source, &p.UpdatedAt, &p.UpdatedBy,
|
||||
)
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return &p, err
|
||||
}
|
||||
|
||||
// Upsert inserts or updates a pricing entry (manual admin override).
|
||||
func (s *PricingStore) Upsert(ctx context.Context, entry *models.PricingEntry) error {
|
||||
_, err := DB.ExecContext(ctx, `
|
||||
INSERT INTO model_pricing (
|
||||
provider_config_id, model_id,
|
||||
input_per_m, output_per_m, cache_create_per_m, cache_read_per_m,
|
||||
currency, source, updated_by, updated_at
|
||||
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, NOW())
|
||||
ON CONFLICT (provider_config_id, model_id)
|
||||
DO UPDATE SET
|
||||
input_per_m = EXCLUDED.input_per_m,
|
||||
output_per_m = EXCLUDED.output_per_m,
|
||||
cache_create_per_m = EXCLUDED.cache_create_per_m,
|
||||
cache_read_per_m = EXCLUDED.cache_read_per_m,
|
||||
currency = EXCLUDED.currency,
|
||||
source = EXCLUDED.source,
|
||||
updated_by = EXCLUDED.updated_by,
|
||||
updated_at = NOW()
|
||||
`,
|
||||
entry.ProviderConfigID, entry.ModelID,
|
||||
entry.InputPerM, entry.OutputPerM, entry.CacheCreatePerM, entry.CacheReadPerM,
|
||||
entry.Currency, entry.Source, entry.UpdatedBy,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
// UpsertFromCatalog inserts or updates pricing from a catalog sync.
|
||||
// Manual overrides (source='manual') are never overwritten.
|
||||
func (s *PricingStore) UpsertFromCatalog(ctx context.Context, providerConfigID, modelID string, pricing *models.ModelPricing) error {
|
||||
if pricing == nil || (pricing.InputPerM == 0 && pricing.OutputPerM == 0) {
|
||||
return nil // No pricing to store
|
||||
}
|
||||
|
||||
currency := pricing.Currency
|
||||
if currency == "" {
|
||||
currency = "USD"
|
||||
}
|
||||
|
||||
_, err := DB.ExecContext(ctx, `
|
||||
INSERT INTO model_pricing (
|
||||
provider_config_id, model_id,
|
||||
input_per_m, output_per_m,
|
||||
currency, source, updated_at
|
||||
) VALUES ($1, $2, $3, $4, $5, 'catalog', NOW())
|
||||
ON CONFLICT (provider_config_id, model_id)
|
||||
DO UPDATE SET
|
||||
input_per_m = EXCLUDED.input_per_m,
|
||||
output_per_m = EXCLUDED.output_per_m,
|
||||
currency = EXCLUDED.currency,
|
||||
updated_at = NOW()
|
||||
WHERE model_pricing.source = 'catalog'
|
||||
`,
|
||||
providerConfigID, modelID,
|
||||
pricing.InputPerM, pricing.OutputPerM,
|
||||
currency,
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
// List returns pricing entries for admin-managed providers (global + team).
|
||||
// Excludes personal BYOK providers — those are not the admin's concern.
|
||||
func (s *PricingStore) List(ctx context.Context) ([]models.PricingEntry, error) {
|
||||
rows, err := DB.QueryContext(ctx, `
|
||||
SELECT mp.id, mp.provider_config_id, mp.model_id,
|
||||
mp.input_per_m, mp.output_per_m, mp.cache_create_per_m, mp.cache_read_per_m,
|
||||
mp.currency, mp.source, mp.updated_at, mp.updated_by
|
||||
FROM model_pricing mp
|
||||
JOIN provider_configs pc ON pc.id = mp.provider_config_id
|
||||
WHERE pc.scope != 'personal'
|
||||
ORDER BY mp.provider_config_id, mp.model_id
|
||||
`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var results []models.PricingEntry
|
||||
for rows.Next() {
|
||||
var p models.PricingEntry
|
||||
if err := rows.Scan(
|
||||
&p.ID, &p.ProviderConfigID, &p.ModelID,
|
||||
&p.InputPerM, &p.OutputPerM, &p.CacheCreatePerM, &p.CacheReadPerM,
|
||||
&p.Currency, &p.Source, &p.UpdatedAt, &p.UpdatedBy,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
results = append(results, p)
|
||||
}
|
||||
return results, rows.Err()
|
||||
}
|
||||
|
||||
// Delete removes a pricing entry.
|
||||
func (s *PricingStore) Delete(ctx context.Context, providerConfigID, modelID string) error {
|
||||
_, err := DB.ExecContext(ctx, `
|
||||
DELETE FROM model_pricing WHERE provider_config_id = $1 AND model_id = $2
|
||||
`, providerConfigID, modelID)
|
||||
return err
|
||||
}
|
||||
@@ -23,5 +23,7 @@ func NewStores(db *sql.DB) store.Stores {
|
||||
Audit: NewAuditStore(),
|
||||
Notes: NewNoteStore(),
|
||||
GlobalConfig: NewGlobalConfigStore(),
|
||||
Usage: NewUsageStore(),
|
||||
Pricing: NewPricingStore(),
|
||||
}
|
||||
}
|
||||
|
||||
220
server/store/postgres/usage.go
Normal file
220
server/store/postgres/usage.go
Normal file
@@ -0,0 +1,220 @@
|
||||
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()
|
||||
}
|
||||
Reference in New Issue
Block a user