diff --git a/STATISTIC.md b/STATISTIC.md new file mode 100644 index 0000000..435f11f --- /dev/null +++ b/STATISTIC.md @@ -0,0 +1,253 @@ +statistics API response notes: +- all statistics below use only data not older than the last 5 months +- `costs.1month` and `costs.5month` remain cumulative spend windows +- distributions are now time-series points, not histograms +- each point is calculated from unique user spend aggregated inside a calendar period +- calendar boundaries are used: + - day: from start to end of day + - week: from start to end of ISO-style week (Monday to Sunday) + - month: from start to end of month +- returned windows: + - `daily*Distribution`: one point per day for the last month + - `weekly*Distribution`: one point per week for the last 5 months + - `monthly*Distribution`: one point per month for the last 5 months + +cost: +{ + "1month": float, + "5month": float +} + +costPack: +{ + "codioProvided": cost, + "codioSpecial": cost +} + +dailyDistributionDataPoint: +{ + "maxUserSpend": float, + "medianUserSpend": float, + "avgUserSpend": float, + "p95UserSpend": float, + "p99UserSpend": float, + "date": "YYYY-MM-DD" +} + +periodDistributionDataPoint: +{ + "maxUserSpend": float, + "medianUserSpend": float, + "avgUserSpend": float, + "p95UserSpend": float, + "p99UserSpend": float, + "datePeriod": string +} + +----------------------- + +all: +{ + "total": costPack, + "orgs": [ + { + "id": string, + "costs": costPack + } + ] +} + +--------------------- + +topFiveItem: +{ + "userId": string, + "spendLastMonth": float +} + +org: +{ + "id": string, + "costs": costPack, + "courses": [ + { + "id": string, + "costs": costPack + } + ], + "dailySpecialDistribution": [dailyDistributionDataPoint], + "weeklySpecialDistribution": [periodDistributionDataPoint], + "monthlySpecialDistribution": [periodDistributionDataPoint], + "dailyCodioProvidedDistribution": [dailyDistributionDataPoint], + "weeklyCodioProvidedDistribution": [periodDistributionDataPoint], + "monthlyCodioProvidedDistribution": [periodDistributionDataPoint], + "topFive": [topFiveItem] +} + +--------------------- + +course: +{ + "id": string, + "costs": costPack, + "dailySpecialDistribution": [dailyDistributionDataPoint], + "weeklySpecialDistribution": [periodDistributionDataPoint], + "monthlySpecialDistribution": [periodDistributionDataPoint], + "dailyCodioProvidedDistribution": [dailyDistributionDataPoint], + "weeklyCodioProvidedDistribution": [periodDistributionDataPoint], + "monthlyCodioProvidedDistribution": [periodDistributionDataPoint], + "topFive": [topFiveItem] +} + +----------------------------- + +How to interpret the returned statistics + +General principles: +- each point is calculated from unique users active in that calendar period +- for one period, each user contributes one aggregated spend value +- metrics are then calculated across the set of user totals for that period +- if a period has no activity, the API still returns a point with zero values so the frontend can render a continuous series + +What `costs` means: +- `costs.codioProvided.1month` — cumulative spend for codio-provided traffic over the last 1 month +- `costs.codioProvided.5month` — cumulative spend for codio-provided traffic over the last 5 months +- `costs.codioSpecial.1month` — cumulative spend for codio-special traffic over the last 1 month +- `costs.codioSpecial.5month` — cumulative spend for codio-special traffic over the last 5 months +- this block is useful for total budget visibility, while the distribution arrays are useful for trend analysis + +What each distribution point means: +- `maxUserSpend` — the highest per-user spend in that day/week/month +- `medianUserSpend` — 50% of active users in that period spent at or below this value +- `avgUserSpend` — arithmetic average per-user spend in that period +- `p95UserSpend` — 95% of active users in that period spent at or below this value +- `p99UserSpend` — 99% of active users in that period spent at or below this value +- `date` is used for daily points +- `datePeriod` is used for weekly/monthly points +- `topFive` is a companion list for the current scope and contains the top 5 users by spend for the last month + +Date labels: +- daily uses `YYYY-MM-DD` +- weekly uses `YYYY-MM-DD/YYYY-MM-DD` where the label represents week start and week end +- monthly uses `YYYY-MM` + +Examples: +- `dailySpecialDistribution[i]` shows one calendar day from the last month +- `weeklySpecialDistribution[i]` shows one calendar week from the last 5 months +- `monthlyCodioProvidedDistribution[i]` shows one calendar month from the last 5 months +- `topFive[i]` shows one of the five users with the highest spend for the last month in the current org or course scope + +How to read `topFive` + +- `topFive` is calculated for the last month +- it is scoped to the current entity: + - for org statistics, top users inside that organization + - for course statistics, top users inside that course +- it is not split by `codio-special` or `codio-provided`; it reflects total spend in scope for the last month +- it is useful for quickly identifying the most expensive users for investigation, outreach, or manual policy review + +How to read the charts + +Daily chart: +- use it to see short-term volatility and spikes +- `maxUserSpend` highlights strongest single-user bursts in a day +- `medianUserSpend` and `avgUserSpend` show whether broad usage is rising or only a few users are spiking + +Weekly chart: +- use it to see medium-term usage stabilization +- compare weekly `p95UserSpend` and `p99UserSpend` across weeks +- this is useful for tuning weekly guardrails + +Monthly chart: +- use it for budget policy and allowance planning +- monthly changes are less noisy and better reflect stable behavior +- a rising monthly `medianUserSpend` means the typical user is genuinely spending more + +How to use these metrics for limits + +Daily limits: +- look at `daily...Distribution` +- a high `p99UserSpend` with a low median usually means a few strong outliers +- useful for anti-spike or abuse protection + +Weekly limits: +- look at `weekly...Distribution` +- stable weekly `p95UserSpend` can guide soft-limit candidates +- rising weekly `maxUserSpend` may justify a hard cap or alerts + +Monthly limits: +- look at `monthly...Distribution` +- this is the best signal for recurring allowance defaults +- if monthly `medianUserSpend` stays low but `p99UserSpend` rises, the tail is getting heavier without broad adoption + +Practical recommendation flow + +1. Start with monthly series: +- review the trend of `medianUserSpend`, `p95UserSpend`, and `p99UserSpend` +- use this to set or revise monthly allowances + +2. Check weekly series: +- see whether usage changes smoothly week to week or has temporary bursts +- if weekly p95 stays stable but max jumps, use alerts before stricter limits + +3. Check daily series: +- use it for operational safety and anomaly detection +- daily max and p99 are especially useful for identifying runaway prompts or abuse + +Suggested interpretation patterns + +Pattern A: Low median, high max, high p99 only on a few days +- normal usage is cheap +- spikes are rare and sharp +- recommended action: + - keep daily hard caps + - avoid lowering broad monthly limits unnecessarily + +Pattern B: Median and avg both trend upward over weeks and months +- usage is growing across the user base, not just in outliers +- recommended action: + - revise default weekly/monthly budgets upward if product usage is healthy + +Pattern C: Monthly p95 rises while median stays flat +- most users are stable, but heavy-user tail is getting more expensive +- recommended action: + - keep default limits, but strengthen tail controls for power users + +Pattern D: Weekly and monthly max rise together +- expensive usage is not only a single-day anomaly +- recommended action: + - investigate long-running high-cost users and review policy settings + +Limitations of the current model + +- these arrays show metric trends over time, not spend histograms by bucket +- percentiles can be noisy when the number of active users in a period is low +- daily values are naturally more volatile than weekly/monthly values +- recommendations should still be combined with product and business context + +Suggested UI usage + +For each org or course, show: +- a `topFive` table or side panel with: + - `userId` + - `spendLastMonth` +- separate charts for: +- daily special +- weekly special +- monthly special +- daily codio-provided +- weekly codio-provided +- monthly codio-provided + +For each chart, plot one or more lines: +- `medianUserSpend` +- `p95UserSpend` +- `p99UserSpend` +- optionally `maxUserSpend` + +Suggested helper text: +- "Daily p99 shows near-worst-case per-user spend for a single day" +- "Weekly median shows the typical user spend for a full calendar week" +- "Monthly p95 is a strong candidate input for limit policy reviews" + +This makes the dashboard useful for monitoring trend shifts, tuning limits, and spotting anomalous growth. diff --git a/cmd/bricksllm/main.go b/cmd/bricksllm/main.go index 40913bc..f25cf86 100644 --- a/cmd/bricksllm/main.go +++ b/cmd/bricksllm/main.go @@ -297,6 +297,14 @@ func main() { log.Sugar().Fatalf("error connecting to secondary keys redis storage: %v", err) } + statisticsRedisCache := redis.NewClient(defaultRedisOption(cfg, 13)) + + ctx, cancel = context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if err = statisticsRedisCache.Ping(ctx).Err(); err != nil { + log.Sugar().Fatalf("error connecting to statistics redis storage: %v", err) + } + rateLimitCache := redisStorage.NewCache(rateLimitRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) costLimitCache := redisStorage.NewCache(costLimitRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) costStorage := redisStorage.NewStore(costRedisStorage, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) @@ -312,22 +320,23 @@ func main() { keysCache := redisStorage.NewKeysCache(keysRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) secondaryKeysCache := redisStorage.NewSecondaryKeysCache(secondaryKeysRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) requestsLimitStorage := redisStorage.NewStore(requestsLimitRedisStorage, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) + statisticsCache := redisStorage.NewStatisticCache(statisticsRedisCache, cfg.RedisWriteTimeout, cfg.RedisReadTimeout) - encryptor, err := encryptor.NewEncryptor(cfg.DecryptionEndpoint, cfg.EncryptionEndpoint, cfg.EnableEncrytion, cfg.EncryptionTimeout, cfg.Audience) + encrypt, err := encryptor.NewEncryptor(cfg.DecryptionEndpoint, cfg.EncryptionEndpoint, cfg.EnableEncrytion, cfg.EncryptionTimeout, cfg.Audience) if cfg.EnableEncrytion && err != nil { log.Sugar().Fatalf("error creating encryption client: %v", err) } v := validator.NewValidator(costLimitCache, rateLimitCache, costStorage, requestsLimitStorage) m := manager.NewManager(store, costLimitCache, rateLimitCache, accessCache, keysCache, secondaryKeysCache, requestsLimitStorage) - krm := manager.NewReportingManager(costStorage, store, store, v) - psm := manager.NewProviderSettingsManager(store, psCache, encryptor) + krm := manager.NewReportingManager(costStorage, store, store, v, statisticsCache) + psm := manager.NewProviderSettingsManager(store, psCache, encrypt) cpm := manager.NewCustomProvidersManager(store, cpMemStore) rm := manager.NewRouteManager(store, store, rMemStore, psm) pm := manager.NewPolicyManager(store, rMemStore) um := manager.NewUserManager(store, store) - as, err := admin.NewAdminServer(log, *modePtr, m, krm, psm, cpm, rm, pm, um, cfg.AdminPass, cfg.XCodioSignSecret) + as, err := admin.NewAdminServer(log, *modePtr, m, krm, psm, cpm, rm, pm, um) if err != nil { log.Sugar().Fatalf("error creating admin http server: %v", err) } @@ -358,7 +367,7 @@ func main() { rec := recorder.NewRecorder(costStorage, userCostStorage, costLimitCache, userCostLimitCache, ce, store, requestsLimitStorage) rlm := manager.NewRateLimitManager(rateLimitCache, userRateLimitCache) - a := auth.NewAuthenticator(psm, m, rm, store, encryptor) + a := auth.NewAuthenticator(psm, m, rm, store, encrypt) c := cache.NewCache(apiCache) diff --git a/go.mod b/go.mod index 0789b4a..ac0927b 100644 --- a/go.mod +++ b/go.mod @@ -12,6 +12,7 @@ require ( github.com/cenkalti/backoff/v4 v4.3.0 github.com/fatih/color v1.18.0 github.com/gin-gonic/gin v1.11.0 + github.com/google/tink/go v1.7.0 github.com/google/uuid v1.6.0 github.com/joho/godotenv v1.5.1 github.com/lib/pq v1.10.9 diff --git a/go.sum b/go.sum index df2531b..09df7af 100644 --- a/go.sum +++ b/go.sum @@ -115,6 +115,8 @@ github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= +github.com/google/tink/go v1.7.0 h1:6Eox8zONGebBFcCBqkVmt60LaWZa6xg1cl/DwAh/J1w= +github.com/google/tink/go v1.7.0/go.mod h1:GAUOd+QE3pgj9q8VKIGTCP33c/B7eb4NhxLcgTJZStM= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/googleapis/enterprise-certificate-proxy v0.3.6 h1:GW/XbdyBFQ8Qe+YAmFU9uHLo7OnF5tL52HFAgMmyrf4= diff --git a/internal/event/key_reporting.go b/internal/event/key_reporting.go index 6ebf3f4..7aa9c67 100644 --- a/internal/event/key_reporting.go +++ b/internal/event/key_reporting.go @@ -1,5 +1,11 @@ package event +import ( + "strings" + + internalErrors "github.com/bricks-cloud/bricksllm/internal/errors" +) + type KeyDataPoint struct { KeyId string `json:"keyId"` CostInUsd float64 `json:"costInUsd"` @@ -76,3 +82,154 @@ type UsageData struct { type UsageReportingResponse struct { UsageData *UsageData `json:"usageData"` } + +type StatisticsRequest struct { + Level string `json:"level"` // all | org | course + Id *string `json:"id"` +} + +func (r *StatisticsRequest) GetCacheKey() string { + if r.Level == "all" { + return r.Level + } + if r.Id != nil { + return r.Level + ":" + *r.Id + } + return r.Level +} + +type StatisticLevel string + +var StaticLevels = struct { + Unknown StatisticLevel + All StatisticLevel + Org StatisticLevel + Course StatisticLevel +}{ + Unknown: "unknown", + All: "all", + Org: "org", + Course: "course", +} + +func StatisticLevelFromStr(s string) StatisticLevel { + switch s { + case "all": + return StaticLevels.All + case "org": + return StaticLevels.Org + case "course": + return StaticLevels.Course + default: + return StaticLevels.Unknown + } +} + +func (r *StatisticsRequest) Validate() error { + if r.Level != "all" && r.Level != "org" && r.Level != "course" { + return internalErrors.NewValidationError("level must be one of 'all', 'org', or 'course'") + } + if (r.Level == "org" || r.Level == "course") && (r.Id == nil || strings.TrimSpace(*r.Id) == "") { + return internalErrors.NewValidationError("id must be provided when level is 'org' or 'course'") + } + return nil +} + +func (r *StatisticsRequest) GetLevel() StatisticLevel { + return StatisticLevelFromStr(r.Level) +} + +type StatisticsResponse struct { + StatisticsData *StatisticsData `json:"statisticsData"` +} + +type StatisticsData struct { + AllStatisticsData *AllStatisticsData `json:"allStatisticsData,omitempty"` + OrgStatisticsData *OrgStatisticsData `json:"orgStatisticsData,omitempty"` + CourseStatisticsData *CourseStatisticsData `json:"courseStatisticsData,omitempty"` +} + +type Cost struct { + OneMonth float64 `json:"1month"` + FiveMonth float64 `json:"5month"` +} + +type CostPack struct { + CodioProvided Cost `json:"codioProvided"` + CodioSpecial Cost `json:"codioSpecial"` +} + +type FinancialKpis struct { + MaxUserSpend float64 `json:"maxUserSpend"` + MedianUserSpend float64 `json:"medianUserSpend"` + AvgUserSpend float64 `json:"avgUserSpend"` + P90UserSpend float64 `json:"p90UserSpend"` + P95UserSpend float64 `json:"p95UserSpend"` + P99UserSpend float64 `json:"p99UserSpend"` + SampleCount int `json:"sampleCount"` + RecommendedSoftLimitUsd float64 `json:"recommendedSoftLimitUsd"` + RecommendedHardLimitUsd float64 `json:"recommendedHardLimitUsd"` +} + +type DailySpendDistributionDataPoint struct { + MaxUserSpend float64 `json:"maxUserSpend"` + MedianUserSpend float64 `json:"medianUserSpend"` + AvgUserSpend float64 `json:"avgUserSpend"` + P95UserSpend float64 `json:"p95UserSpend"` + P99UserSpend float64 `json:"p99UserSpend"` + Date string `json:"date"` +} + +type PeriodSpendDistributionDataPoint struct { + MaxUserSpend float64 `json:"maxUserSpend"` + MedianUserSpend float64 `json:"medianUserSpend"` + AvgUserSpend float64 `json:"avgUserSpend"` + P95UserSpend float64 `json:"p95UserSpend"` + P99UserSpend float64 `json:"p99UserSpend"` + DatePeriod string `json:"datePeriod"` +} + +type AllStatisticsData struct { + Total CostPack `json:"total"` + Orgs []ShortOrgStatisticsData `json:"orgs"` +} + +type ShortOrgStatisticsData struct { + Id string `json:"id"` + Costs CostPack `json:"costs"` +} + +type ShortCourseStatisticsData struct { + Id string `json:"id"` + Costs CostPack `json:"costs"` +} + +type TopFiveUserSpend struct { + UserId string `json:"userId"` + SpendLastMonth float64 `json:"spendLastMonth"` +} + +type OrgStatisticsData struct { + Id string `json:"id"` + Costs CostPack `json:"costs"` + Courses []ShortCourseStatisticsData `json:"courses"` + DailySpecialDistribution []DailySpendDistributionDataPoint `json:"dailySpecialDistribution"` + WeeklySpecialDistribution []PeriodSpendDistributionDataPoint `json:"weeklySpecialDistribution"` + MonthlySpecialDistribution []PeriodSpendDistributionDataPoint `json:"monthlySpecialDistribution"` + DailyCodioProvidedDistribution []DailySpendDistributionDataPoint `json:"dailyCodioProvidedDistribution"` + WeeklyCodioProvidedDistribution []PeriodSpendDistributionDataPoint `json:"weeklyCodioProvidedDistribution"` + MonthlyCodioProvidedDistribution []PeriodSpendDistributionDataPoint `json:"monthlyCodioProvidedDistribution"` + TopFive []TopFiveUserSpend `json:"topFive"` +} + +type CourseStatisticsData struct { + Id string `json:"id"` + Costs CostPack `json:"costs"` + DailySpecialDistribution []DailySpendDistributionDataPoint `json:"dailySpecialDistribution"` + WeeklySpecialDistribution []PeriodSpendDistributionDataPoint `json:"weeklySpecialDistribution"` + MonthlySpecialDistribution []PeriodSpendDistributionDataPoint `json:"monthlySpecialDistribution"` + DailyCodioProvidedDistribution []DailySpendDistributionDataPoint `json:"dailyCodioProvidedDistribution"` + WeeklyCodioProvidedDistribution []PeriodSpendDistributionDataPoint `json:"weeklyCodioProvidedDistribution"` + MonthlyCodioProvidedDistribution []PeriodSpendDistributionDataPoint `json:"monthlyCodioProvidedDistribution"` + TopFive []TopFiveUserSpend `json:"topFive"` +} diff --git a/internal/manager/reporting.go b/internal/manager/reporting.go index ffe2578..37c9679 100644 --- a/internal/manager/reporting.go +++ b/internal/manager/reporting.go @@ -2,6 +2,7 @@ package manager import ( "strings" + "time" internal_errors "github.com/bricks-cloud/bricksllm/internal/errors" "github.com/bricks-cloud/bricksllm/internal/event" @@ -21,6 +22,13 @@ type keyValidator interface { Validate(k *key.ResponseKey, promptCost float64) error } +type StatisticsCache interface { + Get(key string) (*event.StatisticsData, error) + Set(key string, val *event.StatisticsData, ttl time.Duration) error + TryMarkInProgress(key string) (bool, error) + DeleteInProgress(key string) error +} + type eventStorage interface { GetEvents(userId, customId string, keyIds []string, start, end int64) ([]*event.Event, error) GetEventsV2(req *event.EventRequest) (*event.EventResponse, error) @@ -33,6 +41,7 @@ type eventStorage interface { GetTopKeyRingDataPoints(start, end int64, tags []string, order string, limit, offset int, revoked *bool, topBy string) ([]*event.KeyRingDataPoint, error) GetUsageData(tags []string) (*event.UsageData, error) + GetStatisticsData(level event.StatisticLevel, id *string) (*event.StatisticsData, error) } type ReportingManager struct { @@ -40,14 +49,16 @@ type ReportingManager struct { cs costStorage ks keyStorage kv keyValidator + sc StatisticsCache } -func NewReportingManager(cs costStorage, ks keyStorage, es eventStorage, kv keyValidator) *ReportingManager { +func NewReportingManager(cs costStorage, ks keyStorage, es eventStorage, kv keyValidator, sc StatisticsCache) *ReportingManager { return &ReportingManager{ cs: cs, ks: ks, es: es, kv: kv, + sc: sc, } } @@ -187,6 +198,57 @@ func (rm *ReportingManager) GetUsageReporting(r *event.UsageReportingRequest) (* }, nil } +func (rm *ReportingManager) GetStatistic(r *event.StatisticsRequest) (*event.StatisticsResponse, error) { + if r == nil { + return nil, internal_errors.NewValidationError("statistics request cannot be nil") + } + + if err := r.Validate(); err != nil { + return nil, err + } + + cacheKey := r.GetCacheKey() + if data, err := rm.sc.Get(cacheKey); err == nil { + return &event.StatisticsResponse{ + StatisticsData: data, + }, nil + } + + go rm.backgroundCollectStatisticsData(cacheKey, r.GetLevel(), r.Id) + + timeout := time.After(10 * time.Second) + ticker := time.NewTicker(1 * time.Second) + defer ticker.Stop() + + for { + select { + case <-timeout: + return nil, internal_errors.NewNotFoundError("statistics data is not ready yet, please try again later") + case <-ticker.C: + if data, err := rm.sc.Get(cacheKey); err == nil { + return &event.StatisticsResponse{ + StatisticsData: data, + }, nil + } + } + } +} + +func (rm *ReportingManager) backgroundCollectStatisticsData(cacheKey string, level event.StatisticLevel, id *string) { + claimed, err := rm.sc.TryMarkInProgress(cacheKey) + if err != nil || !claimed { + return + } + defer rm.sc.DeleteInProgress(cacheKey) + + statisticsData, err := rm.es.GetStatisticsData(level, id) + + if err != nil { + return + } + _ = rm.sc.Set(cacheKey, statisticsData, time.Hour*24) +} + func (rm *ReportingManager) GetCustomIds(keyId string) ([]string, error) { return rm.es.GetCustomIds(keyId) } diff --git a/internal/server/web/admin/admin.go b/internal/server/web/admin/admin.go index e83e846..fb089ff 100644 --- a/internal/server/web/admin/admin.go +++ b/internal/server/web/admin/admin.go @@ -43,6 +43,7 @@ type KeyReportingManager interface { GetTopKeyRingReporting(r *event.KeyRingReportingRequest) (*event.KeyRingReportingResponse, error) GetSpentKeyReporting(r *event.SpentKeyReportingRequest) (*event.SpentKeyReportingResponse, error) GetUsageReporting(r *event.UsageReportingRequest) (*event.UsageReportingResponse, error) + GetStatistic(r *event.StatisticsRequest) (*event.StatisticsResponse, error) GetKeyReporting(keyId string) (*key.KeyReporting, error) GetEvents(userId, customId string, keyIds []string, start int64, end int64) ([]*event.Event, error) @@ -73,12 +74,12 @@ type AdminServer struct { m KeyManager } -func NewAdminServer(log *zap.Logger, mode string, m KeyManager, krm KeyReportingManager, psm ProviderSettingsManager, cpm CustomProvidersManager, rm RouteManager, pm PoliciesManager, um UserManager, adminPass, xCodioSignSecret string) (*AdminServer, error) { +func NewAdminServer(log *zap.Logger, mode string, m KeyManager, krm KeyReportingManager, psm ProviderSettingsManager, cpm CustomProvidersManager, rm RouteManager, pm PoliciesManager, um UserManager) (*AdminServer, error) { router := gin.New() prod := mode == "production" - router.Use(getAdminLoggerMiddleware(log, "admin", prod, adminPass)) - router.Use(getAdminSignRequestMiddleware(prod, xCodioSignSecret)) + router.Use(getAdminLoggerMiddleware(log, "admin", prod)) + router.Use(getAdminSignRequestMiddleware(prod)) router.GET("/api/health", getGetHealthCheckHandler()) @@ -103,6 +104,8 @@ func NewAdminServer(log *zap.Logger, mode string, m KeyManager, krm KeyReporting router.POST("/api/reporting/spent-keys", getGetSpentKeyMetricsHandler(krm, prod)) router.POST("/api/reporting/usage", getGetUsageMetricsHandler(krm, prod)) + router.POST("/api/reporting/statistic", getGetStatisticHandler(krm, prod)) + router.GET("/api/reporting/custom-ids", getGetCustomIdsHandler(krm, prod)) router.PUT("/api/provider-settings", getCreateProviderSettingHandler(psm, prod)) diff --git a/internal/server/web/admin/middleware.go b/internal/server/web/admin/middleware.go index 23498e5..c050363 100644 --- a/internal/server/web/admin/middleware.go +++ b/internal/server/web/admin/middleware.go @@ -2,27 +2,19 @@ package admin import ( "bytes" - "crypto/hmac" - "crypto/sha1" - "encoding/base64" "fmt" "io" "net/http" "time" "github.com/bricks-cloud/bricksllm/internal/util" + macverification "github.com/bricks-cloud/bricksllm/internal/util/mac-verification" "github.com/gin-gonic/gin" "go.uber.org/zap" ) -func getAdminLoggerMiddleware(log *zap.Logger, prefix string, prod bool, adminPass string) gin.HandlerFunc { +func getAdminLoggerMiddleware(log *zap.Logger, prefix string, prod bool) gin.HandlerFunc { return func(c *gin.Context) { - if len(adminPass) != 0 && c.Request.Header.Get("X-API-KEY") != adminPass { - c.Status(200) - c.Abort() - return - } - cid := util.NewUuid() c.Set(util.STRING_CORRELATION_ID, cid) logWithCid := log.With(zap.String(util.STRING_CORRELATION_ID, cid)) @@ -40,13 +32,13 @@ func getAdminLoggerMiddleware(log *zap.Logger, prefix string, prod bool, adminPa zap.Int("code", c.Writer.Status()), zap.String("method", c.Request.Method), zap.String("path", c.FullPath()), - zap.Int64("lantecyInMs", latency), + zap.Int64("latencyInMs", latency), ) } } } -func getAdminSignRequestMiddleware(prod bool, xCodioSignSecret string) gin.HandlerFunc { +func getAdminSignRequestMiddleware(prod bool) gin.HandlerFunc { return func(c *gin.Context) { log := util.GetLogFromCtx(c) @@ -54,9 +46,12 @@ func getAdminSignRequestMiddleware(prod bool, xCodioSignSecret string) gin.Handl c.Next() return } - sign := c.Request.Header.Get("X-Codio-Sign") - timestamp := c.Request.Header.Get("X-Codio-Sign-Timestamp") - if len(sign) == 0 || len(timestamp) == 0 || len(xCodioSignSecret) == 0 { + + timestamp := c.GetHeader("X-Codio-Sign-Timestamp") + token := c.GetHeader("X-Codio-Sign") + provider := c.GetHeader("X-Codio-Provider") + + if len(token) == 0 || len(timestamp) == 0 || len(provider) == 0 { c.Status(403) c.Abort() return @@ -74,9 +69,10 @@ func getAdminSignRequestMiddleware(prod bool, xCodioSignSecret string) gin.Handl }) return } - data := fmt.Sprintf("%s%s", timestamp, body) + data := fmt.Sprintf("%s%s%s", timestamp, body, provider) c.Request.Body = io.NopCloser(bytes.NewReader(body)) - if !validSign([]byte(data), []byte(xCodioSignSecret), sign) { + + if !macverification.VerifySign(provider, []byte(data), token) { c.Status(403) c.Abort() return @@ -84,11 +80,3 @@ func getAdminSignRequestMiddleware(prod bool, xCodioSignSecret string) gin.Handl c.Next() } } - -func validSign(message, key []byte, messageSign string) bool { - mac := hmac.New(sha1.New, key) - mac.Write(message) - expectedMAC := mac.Sum(nil) - expectedSign := base64.StdEncoding.EncodeToString(expectedMAC) - return messageSign == expectedSign -} diff --git a/internal/server/web/admin/reporting.go b/internal/server/web/admin/reporting.go index fdbdaaf..1ed884c 100644 --- a/internal/server/web/admin/reporting.go +++ b/internal/server/web/admin/reporting.go @@ -7,6 +7,7 @@ import ( "net/http" "time" + "github.com/bricks-cloud/bricksllm/internal/errors" "github.com/bricks-cloud/bricksllm/internal/event" "github.com/bricks-cloud/bricksllm/internal/telemetry" "github.com/bricks-cloud/bricksllm/internal/util" @@ -550,3 +551,82 @@ func getGetUsageMetricsHandler(m KeyReportingManager, prod bool) gin.HandlerFunc c.JSON(http.StatusOK, reportingResponse) } } + +func getGetStatisticHandler(m KeyReportingManager, prod bool) gin.HandlerFunc { + return func(c *gin.Context) { + log := util.GetLogFromCtx(c) + telemetry.Incr("bricksllm.admin.get_get_statistic_handler.requests", nil, 1) + + start := time.Now() + defer func() { + dur := time.Since(start) + telemetry.Timing("bricksllm.admin.get_get_statistic_handler.latency", dur, nil, 1) + }() + + path := "/api/reporting/statistic" + + if c == nil || c.Request == nil { + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/empty-context", + Title: "context is empty error", + Status: http.StatusInternalServerError, + Detail: "gin context is empty", + Instance: path, + }) + return + } + + data, err := io.ReadAll(c.Request.Body) + if err != nil { + logError(log, "error when reading statistics request body", prod, err) + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/request-body-read", + Title: "request body reader error", + Status: http.StatusInternalServerError, + Detail: err.Error(), + Instance: path, + }) + return + } + + request := &event.StatisticsRequest{} + err = json.Unmarshal(data, request) + if err != nil { + logError(log, "error when unmarshalling statistics request body", prod, err) + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/json-unmarshal", + Title: "json unmarshaller error", + Status: http.StatusInternalServerError, + Detail: err.Error(), + Instance: path, + }) + return + } + + statisticResponse, err := m.GetStatistic(request) + if err != nil { + telemetry.Incr("bricksllm.admin.get_get_statistic_handler.get_statistic", nil, 1) + + logError(log, "error when getting statistics", prod, err) + + if _, ok := err.(*errors.NotFoundError); ok { + fmt.Println("NotFoundError:", err.Error()) + c.JSON(http.StatusAccepted, &gin.H{"status": "in_progress", "message": "statistics data is being collected, please try again later"}) + return + } + + c.JSON(http.StatusInternalServerError, &ErrorResponse{ + Type: "/errors/event-reporting-manager", + Title: "statistics reporting error", + Status: http.StatusInternalServerError, + Detail: err.Error(), + Instance: path, + }) + return + } + + telemetry.Incr("bricksllm.admin.get_get_statistic_handler.success", nil, 1) + + c.JSON(http.StatusOK, statisticResponse) + } +} diff --git a/internal/storage/postgresql/event.go b/internal/storage/postgresql/event.go index c3877e6..8a37bcc 100644 --- a/internal/storage/postgresql/event.go +++ b/internal/storage/postgresql/event.go @@ -651,6 +651,539 @@ func (s *Store) GetUsageData(tags []string) (*event.UsageData, error) { return data, nil } +func (s *Store) GetStatisticsData(level event.StatisticLevel, id *string) (*event.StatisticsData, error) { + result := &event.StatisticsData{} + + switch level { + case event.StaticLevels.All: + total, err := s.getTotalCostPack(nil) + if err != nil { + return nil, err + } + + orgPacks, err := s.GetOrgsCostPack() + if err != nil { + return nil, err + } + + orgs := make([]event.ShortOrgStatisticsData, 0, len(orgPacks)) + for orgID, costs := range orgPacks { + orgs = append(orgs, event.ShortOrgStatisticsData{Id: orgID, Costs: costs}) + } + slices.SortFunc(orgs, func(a, b event.ShortOrgStatisticsData) int { + return strings.Compare(a.Id, b.Id) + }) + + result.AllStatisticsData = &event.AllStatisticsData{ + Total: *total, + Orgs: orgs, + } + case event.StaticLevels.Org: + if id == nil || len(*id) == 0 { + return nil, internal_errors.NewValidationError("id must be provided when level is 'org'") + } + + costs, err := s.getTotalCostPack([]string{orgTagPrefix + *id}) + if err != nil { + return nil, err + } + + coursePacks, err := s.GetCoursesCostPack(*id) + if err != nil { + return nil, err + } + + courses := make([]event.ShortCourseStatisticsData, 0, len(coursePacks)) + for courseID, courseCosts := range coursePacks { + courses = append(courses, event.ShortCourseStatisticsData{Id: courseID, Costs: courseCosts}) + } + slices.SortFunc(courses, func(a, b event.ShortCourseStatisticsData) int { + return strings.Compare(a.Id, b.Id) + }) + + dailySpecial, err := s.getDailyDistribution([]string{orgTagPrefix + *id}, codioSpecialTag) + if err != nil { + return nil, err + } + weeklySpecial, err := s.getPeriodDistribution([]string{orgTagPrefix + *id}, codioSpecialTag, "week") + if err != nil { + return nil, err + } + monthlySpecial, err := s.getPeriodDistribution([]string{orgTagPrefix + *id}, codioSpecialTag, "month") + if err != nil { + return nil, err + } + dailyProvided, err := s.getDailyDistribution([]string{orgTagPrefix + *id}, codioProvidedTag) + if err != nil { + return nil, err + } + weeklyProvided, err := s.getPeriodDistribution([]string{orgTagPrefix + *id}, codioProvidedTag, "week") + if err != nil { + return nil, err + } + monthlyProvided, err := s.getPeriodDistribution([]string{orgTagPrefix + *id}, codioProvidedTag, "month") + if err != nil { + return nil, err + } + topFive, err := s.getTopFiveUserSpends([]string{orgTagPrefix + *id}) + if err != nil { + return nil, err + } + + result.OrgStatisticsData = &event.OrgStatisticsData{ + Id: *id, + Costs: *costs, + Courses: courses, + DailySpecialDistribution: dailySpecial, + WeeklySpecialDistribution: weeklySpecial, + MonthlySpecialDistribution: monthlySpecial, + DailyCodioProvidedDistribution: dailyProvided, + WeeklyCodioProvidedDistribution: weeklyProvided, + MonthlyCodioProvidedDistribution: monthlyProvided, + TopFive: topFive, + } + case event.StaticLevels.Course: + if id == nil || len(*id) == 0 { + return nil, internal_errors.NewValidationError("id must be provided when level is 'course'") + } + + costs, err := s.GetCourseCostPack(*id) + if err != nil { + return nil, err + } + + dailySpecial, err := s.getDailyDistribution([]string{courseTagPrefix + *id}, codioSpecialTag) + if err != nil { + return nil, err + } + weeklySpecial, err := s.getPeriodDistribution([]string{courseTagPrefix + *id}, codioSpecialTag, "week") + if err != nil { + return nil, err + } + monthlySpecial, err := s.getPeriodDistribution([]string{courseTagPrefix + *id}, codioSpecialTag, "month") + if err != nil { + return nil, err + } + dailyProvided, err := s.getDailyDistribution([]string{courseTagPrefix + *id}, codioProvidedTag) + if err != nil { + return nil, err + } + weeklyProvided, err := s.getPeriodDistribution([]string{courseTagPrefix + *id}, codioProvidedTag, "week") + if err != nil { + return nil, err + } + monthlyProvided, err := s.getPeriodDistribution([]string{courseTagPrefix + *id}, codioProvidedTag, "month") + if err != nil { + return nil, err + } + topFive, err := s.getTopFiveUserSpends([]string{courseTagPrefix + *id}) + if err != nil { + return nil, err + } + + result.CourseStatisticsData = &event.CourseStatisticsData{ + Id: *id, + Costs: *costs, + DailySpecialDistribution: dailySpecial, + WeeklySpecialDistribution: weeklySpecial, + MonthlySpecialDistribution: monthlySpecial, + DailyCodioProvidedDistribution: dailyProvided, + WeeklyCodioProvidedDistribution: weeklyProvided, + MonthlyCodioProvidedDistribution: monthlyProvided, + TopFive: topFive, + } + default: + return nil, internal_errors.NewValidationError("invalid level") + } + + return result, nil +} + +const ( + orgTagPrefix = "org-tag-" + userTagPrefix = "user-tag-" + courseTagPrefix = "course-tag-" + codioSpecialTag = "codio-special" + codioProvidedTag = "codio-provided" + statisticsLookbackAge = -5 * 30 * 24 * time.Hour +) + +func (s *Store) getGroupedCostPack(groupBy string, filterTags []string) (map[string]event.CostPack, error) { + now := time.Now() + oneMonthAgo := now.Add(-30 * 24 * time.Hour).Unix() + fiveMonthsAgo := now.Add(-5 * 30 * 24 * time.Hour).Unix() + + args := []any{groupBy, oneMonthAgo, fiveMonthsAgo} + conditions := []string{"e.created_at >= $3"} + index := 4 + + for _, tag := range filterTags { + conditions = append(conditions, fmt.Sprintf("e.tags @> $%d", index)) + args = append(args, pq.Array([]string{tag})) + index++ + } + + query := fmt.Sprintf(` + SELECT + regexp_replace(tag.tag, '^' || $1, '') AS grouped_id, + COALESCE(SUM(e.cost_in_usd) FILTER (WHERE e.created_at >= $2 AND e.tags @> ARRAY['%s']::varchar[]), 0) AS codio_provided_one_month, + COALESCE(SUM(e.cost_in_usd) FILTER (WHERE e.tags @> ARRAY['%s']::varchar[]), 0) AS codio_provided_five_month, + COALESCE(SUM(e.cost_in_usd) FILTER (WHERE e.created_at >= $2 AND e.tags @> ARRAY['%s']::varchar[]), 0) AS codio_special_one_month, + COALESCE(SUM(e.cost_in_usd) FILTER (WHERE e.tags @> ARRAY['%s']::varchar[]), 0) AS codio_special_five_month + FROM events e, + LATERAL unnest(e.tags) AS tag(tag) + WHERE tag.tag LIKE $1 || '%%' + AND %s + GROUP BY regexp_replace(tag.tag, '^' || $1, '') + `, codioProvidedTag, codioProvidedTag, codioSpecialTag, codioSpecialTag, strings.Join(conditions, " AND ")) + + ctx, cancel := context.WithTimeout(context.Background(), s.rt) + defer cancel() + + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + result := make(map[string]event.CostPack) + for rows.Next() { + var id string + var pack event.CostPack + + if err2 := rows.Scan( + &id, + &pack.CodioProvided.OneMonth, + &pack.CodioProvided.FiveMonth, + &pack.CodioSpecial.OneMonth, + &pack.CodioSpecial.FiveMonth, + ); err2 != nil { + return nil, err2 + } + + result[id] = pack + } + if err2 := rows.Err(); err2 != nil { + return nil, err2 + } + + return result, nil +} + +func (s *Store) getTotalCostPack(filterTags []string) (*event.CostPack, error) { + now := time.Now() + oneMonthAgo := now.Add(-30 * 24 * time.Hour).Unix() + fiveMonthsAgo := now.Add(-5 * 30 * 24 * time.Hour).Unix() + + args := []any{oneMonthAgo, fiveMonthsAgo} + conditions := []string{"created_at >= $2"} + index := 3 + + for _, tag := range filterTags { + conditions = append(conditions, fmt.Sprintf("tags @> $%d", index)) + args = append(args, pq.Array([]string{tag})) + index++ + } + + query := fmt.Sprintf(` + SELECT + COALESCE(SUM(cost_in_usd) FILTER (WHERE created_at >= $1 AND tags @> ARRAY['%s']::varchar[]), 0) AS codio_provided_one_month, + COALESCE(SUM(cost_in_usd) FILTER (WHERE tags @> ARRAY['%s']::varchar[]), 0) AS codio_provided_five_month, + COALESCE(SUM(cost_in_usd) FILTER (WHERE created_at >= $1 AND tags @> ARRAY['%s']::varchar[]), 0) AS codio_special_one_month, + COALESCE(SUM(cost_in_usd) FILTER (WHERE tags @> ARRAY['%s']::varchar[]), 0) AS codio_special_five_month + FROM events + WHERE %s + `, codioProvidedTag, codioProvidedTag, codioSpecialTag, codioSpecialTag, strings.Join(conditions, " AND ")) + + ctx, cancel := context.WithTimeout(context.Background(), s.rt) + defer cancel() + + pack := &event.CostPack{} + if err := s.db.QueryRowContext(ctx, query, args...).Scan( + &pack.CodioProvided.OneMonth, + &pack.CodioProvided.FiveMonth, + &pack.CodioSpecial.OneMonth, + &pack.CodioSpecial.FiveMonth, + ); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return pack, nil + } + return nil, err + } + + return pack, nil +} + +func (s *Store) GetOrgsCostPack() (map[string]event.CostPack, error) { + return s.getGroupedCostPack(orgTagPrefix, nil) +} + +func (s *Store) GetCoursesCostPack(orgId string) (map[string]event.CostPack, error) { + return s.getGroupedCostPack(courseTagPrefix, []string{orgTagPrefix + orgId}) +} + +func (s *Store) GetCourseCostPack(courseId string) (*event.CostPack, error) { + return s.getTotalCostPack([]string{courseTagPrefix + courseId}) +} + +func (s *Store) getDailyDistribution(filterTags []string, costTypeTag string) ([]event.DailySpendDistributionDataPoint, error) { + start := beginningOfDay(time.Now().UTC()).AddDate(0, 0, -29) + end := beginningOfDay(time.Now().UTC()).AddDate(0, 0, 1) + + rows, err := s.getPeriodKpisRows(filterTags, costTypeTag, "day", start, end) + if err != nil { + return nil, err + } + + rowMap := make(map[string]financialKpiRow, len(rows)) + for _, row := range rows { + rowMap[row.periodStart.Format("Jan-02")] = row + } + + result := make([]event.DailySpendDistributionDataPoint, 0, 30) + for current := start; current.Before(end); current = current.AddDate(0, 0, 1) { + key := current.Format("Jan-02") + row, ok := rowMap[key] + if !ok { + result = append(result, event.DailySpendDistributionDataPoint{Date: key}) + continue + } + result = append(result, event.DailySpendDistributionDataPoint{ + MaxUserSpend: row.max, + MedianUserSpend: row.median, + AvgUserSpend: row.avg, + P95UserSpend: row.p95, + P99UserSpend: row.p99, + Date: key, + }) + } + + return result, nil +} + +func (s *Store) getPeriodDistribution(filterTags []string, costTypeTag, period string) ([]event.PeriodSpendDistributionDataPoint, error) { + var start time.Time + var end time.Time + + now := time.Now().UTC() + switch period { + case "week": + start = beginningOfWeek(now).AddDate(0, 0, -7*20) + end = beginningOfWeek(now).AddDate(0, 0, 7) + case "month": + start = beginningOfMonth(now).AddDate(0, -4, 0) + end = beginningOfMonth(now).AddDate(0, 1, 0) + default: + return nil, internal_errors.NewValidationError("unsupported period") + } + + lookbackStart := time.Now().UTC().Add(statisticsLookbackAge) + if start.Before(lookbackStart) { + start = truncateToPeriodStart(lookbackStart, period) + } + + rows, err := s.getPeriodKpisRows(filterTags, costTypeTag, period, start, end) + if err != nil { + return nil, err + } + + rowMap := make(map[string]financialKpiRow, len(rows)) + for _, row := range rows { + rowMap[formatPeriodLabel(row.periodStart, period)] = row + } + + result := []event.PeriodSpendDistributionDataPoint{} + for current := start; current.Before(end); current = addPeriod(current, period) { + label := formatPeriodLabel(current, period) + row, ok := rowMap[label] + if !ok { + result = append(result, event.PeriodSpendDistributionDataPoint{DatePeriod: label}) + continue + } + result = append(result, event.PeriodSpendDistributionDataPoint{ + MaxUserSpend: row.max, + MedianUserSpend: row.median, + AvgUserSpend: row.avg, + P95UserSpend: row.p95, + P99UserSpend: row.p99, + DatePeriod: label, + }) + } + + return result, nil +} + +type financialKpiRow struct { + periodStart time.Time + max float64 + median float64 + avg float64 + p95 float64 + p99 float64 +} + +func (s *Store) getPeriodKpisRows(filterTags []string, costTypeTag, period string, start, end time.Time) ([]financialKpiRow, error) { + args := []any{userTagPrefix, pq.Array([]string{costTypeTag}), start.Unix(), end.Unix()} + conditions := []string{"tag.tag LIKE $1 || '%'", "e.tags @> $2", "e.created_at >= $3", "e.created_at < $4"} + index := 5 + + for _, tag := range filterTags { + conditions = append(conditions, fmt.Sprintf("e.tags @> $%d", index)) + args = append(args, pq.Array([]string{tag})) + index++ + } + + periodExpr := fmt.Sprintf("date_trunc('%s', timezone('UTC', to_timestamp(e.created_at)))", period) + query := fmt.Sprintf(` + WITH per_user_period AS ( + SELECT + %s AS period_start, + regexp_replace(tag.tag, '^' || $1, '') AS user_id, + COALESCE(SUM(e.cost_in_usd), 0) AS user_period_cost + FROM events e, + LATERAL unnest(e.tags) AS tag(tag) + WHERE %s + GROUP BY %s, regexp_replace(tag.tag, '^' || $1, '') + ) + SELECT + period_start, + COALESCE(MAX(user_period_cost), 0) AS max_user_spend, + COALESCE(percentile_cont(0.5) WITHIN GROUP (ORDER BY user_period_cost), 0) AS median_user_spend, + COALESCE(AVG(user_period_cost), 0) AS avg_user_spend, + COALESCE(percentile_cont(0.95) WITHIN GROUP (ORDER BY user_period_cost), 0) AS p95_user_spend, + COALESCE(percentile_cont(0.99) WITHIN GROUP (ORDER BY user_period_cost), 0) AS p99_user_spend + FROM per_user_period + GROUP BY period_start + ORDER BY period_start + `, periodExpr, strings.Join(conditions, " AND "), periodExpr) + + ctx, cancel := context.WithTimeout(context.Background(), s.rt) + defer cancel() + + queryRows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer queryRows.Close() + + result := []financialKpiRow{} + for queryRows.Next() { + var row financialKpiRow + if err := queryRows.Scan(&row.periodStart, &row.max, &row.median, &row.avg, &row.p95, &row.p99); err != nil { + return nil, err + } + result = append(result, row) + } + if err := queryRows.Err(); err != nil { + return nil, err + } + + return result, nil +} + +func (s *Store) getTopFiveUserSpends(filterTags []string) ([]event.TopFiveUserSpend, error) { + end := time.Now().UTC() + start := end.Add(-30 * 24 * time.Hour) + + args := []any{userTagPrefix, start.Unix(), end.Unix()} + conditions := []string{"tag.tag LIKE $1 || '%'", "e.created_at >= $2", "e.created_at < $3"} + index := 4 + + for _, tag := range filterTags { + conditions = append(conditions, fmt.Sprintf("e.tags @> $%d", index)) + args = append(args, pq.Array([]string{tag})) + index++ + } + + query := fmt.Sprintf(` + SELECT + regexp_replace(tag.tag, '^' || $1, '') AS user_id, + COALESCE(SUM(e.cost_in_usd), 0) AS spend_last_month + FROM events e, + LATERAL unnest(e.tags) AS tag(tag) + WHERE %s + GROUP BY regexp_replace(tag.tag, '^' || $1, '') + ORDER BY spend_last_month DESC, user_id ASC + LIMIT 5 + `, strings.Join(conditions, " AND ")) + + ctx, cancel := context.WithTimeout(context.Background(), s.rt) + defer cancel() + + rows, err := s.db.QueryContext(ctx, query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + result := []event.TopFiveUserSpend{} + for rows.Next() { + var row event.TopFiveUserSpend + if err := rows.Scan(&row.UserId, &row.SpendLastMonth); err != nil { + return nil, err + } + result = append(result, row) + } + if err := rows.Err(); err != nil { + return nil, err + } + + return result, nil +} + +func beginningOfDay(t time.Time) time.Time { + return time.Date(t.Year(), t.Month(), t.Day(), 0, 0, 0, 0, time.UTC) +} + +func beginningOfWeek(t time.Time) time.Time { + dayStart := beginningOfDay(t) + weekday := int(dayStart.Weekday()) + if weekday == 0 { + weekday = 7 + } + return dayStart.AddDate(0, 0, -(weekday - 1)) +} + +func beginningOfMonth(t time.Time) time.Time { + return time.Date(t.Year(), t.Month(), 1, 0, 0, 0, 0, time.UTC) +} + +func truncateToPeriodStart(t time.Time, period string) time.Time { + switch period { + case "day": + return beginningOfDay(t) + case "week": + return beginningOfWeek(t) + case "month": + return beginningOfMonth(t) + default: + return beginningOfDay(t) + } +} + +func addPeriod(t time.Time, period string) time.Time { + switch period { + case "week": + return t.AddDate(0, 0, 7) + case "month": + return t.AddDate(0, 1, 0) + default: + return t.AddDate(0, 0, 1) + } +} + +func formatPeriodLabel(start time.Time, period string) string { + switch period { + case "week": + end := start.AddDate(0, 0, 6) + return fmt.Sprintf("%s/%s", start.Format("Jan-02"), end.Format("Jan-02")) + case "month": + return start.Format("Jan") + default: + return start.Format("Jan-02") + } +} + func (s *Store) GetAggregatedEventByDayDataPoints(start, end int64, keyIds []string) ([]*event.DataPointV2, error) { conditionBlock := fmt.Sprintf("WHERE time_stamp >= %d AND time_stamp < %d ", start, end) if len(keyIds) != 0 { diff --git a/internal/storage/redis/statistic-cache.go b/internal/storage/redis/statistic-cache.go new file mode 100644 index 0000000..e3a7a3b --- /dev/null +++ b/internal/storage/redis/statistic-cache.go @@ -0,0 +1,98 @@ +package redis + +import ( + "context" + "encoding/json" + "errors" + "time" + + "github.com/bricks-cloud/bricksllm/internal/event" + "github.com/redis/go-redis/v9" +) + +const inProgressKeyPrefix = "statistic_in_progress_" + +type StatisticCache struct { + client *redis.Client + wt time.Duration + rt time.Duration +} + +func NewStatisticCache(c *redis.Client, wt time.Duration, rt time.Duration) *StatisticCache { + return &StatisticCache{ + client: c, + wt: wt, + rt: rt, + } +} + +func (c *StatisticCache) Set(key string, value *event.StatisticsData, ttl time.Duration) error { + if value == nil { + return errors.New("statistic data to set is nil") + } + ctx, cancel := context.WithTimeout(context.Background(), c.wt) + defer cancel() + bs, err := json.Marshal(value) + if err != nil { + return err + } + err = c.client.Set(ctx, key, bs, ttl).Err() + if err != nil { + return err + } + + return nil +} + +func (c *StatisticCache) Delete(key string) error { + ctx, cancel := context.WithTimeout(context.Background(), c.wt) + defer cancel() + err := c.client.Del(ctx, key).Err() + if err != nil { + return err + } + + return nil +} + +func (c *StatisticCache) Get(key string) (*event.StatisticsData, error) { + ctx, cancel := context.WithTimeout(context.Background(), c.rt) + defer cancel() + + result := c.client.Get(ctx, key) + err := result.Err() + if err != nil { + return nil, err + } + + bs, err := result.Bytes() + if err != nil { + return nil, err + } + + var stat event.StatisticsData + err = json.Unmarshal(bs, &stat) + if err != nil { + return nil, err + } + + return &stat, nil +} + +func (c *StatisticCache) TryMarkInProgress(key string) (bool, error) { + k := inProgressKeyPrefix + key + ctx, cancel := context.WithTimeout(context.Background(), c.wt) + defer cancel() + return c.client.SetNX(ctx, k, true, time.Minute*10).Result() +} + +func (c *StatisticCache) DeleteInProgress(key string) error { + k := inProgressKeyPrefix + key + ctx, cancel := context.WithTimeout(context.Background(), c.wt) + defer cancel() + err := c.client.Del(ctx, k).Err() + if err != nil { + return err + } + return nil +} diff --git a/internal/util/mac-verification/verification.go b/internal/util/mac-verification/verification.go new file mode 100644 index 0000000..f318177 --- /dev/null +++ b/internal/util/mac-verification/verification.go @@ -0,0 +1,86 @@ +package mac_verification + +import ( + "encoding/base64" + "fmt" + "log" + "os" + "path/filepath" + + "github.com/google/tink/go/insecurecleartextkeyset" + "github.com/google/tink/go/keyset" + "github.com/google/tink/go/mac" + "github.com/google/tink/go/tink" +) + +const keySetSubDir = "./tink" + +var macCache map[string]tink.MAC = make(map[string]tink.MAC) + +func init() { + err := loadKeySets() + if err != nil { + log.Fatal(err) + } +} + +func loadKeySets() error { + execDir, err := os.Executable() + if err != nil { + return fmt.Errorf("error getting executable path: %v", err) + } + execDir = filepath.Dir(execDir) + keySetDir := filepath.Join(execDir, keySetSubDir) + files, err := os.ReadDir(keySetDir) + if err != nil { + return fmt.Errorf("error reading directory %s: %v", keySetDir, err) + } + for _, file := range files { + if file.IsDir() { + continue + } + filePath := filepath.Join(keySetDir, file.Name()) + f, err := os.Open(filePath) + if err != nil { + return fmt.Errorf("error opening keyset file %s: %v", filePath, err) + } + defer f.Close() + + jsonReader := keyset.NewJSONReader(f) + + kh, err := insecurecleartextkeyset.Read(jsonReader) + if err != nil { + return fmt.Errorf("error reading keyset from file %s: %v", filePath, err) + } + + macPrimitive, err := mac.New(kh) + if err != nil { + return fmt.Errorf("error creating MAC primitive from keyset %s: %v", filePath, err) + } + + providerName := file.Name() + if len(providerName) > 5 && providerName[len(providerName)-5:] == ".json" { + providerName = providerName[:len(providerName)-5] + } + macCache[providerName] = macPrimitive + } + return nil +} + +func VerifySign(provider string, data []byte, token string) bool { + decodedSignature, err := base64.StdEncoding.DecodeString(token) + if err != nil { + fmt.Printf("Error decoding token from Base64: %v\n", err) + return false + } + macPrimitive, ok := macCache[provider] + if !ok { + fmt.Printf("MAC primitive for provider %s not found\n", provider) + return false + } + err = macPrimitive.VerifyMAC(decodedSignature, data) + if err != nil { + return false + } + return true +}