diff --git a/internal/errors/cost_limit_err.go b/internal/errors/cost_limit_err.go index 15c92deb..ae84ab47 100644 --- a/internal/errors/cost_limit_err.go +++ b/internal/errors/cost_limit_err.go @@ -1,17 +1,26 @@ package errors type CostLimitError struct { - message string + message string + timeUnit string } -func NewCostLimitError(msg string) *CostLimitError { - return &CostLimitError{ +func NewCostLimitError(msg string, timeUnit ...string) *CostLimitError { + cle := &CostLimitError{ message: msg, } + if len(timeUnit) != 0 { + cle.timeUnit = timeUnit[0] + } + return cle } func (cle *CostLimitError) Error() string { return cle.message } +func (cle *CostLimitError) TimeUnit() string { + return cle.timeUnit +} + func (rle *CostLimitError) CostLimit() {} diff --git a/internal/key/key.go b/internal/key/key.go index 94eb537c..13e629a2 100644 --- a/internal/key/key.go +++ b/internal/key/key.go @@ -3,6 +3,7 @@ package key import ( "errors" "fmt" + "slices" "strings" "time" @@ -11,27 +12,41 @@ import ( const RevokedReasonExpired string = "expired" +type ExtendedBudgetLimitItem struct { + LimitInUsdOverTime float64 `json:"costLimitOverTime"` + Unit TimeUnit `json:"unit"` +} + +func (i *ExtendedBudgetLimitItem) ExtendedKey(k string) string { + return fmt.Sprintf("%s-%s", k, i.Unit) +} + +type ExtendedBudgetLimit struct { + Items []ExtendedBudgetLimitItem `json:"items"` +} + type UpdateKey struct { - Name string `json:"name"` - UpdatedAt int64 `json:"updatedAt"` - Tags []string `json:"tags"` - Revoked *bool `json:"revoked"` - RevokedReason string `json:"revokedReason"` - Key string `json:"key"` - SettingId string `json:"settingId"` - SettingIds []string `json:"settingIds"` - CostLimitInUsd *float64 `json:"costLimitInUsd"` - CostLimitInUsdOverTime *float64 `json:"costLimitInUsdOverTime"` - CostLimitInUsdUnit *TimeUnit `json:"costLimitInUsdUnit"` - RateLimitOverTime *int `json:"rateLimitOverTime"` - RateLimitUnit *TimeUnit `json:"rateLimitUnit"` - RequestsLimit *int `json:"requestsLimit"` - AllowedPaths *[]PathConfig `json:"allowedPaths,omitempty"` - ShouldLogRequest *bool `json:"shouldLogRequest"` - ShouldLogResponse *bool `json:"shouldLogResponse"` - RotationEnabled *bool `json:"rotationEnabled"` - PolicyId *string `json:"policyId"` - IsKeyNotHashed *bool `json:"isKeyNotHashed"` + Name string `json:"name"` + UpdatedAt int64 `json:"updatedAt"` + Tags []string `json:"tags"` + Revoked *bool `json:"revoked"` + RevokedReason string `json:"revokedReason"` + Key string `json:"key"` + SettingId string `json:"settingId"` + SettingIds []string `json:"settingIds"` + CostLimitInUsd *float64 `json:"costLimitInUsd"` + CostLimitInUsdOverTime *float64 `json:"costLimitInUsdOverTime"` + CostLimitInUsdUnit *TimeUnit `json:"costLimitInUsdUnit"` + RateLimitOverTime *int `json:"rateLimitOverTime"` + RateLimitUnit *TimeUnit `json:"rateLimitUnit"` + RequestsLimit *int `json:"requestsLimit"` + AllowedPaths *[]PathConfig `json:"allowedPaths,omitempty"` + ShouldLogRequest *bool `json:"shouldLogRequest"` + ShouldLogResponse *bool `json:"shouldLogResponse"` + RotationEnabled *bool `json:"rotationEnabled"` + PolicyId *string `json:"policyId"` + IsKeyNotHashed *bool `json:"isKeyNotHashed"` + ExtendedBudgetLimit *ExtendedBudgetLimit `json:"extendedBudgetLimit"` } func (uk *UpdateKey) Validate() error { @@ -125,7 +140,7 @@ func (uk *UpdateKey) Validate() error { return internal_errors.NewValidationError("rate limit unit can not be empty if rate limit over time is specified") } - if *uk.RateLimitOverTime != 0 && *uk.RateLimitUnit != HourTimeUnit && *uk.RateLimitUnit != MinuteTimeUnit && *uk.RateLimitUnit != SecondTimeUnit && *uk.RateLimitUnit != DayTimeUnit { + if *uk.RateLimitOverTime != 0 && !slices.Contains(AllowedTimeUnits, *uk.RateLimitUnit) { return internal_errors.NewValidationError("rate limit unit can not be identified") } } @@ -143,10 +158,13 @@ func (uk *UpdateKey) Validate() error { return internal_errors.NewValidationError("cost limit unit can not be empty if cost limit over time is specified") } - if *uk.CostLimitInUsdOverTime != 0 && *uk.CostLimitInUsdUnit != DayTimeUnit && *uk.CostLimitInUsdUnit != HourTimeUnit && *uk.CostLimitInUsdUnit != MonthTimeUnit && *uk.CostLimitInUsdUnit != MinuteTimeUnit { + if *uk.CostLimitInUsdOverTime != 0 && !slices.Contains(AllowedTimeUnits, *uk.CostLimitInUsdUnit) { return internal_errors.NewValidationError("cost limit unit can not be identified") } } + if err := validateExtendedBudgetLimit(uk.ExtendedBudgetLimit); err != nil { + return err + } return nil } @@ -157,28 +175,29 @@ type PathConfig struct { } type RequestKey struct { - Name string `json:"name"` - CreatedAt int64 `json:"createdAt"` - UpdatedAt int64 `json:"updatedAt"` - Tags []string `json:"tags"` - KeyId string `json:"keyId"` - Key string `json:"key"` - CostLimitInUsd float64 `json:"costLimitInUsd"` - CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"` - CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"` - RateLimitOverTime int `json:"rateLimitOverTime"` - RateLimitUnit TimeUnit `json:"rateLimitUnit"` - Ttl string `json:"ttl"` - KeyRing string `json:"keyRing"` - SettingId string `json:"settingId"` - AllowedPaths []PathConfig `json:"allowedPaths"` - SettingIds []string `json:"settingIds"` - ShouldLogRequest bool `json:"shouldLogRequest"` - ShouldLogResponse bool `json:"shouldLogResponse"` - RotationEnabled bool `json:"rotationEnabled"` - PolicyId string `json:"policyId"` - IsKeyNotHashed bool `json:"isKeyNotHashed"` - RequestsLimit int `json:"requestsLimit"` + Name string `json:"name"` + CreatedAt int64 `json:"createdAt"` + UpdatedAt int64 `json:"updatedAt"` + Tags []string `json:"tags"` + KeyId string `json:"keyId"` + Key string `json:"key"` + CostLimitInUsd float64 `json:"costLimitInUsd"` + CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"` + CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"` + RateLimitOverTime int `json:"rateLimitOverTime"` + RateLimitUnit TimeUnit `json:"rateLimitUnit"` + Ttl string `json:"ttl"` + KeyRing string `json:"keyRing"` + SettingId string `json:"settingId"` + AllowedPaths []PathConfig `json:"allowedPaths"` + SettingIds []string `json:"settingIds"` + ShouldLogRequest bool `json:"shouldLogRequest"` + ShouldLogResponse bool `json:"shouldLogResponse"` + RotationEnabled bool `json:"rotationEnabled"` + PolicyId string `json:"policyId"` + IsKeyNotHashed bool `json:"isKeyNotHashed"` + RequestsLimit int `json:"requestsLimit"` + ExtendedBudgetLimit *ExtendedBudgetLimit `json:"extendedBudgetLimit"` } func (rk *RequestKey) Validate() error { @@ -285,7 +304,7 @@ func (rk *RequestKey) Validate() error { return internal_errors.NewValidationError("rate limit unit can not be empty if rate limit over time is specified") } - if rk.RateLimitUnit != HourTimeUnit && rk.RateLimitUnit != MinuteTimeUnit && rk.RateLimitUnit != SecondTimeUnit && rk.RateLimitUnit != DayTimeUnit { + if !slices.Contains(AllowedTimeUnits, rk.RateLimitUnit) { return internal_errors.NewValidationError("rate limit unit can not be identified") } } @@ -295,11 +314,40 @@ func (rk *RequestKey) Validate() error { return internal_errors.NewValidationError("cost limit unit can not be empty if cost limit over time is specified") } - if rk.CostLimitInUsdUnit != DayTimeUnit && rk.CostLimitInUsdUnit != HourTimeUnit && rk.CostLimitInUsdUnit != MonthTimeUnit && rk.CostLimitInUsdUnit != MinuteTimeUnit { + if !slices.Contains(AllowedTimeUnits, rk.CostLimitInUsdUnit) { return internal_errors.NewValidationError("cost limit unit can not be identified") } } + if err := validateExtendedBudgetLimit(rk.ExtendedBudgetLimit); err != nil { + return err + } + return nil +} + +func validateExtendedBudgetLimit(ebl *ExtendedBudgetLimit) error { + if ebl == nil { + return nil + } + + seenUnits := make(map[TimeUnit]struct{}, len(ebl.Items)) + for index, item := range ebl.Items { + if item.LimitInUsdOverTime < 0 { + return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].costLimitOverTime is invalid", index)) + } + if len(item.Unit) == 0 { + return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].unit is invalid", index)) + } + + if !slices.Contains(AllowedTimeUnits, item.Unit) { + return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].unit can not be identified", index)) + } + + if _, exists := seenUnits[item.Unit]; exists { + return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].unit is duplicated", index)) + } + seenUnits[item.Unit] = struct{}{} + } return nil } @@ -310,34 +358,45 @@ const ( MinuteTimeUnit TimeUnit = "m" SecondTimeUnit TimeUnit = "s" DayTimeUnit TimeUnit = "d" + WeekTimeUnit TimeUnit = "w" MonthTimeUnit TimeUnit = "mo" ) +var AllowedTimeUnits = []TimeUnit{ + HourTimeUnit, + MinuteTimeUnit, + SecondTimeUnit, + DayTimeUnit, + WeekTimeUnit, + MonthTimeUnit, +} + type ResponseKey struct { - Name string `json:"name"` - CreatedAt int64 `json:"createdAt"` - UpdatedAt int64 `json:"updatedAt"` - Tags []string `json:"tags"` - KeyId string `json:"keyId"` - Revoked bool `json:"revoked"` - Key string `json:"key"` - RevokedReason string `json:"revokedReason"` - CostLimitInUsd float64 `json:"costLimitInUsd"` - CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"` - CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"` - RateLimitOverTime int `json:"rateLimitOverTime"` - RateLimitUnit TimeUnit `json:"rateLimitUnit"` - RequestsLimit int `json:"requestsLimit"` - Ttl string `json:"ttl"` - KeyRing string `json:"keyRing"` - SettingId string `json:"settingId"` - AllowedPaths []PathConfig `json:"allowedPaths"` - SettingIds []string `json:"settingIds"` - ShouldLogRequest bool `json:"shouldLogRequest"` - ShouldLogResponse bool `json:"shouldLogResponse"` - RotationEnabled bool `json:"rotationEnabled"` - PolicyId string `json:"policyId"` - IsKeyNotHashed bool `json:"isKeyNotHashed"` + Name string `json:"name"` + CreatedAt int64 `json:"createdAt"` + UpdatedAt int64 `json:"updatedAt"` + Tags []string `json:"tags"` + KeyId string `json:"keyId"` + Revoked bool `json:"revoked"` + Key string `json:"key"` + RevokedReason string `json:"revokedReason"` + CostLimitInUsd float64 `json:"costLimitInUsd"` + CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"` + CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"` + RateLimitOverTime int `json:"rateLimitOverTime"` + RateLimitUnit TimeUnit `json:"rateLimitUnit"` + RequestsLimit int `json:"requestsLimit"` + Ttl string `json:"ttl"` + KeyRing string `json:"keyRing"` + SettingId string `json:"settingId"` + AllowedPaths []PathConfig `json:"allowedPaths"` + SettingIds []string `json:"settingIds"` + ShouldLogRequest bool `json:"shouldLogRequest"` + ShouldLogResponse bool `json:"shouldLogResponse"` + RotationEnabled bool `json:"rotationEnabled"` + PolicyId string `json:"policyId"` + IsKeyNotHashed bool `json:"isKeyNotHashed"` + ExtendedBudgetLimit *ExtendedBudgetLimit `json:"extendedBudgetLimit"` } func (rk *ResponseKey) GetSettingIds() []string { diff --git a/internal/key/key_test.go b/internal/key/key_test.go new file mode 100644 index 00000000..bdacafaa --- /dev/null +++ b/internal/key/key_test.go @@ -0,0 +1,57 @@ +package key + +import "testing" + +func TestValidateExtendedBudgetLimit(t *testing.T) { + tests := []struct { + name string + ebl *ExtendedBudgetLimit + wantErr string + }{ + { + name: "nil limit is valid", + ebl: nil, + wantErr: "", + }, + { + name: "duplicate units are rejected", + ebl: &ExtendedBudgetLimit{ + Items: []ExtendedBudgetLimitItem{ + {LimitInUsdOverTime: 1, Unit: HourTimeUnit}, + {LimitInUsdOverTime: 2, Unit: HourTimeUnit}, + }, + }, + wantErr: "extendedBudgetLimit.items[1].unit is duplicated", + }, + { + name: "unique units are valid", + ebl: &ExtendedBudgetLimit{ + Items: []ExtendedBudgetLimitItem{ + {LimitInUsdOverTime: 1, Unit: HourTimeUnit}, + {LimitInUsdOverTime: 2, Unit: MinuteTimeUnit}, + }, + }, + wantErr: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateExtendedBudgetLimit(tt.ebl) + if tt.wantErr == "" { + if err != nil { + t.Fatalf("expected no error, got %v", err) + } + return + } + + if err == nil { + t.Fatalf("expected error %q, got nil", tt.wantErr) + } + + if err.Error() != tt.wantErr { + t.Fatalf("expected error %q, got %q", tt.wantErr, err.Error()) + } + }) + } +} diff --git a/internal/message/consumer.go b/internal/message/consumer.go index 189043fd..17eaad32 100644 --- a/internal/message/consumer.go +++ b/internal/message/consumer.go @@ -18,7 +18,7 @@ type Consumer struct { } type recorder interface { - RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit) error + RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit, extendedLimits *key.ExtendedBudgetLimit) error RecordUserSpend(userId string, micros int64, costLimitUnit key.TimeUnit) error RecordEvent(e *event.Event) error RecordKeyRequestSpent(keyId string) error diff --git a/internal/message/handler.go b/internal/message/handler.go index 7c973ef3..3a82bee0 100644 --- a/internal/message/handler.go +++ b/internal/message/handler.go @@ -182,6 +182,10 @@ type costLimitError interface { CostLimit() } +type costLimitTimeUnitError interface { + TimeUnit() string +} + type rateLimitError interface { Error() string RateLimit() @@ -247,7 +251,12 @@ func (h *Handler) handleValidationResult(kc *key.ResponseKey, cost float64) erro if _, ok := err.(costLimitError); ok { telemetry.Incr("bricksllm.message.handler.handle_validation_result.cost_limit_error", nil, 1) - err = h.ac.Set(kc.KeyId, kc.CostLimitInUsdUnit) + costLimitUnit := kc.CostLimitInUsdUnit + if cle, ok := err.(costLimitTimeUnitError); ok && len(cle.TimeUnit()) != 0 { + costLimitUnit = key.TimeUnit(cle.TimeUnit()) + } + + err = h.ac.Set(kc.KeyId, costLimitUnit) if err != nil { telemetry.Incr("bricksllm.message.handler.handle_validation_result.set_cost_limit_error", nil, 1) return err @@ -340,7 +349,7 @@ func (h *Handler) HandleEventWithRequestAndResponse(m Message) error { if e.Event.CostInUsd != 0 { micros := int64(e.Event.CostInUsd * 1000000) - err = h.recorder.RecordKeySpend(e.Event.KeyId, micros, e.Key.CostLimitInUsdUnit) + err = h.recorder.RecordKeySpend(e.Event.KeyId, micros, e.Key.CostLimitInUsdUnit, e.Key.ExtendedBudgetLimit) if err != nil { telemetry.Incr("bricksllm.message.handler.handle_event_with_request_and_response.record_key_spend_error", nil, 1) h.log.Debug("error when recording key spend", zap.Error(err)) diff --git a/internal/message/handler_test.go b/internal/message/handler_test.go new file mode 100644 index 00000000..63e87723 --- /dev/null +++ b/internal/message/handler_test.go @@ -0,0 +1,71 @@ +package message + +import ( + "testing" + + internal_errors "github.com/bricks-cloud/bricksllm/internal/errors" + "github.com/bricks-cloud/bricksllm/internal/key" +) + +type stubValidator struct { + err error +} + +func (s *stubValidator) Validate(k *key.ResponseKey, promptCost float64) error { + return s.err +} + +type stubAccessCache struct { + key string + timeUnit key.TimeUnit +} + +func (s *stubAccessCache) Set(key string, timeUnit key.TimeUnit) error { + s.key = key + s.timeUnit = timeUnit + return nil +} + +func TestHandleValidationResultUsesBreachedCostLimitUnitWhenProvided(t *testing.T) { + ac := &stubAccessCache{} + h := &Handler{ + v: &stubValidator{err: internal_errors.NewCostLimitError("extended cost limit reached", string(key.WeekTimeUnit))}, + ac: ac, + } + + err := h.handleValidationResult(&key.ResponseKey{ + KeyId: "test-key", + CostLimitInUsdUnit: key.DayTimeUnit, + }, 0) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + + if ac.key != "test-key" { + t.Fatalf("expected cache key %q, got %q", "test-key", ac.key) + } + + if ac.timeUnit != key.WeekTimeUnit { + t.Fatalf("expected time unit %q, got %q", key.WeekTimeUnit, ac.timeUnit) + } +} + +func TestHandleValidationResultFallsBackToKeyCostLimitUnit(t *testing.T) { + ac := &stubAccessCache{} + h := &Handler{ + v: &stubValidator{err: internal_errors.NewCostLimitError("standard cost limit reached")}, + ac: ac, + } + + err := h.handleValidationResult(&key.ResponseKey{ + KeyId: "test-key", + CostLimitInUsdUnit: key.DayTimeUnit, + }, 0) + if err != nil { + t.Fatalf("expected nil error, got %v", err) + } + + if ac.timeUnit != key.DayTimeUnit { + t.Fatalf("expected time unit %q, got %q", key.DayTimeUnit, ac.timeUnit) + } +} diff --git a/internal/recorder/recorder.go b/internal/recorder/recorder.go index 5c8b73fe..3c3463bd 100644 --- a/internal/recorder/recorder.go +++ b/internal/recorder/recorder.go @@ -60,7 +60,7 @@ func (r *Recorder) RecordUserSpend(userId string, micros int64, costLimitUnit ke return nil } -func (r *Recorder) RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit) error { +func (r *Recorder) RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit, extendedLimits *key.ExtendedBudgetLimit) error { err := r.s.IncrementCounter(keyId, micros) if err != nil { return err @@ -73,6 +73,15 @@ func (r *Recorder) RecordKeySpend(keyId string, micros int64, costLimitUnit key. } } + if extendedLimits != nil { + for _, item := range extendedLimits.Items { + err = r.c.IncrementCounter(item.ExtendedKey(keyId), item.Unit, int64(micros)) + if err != nil { + return err + } + } + } + return nil } diff --git a/internal/storage/postgresql/key.go b/internal/storage/postgresql/key.go index e92f114f..2cba7896 100644 --- a/internal/storage/postgresql/key.go +++ b/internal/storage/postgresql/key.go @@ -5,14 +5,56 @@ import ( "database/sql" "encoding/json" "fmt" - "github.com/bricks-cloud/bricksllm/internal/event" "strings" + "github.com/bricks-cloud/bricksllm/internal/event" + internal_errors "github.com/bricks-cloud/bricksllm/internal/errors" "github.com/bricks-cloud/bricksllm/internal/key" "github.com/lib/pq" ) +func marshalExtendedBudgetLimit(limit *key.ExtendedBudgetLimit) ([]byte, error) { + if limit == nil { + return nil, nil + } + + return json.Marshal(limit) +} + +func unmarshalExtendedBudgetLimit(data []byte) (*key.ExtendedBudgetLimit, error) { + if data == nil { + return nil, nil + } + + var limit *key.ExtendedBudgetLimit + if err := json.Unmarshal(data, &limit); err != nil { + return nil, err + } + + return limit, nil +} + +func hydrateKeyJSONFields(k *key.ResponseKey, allowedPathsData, extendedBudgetLimitData []byte) error { + if len(allowedPathsData) != 0 { + pathConfigs := []key.PathConfig{} + if err := json.Unmarshal(allowedPathsData, &pathConfigs); err != nil { + return err + } + + k.AllowedPaths = pathConfigs + } + + extendedBudgetLimit, err := unmarshalExtendedBudgetLimit(extendedBudgetLimitData) + if err != nil { + return err + } + + k.ExtendedBudgetLimit = extendedBudgetLimit + + return nil +} + func (s *Store) CreateKeysTable() error { createTableQuery := ` CREATE TABLE IF NOT EXISTS keys ( @@ -57,7 +99,7 @@ func (s *Store) AlterKeysTable() error { END IF; END $$; - ALTER TABLE keys ADD COLUMN IF NOT EXISTS setting_id VARCHAR(255), ADD COLUMN IF NOT EXISTS allowed_paths JSONB, ADD COLUMN IF NOT EXISTS setting_ids VARCHAR(255)[] NOT NULL DEFAULT ARRAY[]::VARCHAR(255)[], ADD COLUMN IF NOT EXISTS should_log_request BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS should_log_response BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS rotation_enabled BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS policy_id VARCHAR(255) NOT NULL DEFAULT '', ADD COLUMN IF NOT EXISTS is_key_not_hashed BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS requests_limit INT NOT NULL DEFAULT 0; + ALTER TABLE keys ADD COLUMN IF NOT EXISTS setting_id VARCHAR(255), ADD COLUMN IF NOT EXISTS allowed_paths JSONB, ADD COLUMN IF NOT EXISTS setting_ids VARCHAR(255)[] NOT NULL DEFAULT ARRAY[]::VARCHAR(255)[], ADD COLUMN IF NOT EXISTS should_log_request BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS should_log_response BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS rotation_enabled BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS policy_id VARCHAR(255) NOT NULL DEFAULT '', ADD COLUMN IF NOT EXISTS is_key_not_hashed BOOLEAN NOT NULL DEFAULT FALSE, ADD COLUMN IF NOT EXISTS requests_limit INT NOT NULL DEFAULT 0, ADD COLUMN IF NOT EXISTS extended_budget_limit JSONB; ` ctxTimeout, cancel := context.WithTimeout(context.Background(), s.wt) @@ -166,6 +208,7 @@ func (s *Store) GetKeys(tags, keyIds []string, provider string) ([]*key.Response var k key.ResponseKey var settingId sql.NullString var data []byte + var extendedBudgetLimitData []byte if err := rows.Scan( &k.Name, &k.CreatedAt, @@ -191,6 +234,7 @@ func (s *Store) GetKeys(tags, keyIds []string, provider string) ([]*key.Response &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &extendedBudgetLimitData, ); err != nil { return nil, err } @@ -198,13 +242,8 @@ func (s *Store) GetKeys(tags, keyIds []string, provider string) ([]*key.Response pk := &k pk.SettingId = settingId.String - if len(data) != 0 { - pathConfigs := []key.PathConfig{} - if err := json.Unmarshal(data, &pathConfigs); err != nil { - return nil, err - } - - pk.AllowedPaths = pathConfigs + if err := hydrateKeyJSONFields(pk, data, extendedBudgetLimitData); err != nil { + return nil, err } keys = append(keys, pk) @@ -292,6 +331,7 @@ func (s *Store) GetKeysV2(tags, keyIds []string, revoked *bool, limit, offset in var k key.ResponseKey var settingId sql.NullString var data []byte + var extendedBudgetLimitData []byte if err := rows.Scan( &k.Name, &k.CreatedAt, @@ -317,6 +357,7 @@ func (s *Store) GetKeysV2(tags, keyIds []string, revoked *bool, limit, offset in &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &extendedBudgetLimitData, ); err != nil { return nil, err } @@ -324,13 +365,8 @@ func (s *Store) GetKeysV2(tags, keyIds []string, revoked *bool, limit, offset in pk := &k pk.SettingId = settingId.String - if len(data) != 0 { - pathConfigs := []key.PathConfig{} - if err := json.Unmarshal(data, &pathConfigs); err != nil { - return nil, err - } - - pk.AllowedPaths = pathConfigs + if err := hydrateKeyJSONFields(pk, data, extendedBudgetLimitData); err != nil { + return nil, err } keys = append(keys, pk) @@ -371,6 +407,7 @@ func (s *Store) GetKeyByHash(hash string) (*key.ResponseKey, error) { var k key.ResponseKey var settingId sql.NullString var data []byte + var extendedBudgetLimitData []byte err := s.db.QueryRowContext(ctxTimeout, "SELECT * FROM keys WHERE key = $1", hash).Scan( &k.Name, @@ -397,6 +434,7 @@ func (s *Store) GetKeyByHash(hash string) (*key.ResponseKey, error) { &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &extendedBudgetLimitData, ) if err != nil { @@ -409,13 +447,8 @@ func (s *Store) GetKeyByHash(hash string) (*key.ResponseKey, error) { k.SettingId = settingId.String - if len(data) != 0 { - pathConfigs := []key.PathConfig{} - if err := json.Unmarshal(data, &pathConfigs); err != nil { - return nil, err - } - - k.AllowedPaths = pathConfigs + if err := hydrateKeyJSONFields(&k, data, extendedBudgetLimitData); err != nil { + return nil, err } return &k, nil @@ -436,6 +469,7 @@ func (s *Store) GetKey(keyId string) (*key.ResponseKey, error) { var k key.ResponseKey var settingId sql.NullString var data []byte + var extendedBudgetLimitData []byte if err := rows.Scan( &k.Name, @@ -462,6 +496,7 @@ func (s *Store) GetKey(keyId string) (*key.ResponseKey, error) { &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &extendedBudgetLimitData, ); err != nil { return nil, err } @@ -469,13 +504,8 @@ func (s *Store) GetKey(keyId string) (*key.ResponseKey, error) { pk := &k pk.SettingId = settingId.String - if len(data) != 0 { - pathConfigs := []key.PathConfig{} - if err := json.Unmarshal(data, &pathConfigs); err != nil { - return nil, err - } - - pk.AllowedPaths = pathConfigs + if err := hydrateKeyJSONFields(pk, data, extendedBudgetLimitData); err != nil { + return nil, err } keys = append(keys, pk) @@ -534,6 +564,7 @@ func (s *Store) GetSpentKeys(tags []string, order string, limit, offset int, val var k key.ResponseKey var settingId sql.NullString var data []byte + var extendedBudgetLimitData []byte if err := rows.Scan( &k.Name, &k.CreatedAt, @@ -559,19 +590,15 @@ func (s *Store) GetSpentKeys(tags []string, order string, limit, offset int, val &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &extendedBudgetLimitData, ); err != nil { return nil, err } pk := &k pk.SettingId = settingId.String - if len(data) != 0 { - pathConfigs := []key.PathConfig{} - if err := json.Unmarshal(data, &pathConfigs); err != nil { - return nil, err - } - - pk.AllowedPaths = pathConfigs + if err := hydrateKeyJSONFields(pk, data, extendedBudgetLimitData); err != nil { + return nil, err } if !validator(pk) { @@ -600,6 +627,7 @@ func (s *Store) GetAllKeys() ([]*key.ResponseKey, error) { var k key.ResponseKey var settingId sql.NullString var data []byte + var extendedBudgetLimitData []byte if err := rows.Scan( &k.Name, &k.CreatedAt, @@ -625,19 +653,15 @@ func (s *Store) GetAllKeys() ([]*key.ResponseKey, error) { &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &extendedBudgetLimitData, ); err != nil { return nil, err } pk := &k pk.SettingId = settingId.String - if len(data) != 0 { - pathConfigs := []key.PathConfig{} - if err := json.Unmarshal(data, &pathConfigs); err != nil { - return nil, err - } - - pk.AllowedPaths = pathConfigs + if err := hydrateKeyJSONFields(pk, data, extendedBudgetLimitData); err != nil { + return nil, err } keys = append(keys, pk) @@ -661,6 +685,7 @@ func (s *Store) GetUpdatedKeys(updatedAt int64) ([]*key.ResponseKey, error) { var k key.ResponseKey var settingId sql.NullString var data []byte + var extendedBudgetLimitData []byte if err := rows.Scan( &k.Name, &k.CreatedAt, @@ -686,19 +711,15 @@ func (s *Store) GetUpdatedKeys(updatedAt int64) ([]*key.ResponseKey, error) { &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &extendedBudgetLimitData, ); err != nil { return nil, err } pk := &k pk.SettingId = settingId.String - if len(data) != 0 { - pathConfigs := []key.PathConfig{} - if err := json.Unmarshal(data, &pathConfigs); err != nil { - return nil, err - } - - pk.AllowedPaths = pathConfigs + if err := hydrateKeyJSONFields(pk, data, extendedBudgetLimitData); err != nil { + return nil, err } keys = append(keys, pk) @@ -792,6 +813,17 @@ func (s *Store) UpdateKey(id string, uk *key.UpdateKey) (*key.ResponseKey, error counter++ } + if uk.ExtendedBudgetLimit != nil { + data, err := marshalExtendedBudgetLimit(uk.ExtendedBudgetLimit) + if err != nil { + return nil, err + } + + values = append(values, data) + fields = append(fields, fmt.Sprintf("extended_budget_limit = $%d", counter)) + counter++ + } + if uk.ShouldLogRequest != nil { values = append(values, *uk.ShouldLogRequest) fields = append(fields, fmt.Sprintf("should_log_request = $%d", counter)) @@ -840,6 +872,7 @@ func (s *Store) UpdateKey(id string, uk *key.UpdateKey) (*key.ResponseKey, error var k key.ResponseKey var settingId sql.NullString var data []byte + var extendedBudgetLimitData []byte if err := s.db.QueryRowContext(ctxTimeout, query, values...).Scan( &k.Name, &k.CreatedAt, @@ -865,6 +898,7 @@ func (s *Store) UpdateKey(id string, uk *key.UpdateKey) (*key.ResponseKey, error &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &extendedBudgetLimitData, ); err != nil { if err == sql.ErrNoRows { return nil, internal_errors.NewNotFoundError(fmt.Sprintf("key not found for id: %s", id)) @@ -875,13 +909,8 @@ func (s *Store) UpdateKey(id string, uk *key.UpdateKey) (*key.ResponseKey, error pk := &k pk.SettingId = settingId.String - if len(data) != 0 { - pathConfigs := []key.PathConfig{} - if err := json.Unmarshal(data, &pathConfigs); err != nil { - return nil, err - } - - pk.AllowedPaths = pathConfigs + if err := hydrateKeyJSONFields(pk, data, extendedBudgetLimitData); err != nil { + return nil, err } return pk, nil @@ -889,8 +918,8 @@ func (s *Store) UpdateKey(id string, uk *key.UpdateKey) (*key.ResponseKey, error func (s *Store) CreateKey(rk *key.RequestKey) (*key.ResponseKey, error) { query := ` - INSERT INTO keys (name, created_at, updated_at, tags, revoked, key_id, key, revoked_reason, cost_limit_in_usd, cost_limit_in_usd_over_time, cost_limit_in_usd_unit, rate_limit_over_time, rate_limit_unit, ttl, key_ring, setting_id, allowed_paths, setting_ids, should_log_request, should_log_response, rotation_enabled, policy_id, is_key_not_hashed, requests_limit) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24) + INSERT INTO keys (name, created_at, updated_at, tags, revoked, key_id, key, revoked_reason, cost_limit_in_usd, cost_limit_in_usd_over_time, cost_limit_in_usd_unit, rate_limit_over_time, rate_limit_unit, ttl, key_ring, setting_id, allowed_paths, setting_ids, should_log_request, should_log_response, rotation_enabled, policy_id, is_key_not_hashed, requests_limit, extended_budget_limit) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19, $20, $21, $22, $23, $24, $25) RETURNING *; ` @@ -899,6 +928,11 @@ func (s *Store) CreateKey(rk *key.RequestKey) (*key.ResponseKey, error) { return nil, err } + extendedBudgetLimitValue, err := marshalExtendedBudgetLimit(rk.ExtendedBudgetLimit) + if err != nil { + return nil, err + } + values := []any{ rk.Name, rk.CreatedAt, @@ -924,6 +958,7 @@ func (s *Store) CreateKey(rk *key.RequestKey) (*key.ResponseKey, error) { rk.PolicyId, rk.IsKeyNotHashed, rk.RequestsLimit, + extendedBudgetLimitValue, } ctxTimeout, cancel := context.WithTimeout(context.Background(), s.wt) @@ -933,6 +968,7 @@ func (s *Store) CreateKey(rk *key.RequestKey) (*key.ResponseKey, error) { var settingId sql.NullString var data []byte + var extendedBudgetLimitData []byte if err := s.db.QueryRowContext(ctxTimeout, query, values...).Scan( &k.Name, &k.CreatedAt, @@ -958,6 +994,7 @@ func (s *Store) CreateKey(rk *key.RequestKey) (*key.ResponseKey, error) { &k.PolicyId, &k.IsKeyNotHashed, &k.RequestsLimit, + &extendedBudgetLimitData, ); err != nil { return nil, err } @@ -965,13 +1002,8 @@ func (s *Store) CreateKey(rk *key.RequestKey) (*key.ResponseKey, error) { pk := &k pk.SettingId = settingId.String - if len(data) != 0 { - pathConfigs := []key.PathConfig{} - if err := json.Unmarshal(data, &pathConfigs); err != nil { - return nil, err - } - - pk.AllowedPaths = pathConfigs + if err := hydrateKeyJSONFields(pk, data, extendedBudgetLimitData); err != nil { + return nil, err } return pk, nil diff --git a/internal/storage/postgresql/key_test.go b/internal/storage/postgresql/key_test.go new file mode 100644 index 00000000..e3717101 --- /dev/null +++ b/internal/storage/postgresql/key_test.go @@ -0,0 +1,71 @@ +package postgresql + +import ( + "reflect" + "testing" + + "github.com/bricks-cloud/bricksllm/internal/key" +) + +func TestMarshalExtendedBudgetLimitNil(t *testing.T) { + data, err := marshalExtendedBudgetLimit(nil) + if err != nil { + t.Fatalf("marshalExtendedBudgetLimit(nil) returned error: %v", err) + } + + if data != nil { + t.Fatalf("marshalExtendedBudgetLimit(nil) = %q, want nil", data) + } +} + +func TestExtendedBudgetLimitJSONRoundTrip(t *testing.T) { + limit := &key.ExtendedBudgetLimit{ + Items: []key.ExtendedBudgetLimitItem{ + {LimitInUsdOverTime: 1.5, Unit: key.HourTimeUnit}, + {LimitInUsdOverTime: 2.5, Unit: key.DayTimeUnit}, + }, + } + + data, err := marshalExtendedBudgetLimit(limit) + if err != nil { + t.Fatalf("marshalExtendedBudgetLimit returned error: %v", err) + } + + roundTripped, err := unmarshalExtendedBudgetLimit(data) + if err != nil { + t.Fatalf("unmarshalExtendedBudgetLimit returned error: %v", err) + } + + if !reflect.DeepEqual(roundTripped, limit) { + t.Fatalf("round trip mismatch: got %#v, want %#v", roundTripped, limit) + } +} + +func TestUnmarshalExtendedBudgetLimitNilHandling(t *testing.T) { + tests := []struct { + name string + data []byte + }{ + {name: "sql null", data: nil}, + {name: "json null", data: []byte("null")}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + limit, err := unmarshalExtendedBudgetLimit(tt.data) + if err != nil { + t.Fatalf("unmarshalExtendedBudgetLimit returned error: %v", err) + } + + if limit != nil { + t.Fatalf("unmarshalExtendedBudgetLimit(%q) = %#v, want nil", tt.data, limit) + } + }) + } +} + +func TestUnmarshalExtendedBudgetLimitInvalidJSON(t *testing.T) { + if _, err := unmarshalExtendedBudgetLimit([]byte("{")); err == nil { + t.Fatal("unmarshalExtendedBudgetLimit should fail for invalid JSON") + } +} diff --git a/internal/storage/redis/cache.go b/internal/storage/redis/cache.go index a60f13b8..cd91a363 100644 --- a/internal/storage/redis/cache.go +++ b/internal/storage/redis/cache.go @@ -111,6 +111,8 @@ func getCounterTtl(rateLimitUnit key.TimeUnit) (time.Time, error) { return now.Truncate(60 * time.Minute).Add(time.Minute * 60).Add(-time.Millisecond), nil case key.DayTimeUnit: return now.Truncate(24 * time.Hour).Add(time.Hour * 24).Add(-time.Millisecond), nil + case key.WeekTimeUnit: + return now.Truncate(7 * 24 * time.Hour).Add(7 * 24 * time.Hour).Add(-time.Millisecond), nil case key.MonthTimeUnit: firstDayOfNextMonth := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC) return firstDayOfNextMonth.Add(-time.Millisecond), nil @@ -130,6 +132,8 @@ func getCounterTimeStamp(rateLimitUnit key.TimeUnit) (int64, error) { return int64(now.Minute()), nil case key.DayTimeUnit: return int64(now.Hour()), nil + case key.WeekTimeUnit: + return int64(now.Day()), nil case key.MonthTimeUnit: return int64(now.Day()), nil } diff --git a/internal/user/user.go b/internal/user/user.go index 0604c651..907bf1d1 100644 --- a/internal/user/user.go +++ b/internal/user/user.go @@ -2,6 +2,7 @@ package user import ( "fmt" + "slices" "strings" "time" @@ -113,7 +114,7 @@ func (u *User) Validate() error { return internal_errors.NewValidationError("rate limit unit can not be empty if rate limit over time is specified") } - if u.RateLimitUnit != key.HourTimeUnit && u.RateLimitUnit != key.MinuteTimeUnit && u.RateLimitUnit != key.SecondTimeUnit && u.RateLimitUnit != key.DayTimeUnit { + if !slices.Contains(key.AllowedTimeUnits, u.RateLimitUnit) { return internal_errors.NewValidationError("rate limit unit can not be identified") } } @@ -123,7 +124,7 @@ func (u *User) Validate() error { return internal_errors.NewValidationError("cost limit unit can not be empty if cost limit over time is specified") } - if u.CostLimitInUsdUnit != key.DayTimeUnit && u.CostLimitInUsdUnit != key.HourTimeUnit && u.CostLimitInUsdUnit != key.MonthTimeUnit && u.CostLimitInUsdUnit != key.MinuteTimeUnit { + if !slices.Contains(key.AllowedTimeUnits, u.CostLimitInUsdUnit) { return internal_errors.NewValidationError("cost limit unit can not be identified") } } @@ -228,7 +229,7 @@ func (uu *UpdateUser) Validate() error { return internal_errors.NewValidationError("rate limit unit can not be empty if rate limit over time is specified") } - if *uu.RateLimitOverTime != 0 && *uu.RateLimitUnit != key.HourTimeUnit && *uu.RateLimitUnit != key.MinuteTimeUnit && *uu.RateLimitUnit != key.SecondTimeUnit && *uu.RateLimitUnit != key.DayTimeUnit { + if *uu.RateLimitOverTime != 0 && !slices.Contains(key.AllowedTimeUnits, *uu.RateLimitUnit) { return internal_errors.NewValidationError("rate limit unit can not be identified") } } @@ -246,7 +247,7 @@ func (uu *UpdateUser) Validate() error { return internal_errors.NewValidationError("cost limit unit can not be empty if cost limit over time is specified") } - if *uu.CostLimitInUsdOverTime != 0 && *uu.CostLimitInUsdUnit != key.DayTimeUnit && *uu.CostLimitInUsdUnit != key.HourTimeUnit && *uu.CostLimitInUsdUnit != key.MonthTimeUnit && *uu.CostLimitInUsdUnit != key.MinuteTimeUnit { + if *uu.CostLimitInUsdOverTime != 0 && !slices.Contains(key.AllowedTimeUnits, *uu.CostLimitInUsdUnit) { return internal_errors.NewValidationError("cost limit unit can not be identified") } } diff --git a/internal/validator/validator.go b/internal/validator/validator.go index c3df819a..a1913627 100644 --- a/internal/validator/validator.go +++ b/internal/validator/validator.go @@ -75,6 +75,11 @@ func (v *Validator) Validate(k *key.ResponseKey, promptCost float64) error { return err } + err = v.validateExtendedCostLimitOverTime(k.KeyId, k.ExtendedBudgetLimit) + if err != nil { + return err + } + err = v.validateCostLimit(k.KeyId, k.CostLimitInUsd) if err != nil { return err @@ -128,6 +133,29 @@ func (v *Validator) validateCostLimitOverTime(keyId string, costLimitOverTime fl return nil } +func (v *Validator) validateExtendedCostLimitOverTime(keyId string, limits *key.ExtendedBudgetLimit) error { + if limits == nil { + return nil + } + if limits.Items == nil || len(limits.Items) == 0 { + return nil + } + + for _, item := range limits.Items { + if item.LimitInUsdOverTime == 0 { + continue + } + err := v.validateCostLimitOverTime(item.ExtendedKey(keyId), item.LimitInUsdOverTime, item.Unit) + if err != nil { + if _, ok := err.(*internal_errors.CostLimitError); ok { + return internal_errors.NewCostLimitError(err.Error(), string(item.Unit)) + } + return err + } + } + return nil +} + func convertDollarToMicroDollars(dollar float64) int64 { return int64(dollar * 1000000) } diff --git a/internal/validator/validator_test.go b/internal/validator/validator_test.go new file mode 100644 index 00000000..39d6f644 --- /dev/null +++ b/internal/validator/validator_test.go @@ -0,0 +1,46 @@ +package validator + +import ( + "testing" + + internal_errors "github.com/bricks-cloud/bricksllm/internal/errors" + "github.com/bricks-cloud/bricksllm/internal/key" +) + +type stubCostLimitCache struct { + counters map[string]int64 +} + +func (s *stubCostLimitCache) GetCounter(keyId string, rateLimitUnit key.TimeUnit) (int64, error) { + return s.counters[keyId], nil +} + +func TestValidateExtendedCostLimitOverTimePreservesBreachedUnit(t *testing.T) { + v := NewValidator( + &stubCostLimitCache{counters: map[string]int64{ + "test-key-w": convertDollarToMicroDollars(1), + }}, + nil, + nil, + nil, + ) + + err := v.validateExtendedCostLimitOverTime("test-key", &key.ExtendedBudgetLimit{ + Items: []key.ExtendedBudgetLimitItem{{ + LimitInUsdOverTime: 1, + Unit: key.WeekTimeUnit, + }}, + }) + if err == nil { + t.Fatal("expected cost limit error, got nil") + } + + cle, ok := err.(*internal_errors.CostLimitError) + if !ok { + t.Fatalf("expected CostLimitError, got %T", err) + } + + if cle.TimeUnit() != string(key.WeekTimeUnit) { + t.Fatalf("expected time unit %q, got %q", key.WeekTimeUnit, cle.TimeUnit()) + } +}