Skip to content

Commit 79247f6

Browse files
danpithsdanpiths
committed
feat: add support for per-model and per-provider level budgeting and rate limiting in governance plugin
1 parent f6b85fa commit 79247f6

File tree

10 files changed

+1145
-62
lines changed

10 files changed

+1145
-62
lines changed

‎framework/configstore/clientconfig.go‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -718,10 +718,12 @@ type AuthConfig struct {
718718
type ConfigMap map[schemas.ModelProvider]ProviderConfig
719719

720720
type GovernanceConfig struct {
721-
VirtualKeys []tables.TableVirtualKey `json:"virtual_keys"`
722-
Teams []tables.TableTeam `json:"teams"`
723-
Customers []tables.TableCustomer `json:"customers"`
724-
Budgets []tables.TableBudget `json:"budgets"`
725-
RateLimits []tables.TableRateLimit `json:"rate_limits"`
726-
AuthConfig *AuthConfig `json:"auth_config,omitempty"`
721+
VirtualKeys []tables.TableVirtualKey `json:"virtual_keys"`
722+
Teams []tables.TableTeam `json:"teams"`
723+
Customers []tables.TableCustomer `json:"customers"`
724+
Budgets []tables.TableBudget `json:"budgets"`
725+
RateLimits []tables.TableRateLimit `json:"rate_limits"`
726+
ModelConfigs []tables.TableModelConfig `json:"model_configs"`
727+
Providers []tables.TableProvider `json:"providers"`
728+
AuthConfig *AuthConfig `json:"auth_config,omitempty"`
727729
}

‎framework/configstore/migrations.go‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,12 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
125125
if err := migrationAddUseForBatchAPIColumnAndS3BucketsConfig(ctx, db); err != nil {
126126
return err
127127
}
128+
if err := migrationAddModelConfigTable(ctx, db); err != nil {
129+
return err
130+
}
131+
if err := migrationAddProviderGovernanceColumns(ctx, db); err != nil {
132+
return err
133+
}
128134
return nil
129135
}
130136

@@ -2080,3 +2086,112 @@ func migrationAddUseForBatchAPIColumnAndS3BucketsConfig(ctx context.Context, db
20802086
}
20812087
return nil
20822088
}
2089+
2090+
// migrationAddModelConfigTable adds the governance_model_configs table
2091+
func migrationAddModelConfigTable(ctx context.Context, db *gorm.DB) error {
2092+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
2093+
ID: "add_model_config_table",
2094+
Migrate: func(tx *gorm.DB) error {
2095+
tx = tx.WithContext(ctx)
2096+
migrator := tx.Migrator()
2097+
if !migrator.HasTable(&tables.TableModelConfig{}) {
2098+
if err := migrator.CreateTable(&tables.TableModelConfig{}); err != nil {
2099+
return err
2100+
}
2101+
}
2102+
return nil
2103+
},
2104+
Rollback: func(tx *gorm.DB) error {
2105+
tx = tx.WithContext(ctx)
2106+
migrator := tx.Migrator()
2107+
if err := migrator.DropTable(&tables.TableModelConfig{}); err != nil {
2108+
return err
2109+
}
2110+
return nil
2111+
},
2112+
}})
2113+
err := m.Migrate()
2114+
if err != nil {
2115+
return fmt.Errorf("error while running add model config table migration: %s", err.Error())
2116+
}
2117+
return nil
2118+
}
2119+
2120+
// migrationAddProviderGovernanceColumns adds budget_id and rate_limit_id columns to config_providers table
2121+
func migrationAddProviderGovernanceColumns(ctx context.Context, db *gorm.DB) error {
2122+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
2123+
ID: "add_provider_governance_columns",
2124+
Migrate: func(tx *gorm.DB) error {
2125+
tx = tx.WithContext(ctx)
2126+
migrator := tx.Migrator()
2127+
provider := &tables.TableProvider{}
2128+
2129+
// Add budget_id column if it doesn't exist
2130+
if !migrator.HasColumn(provider, "budget_id") {
2131+
if err := migrator.AddColumn(provider, "BudgetID"); err != nil {
2132+
return fmt.Errorf("failed to add budget_id column: %w", err)
2133+
}
2134+
// Create index for budget_id
2135+
if !migrator.HasIndex(provider, "idx_provider_budget") {
2136+
if err := tx.Exec("CREATE INDEX IF NOT EXISTS idx_provider_budget ON config_providers (budget_id)").Error; err != nil {
2137+
return fmt.Errorf("failed to create budget_id index: %w", err)
2138+
}
2139+
}
2140+
}
2141+
2142+
// Add rate_limit_id column if it doesn't exist
2143+
if !migrator.HasColumn(provider, "rate_limit_id") {
2144+
if err := migrator.AddColumn(provider, "RateLimitID"); err != nil {
2145+
return fmt.Errorf("failed to add rate_limit_id column: %w", err)
2146+
}
2147+
// Create index for rate_limit_id
2148+
if !migrator.HasIndex(provider, "idx_provider_rate_limit") {
2149+
if err := tx.Exec("CREATE INDEX IF NOT EXISTS idx_provider_rate_limit ON config_providers (rate_limit_id)").Error; err != nil {
2150+
return fmt.Errorf("failed to create rate_limit_id index: %w", err)
2151+
}
2152+
}
2153+
}
2154+
2155+
return nil
2156+
},
2157+
Rollback: func(tx *gorm.DB) error {
2158+
tx = tx.WithContext(ctx)
2159+
migrator := tx.Migrator()
2160+
provider := &tables.TableProvider{}
2161+
2162+
// Drop indexes first
2163+
if migrator.HasIndex(provider, "idx_provider_rate_limit") {
2164+
if err := tx.Exec("DROP INDEX IF EXISTS idx_provider_rate_limit").Error; err != nil {
2165+
return fmt.Errorf("failed to drop rate_limit_id index: %w", err)
2166+
}
2167+
}
2168+
2169+
if migrator.HasIndex(provider, "idx_provider_budget") {
2170+
if err := tx.Exec("DROP INDEX IF EXISTS idx_provider_budget").Error; err != nil {
2171+
return fmt.Errorf("failed to drop budget_id index: %w", err)
2172+
}
2173+
}
2174+
2175+
// Drop rate_limit_id column if it exists
2176+
if migrator.HasColumn(provider, "rate_limit_id") {
2177+
if err := migrator.DropColumn(provider, "RateLimitID"); err != nil {
2178+
return fmt.Errorf("failed to drop rate_limit_id column: %w", err)
2179+
}
2180+
}
2181+
2182+
// Drop budget_id column if it exists
2183+
if migrator.HasColumn(provider, "budget_id") {
2184+
if err := migrator.DropColumn(provider, "BudgetID"); err != nil {
2185+
return fmt.Errorf("failed to drop budget_id column: %w", err)
2186+
}
2187+
}
2188+
2189+
return nil
2190+
},
2191+
}})
2192+
err := m.Migrate()
2193+
if err != nil {
2194+
return fmt.Errorf("error while running add provider governance columns migration: %s", err.Error())
2195+
}
2196+
return nil
2197+
}

‎framework/configstore/rdb.go‎

Lines changed: 126 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -698,6 +698,18 @@ func (s *RDBConfigStore) GetProvidersConfig(ctx context.Context) (map[schemas.Mo
698698
return processedProviders, nil
699699
}
700700

701+
// GetProviders retrieves all providers from the database with their governance relationships.
702+
func (s *RDBConfigStore) GetProviders(ctx context.Context) ([]tables.TableProvider, error) {
703+
var providers []tables.TableProvider
704+
if err := s.db.WithContext(ctx).Preload("Budget").Preload("RateLimit").Find(&providers).Error; err != nil {
705+
if errors.Is(err, gorm.ErrRecordNotFound) {
706+
return nil, ErrNotFound
707+
}
708+
return nil, err
709+
}
710+
return providers, nil
711+
}
712+
701713
// GetMCPConfig retrieves the MCP configuration from the database.
702714
func (s *RDBConfigStore) GetMCPConfig(ctx context.Context) (*schemas.MCPConfig, error) {
703715
var dbMCPClients []tables.TableMCPClient
@@ -1915,13 +1927,112 @@ func (s *RDBConfigStore) UpdateBudget(ctx context.Context, budget *tables.TableB
19151927
return nil
19161928
}
19171929

1930+
// GetModelConfigs retrieves all model configs from the database.
1931+
func (s *RDBConfigStore) GetModelConfigs(ctx context.Context) ([]tables.TableModelConfig, error) {
1932+
var modelConfigs []tables.TableModelConfig
1933+
if err := s.db.WithContext(ctx).Preload("Budget").Preload("RateLimit").Find(&modelConfigs).Error; err != nil {
1934+
if errors.Is(err, gorm.ErrRecordNotFound) {
1935+
return nil, ErrNotFound
1936+
}
1937+
return nil, err
1938+
}
1939+
return modelConfigs, nil
1940+
}
1941+
1942+
// GetModelConfig retrieves a specific model config from the database by model name and optional provider.
1943+
func (s *RDBConfigStore) GetModelConfig(ctx context.Context, modelName string, provider *string) (*tables.TableModelConfig, error) {
1944+
var modelConfig tables.TableModelConfig
1945+
query := s.db.WithContext(ctx).Where("model_name = ?", modelName)
1946+
if provider != nil {
1947+
query = query.Where("provider = ?", *provider)
1948+
} else {
1949+
query = query.Where("provider IS NULL")
1950+
}
1951+
if err := query.Preload("Budget").Preload("RateLimit").First(&modelConfig).Error; err != nil {
1952+
if errors.Is(err, gorm.ErrRecordNotFound) {
1953+
return nil, ErrNotFound
1954+
}
1955+
return nil, err
1956+
}
1957+
return &modelConfig, nil
1958+
}
1959+
1960+
// GetModelConfigByID retrieves a specific model config from the database by ID.
1961+
func (s *RDBConfigStore) GetModelConfigByID(ctx context.Context, id string) (*tables.TableModelConfig, error) {
1962+
var modelConfig tables.TableModelConfig
1963+
if err := s.db.WithContext(ctx).Preload("Budget").Preload("RateLimit").First(&modelConfig, "id = ?", id).Error; err != nil {
1964+
if errors.Is(err, gorm.ErrRecordNotFound) {
1965+
return nil, ErrNotFound
1966+
}
1967+
return nil, err
1968+
}
1969+
return &modelConfig, nil
1970+
}
1971+
1972+
// CreateModelConfig creates a new model config in the database.
1973+
func (s *RDBConfigStore) CreateModelConfig(ctx context.Context, modelConfig *tables.TableModelConfig, tx ...*gorm.DB) error {
1974+
var txDB *gorm.DB
1975+
if len(tx) > 0 {
1976+
txDB = tx[0]
1977+
} else {
1978+
txDB = s.db
1979+
}
1980+
if err := txDB.WithContext(ctx).Create(modelConfig).Error; err != nil {
1981+
return s.parseGormError(err)
1982+
}
1983+
return nil
1984+
}
1985+
1986+
// UpdateModelConfig updates a model config in the database.
1987+
func (s *RDBConfigStore) UpdateModelConfig(ctx context.Context, modelConfig *tables.TableModelConfig, tx ...*gorm.DB) error {
1988+
var txDB *gorm.DB
1989+
if len(tx) > 0 {
1990+
txDB = tx[0]
1991+
} else {
1992+
txDB = s.db
1993+
}
1994+
if err := txDB.WithContext(ctx).Save(modelConfig).Error; err != nil {
1995+
return s.parseGormError(err)
1996+
}
1997+
return nil
1998+
}
1999+
2000+
// UpdateModelConfigs updates multiple model configs in the database.
2001+
func (s *RDBConfigStore) UpdateModelConfigs(ctx context.Context, modelConfigs []*tables.TableModelConfig, tx ...*gorm.DB) error {
2002+
var txDB *gorm.DB
2003+
if len(tx) > 0 {
2004+
txDB = tx[0]
2005+
} else {
2006+
txDB = s.db
2007+
}
2008+
for _, mc := range modelConfigs {
2009+
if err := txDB.WithContext(ctx).Save(mc).Error; err != nil {
2010+
return s.parseGormError(err)
2011+
}
2012+
}
2013+
return nil
2014+
}
2015+
2016+
// DeleteModelConfig deletes a model config from the database.
2017+
func (s *RDBConfigStore) DeleteModelConfig(ctx context.Context, id string) error {
2018+
if err := s.db.WithContext(ctx).Delete(&tables.TableModelConfig{}, "id = ?", id).Error; err != nil {
2019+
if errors.Is(err, gorm.ErrRecordNotFound) {
2020+
return ErrNotFound
2021+
}
2022+
return s.parseGormError(err)
2023+
}
2024+
return nil
2025+
}
2026+
19182027
// GetGovernanceConfig retrieves the governance configuration from the database.
19192028
func (s *RDBConfigStore) GetGovernanceConfig(ctx context.Context) (*GovernanceConfig, error) {
19202029
var virtualKeys []tables.TableVirtualKey
19212030
var teams []tables.TableTeam
19222031
var customers []tables.TableCustomer
19232032
var budgets []tables.TableBudget
19242033
var rateLimits []tables.TableRateLimit
2034+
var modelConfigs []tables.TableModelConfig
2035+
var providers []tables.TableProvider
19252036
var governanceConfigs []tables.TableGovernanceConfig
19262037

19272038
if err := s.db.WithContext(ctx).Preload("ProviderConfigs").Find(&virtualKeys).Error; err != nil {
@@ -1939,12 +2050,18 @@ func (s *RDBConfigStore) GetGovernanceConfig(ctx context.Context) (*GovernanceCo
19392050
if err := s.db.WithContext(ctx).Find(&rateLimits).Error; err != nil {
19402051
return nil, err
19412052
}
2053+
if err := s.db.WithContext(ctx).Find(&modelConfigs).Error; err != nil {
2054+
return nil, err
2055+
}
2056+
if err := s.db.WithContext(ctx).Find(&providers).Error; err != nil {
2057+
return nil, err
2058+
}
19422059
// Fetching governance config for username and password
19432060
if err := s.db.WithContext(ctx).Find(&governanceConfigs).Error; err != nil {
19442061
return nil, err
19452062
}
19462063
// Check if any config is present
1947-
if len(virtualKeys) == 0 && len(teams) == 0 && len(customers) == 0 && len(budgets) == 0 && len(rateLimits) == 0 && len(governanceConfigs) == 0 {
2064+
if len(virtualKeys) == 0 && len(teams) == 0 && len(customers) == 0 && len(budgets) == 0 && len(rateLimits) == 0 && len(modelConfigs) == 0 && len(providers) == 0 && len(governanceConfigs) == 0 {
19482065
return nil, nil
19492066
}
19502067
var authConfig *AuthConfig
@@ -1972,12 +2089,14 @@ func (s *RDBConfigStore) GetGovernanceConfig(ctx context.Context) (*GovernanceCo
19722089
}
19732090
}
19742091
return &GovernanceConfig{
1975-
VirtualKeys: virtualKeys,
1976-
Teams: teams,
1977-
Customers: customers,
1978-
Budgets: budgets,
1979-
RateLimits: rateLimits,
1980-
AuthConfig: authConfig,
2092+
VirtualKeys: virtualKeys,
2093+
Teams: teams,
2094+
Customers: customers,
2095+
Budgets: budgets,
2096+
RateLimits: rateLimits,
2097+
ModelConfigs: modelConfigs,
2098+
Providers: providers,
2099+
AuthConfig: authConfig,
19812100
}, nil
19822101
}
19832102

‎framework/configstore/store.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ type ConfigStore interface {
3333
UpdateProvider(ctx context.Context, provider schemas.ModelProvider, config ProviderConfig, envKeys map[string][]EnvKeyInfo, tx ...*gorm.DB) error
3434
DeleteProvider(ctx context.Context, provider schemas.ModelProvider, tx ...*gorm.DB) error
3535
GetProvidersConfig(ctx context.Context) (map[schemas.ModelProvider]ProviderConfig, error)
36+
GetProviders(ctx context.Context) ([]tables.TableProvider, error)
3637

3738
// MCP config CRUD
3839
GetMCPConfig(ctx context.Context) (*schemas.MCPConfig, error)
@@ -114,6 +115,15 @@ type ConfigStore interface {
114115
UpdateBudget(ctx context.Context, budget *tables.TableBudget, tx ...*gorm.DB) error
115116
UpdateBudgets(ctx context.Context, budgets []*tables.TableBudget, tx ...*gorm.DB) error
116117

118+
// Model config CRUD
119+
GetModelConfigs(ctx context.Context) ([]tables.TableModelConfig, error)
120+
GetModelConfig(ctx context.Context, modelName string, provider *string) (*tables.TableModelConfig, error)
121+
GetModelConfigByID(ctx context.Context, id string) (*tables.TableModelConfig, error)
122+
CreateModelConfig(ctx context.Context, modelConfig *tables.TableModelConfig, tx ...*gorm.DB) error
123+
UpdateModelConfig(ctx context.Context, modelConfig *tables.TableModelConfig, tx ...*gorm.DB) error
124+
UpdateModelConfigs(ctx context.Context, modelConfigs []*tables.TableModelConfig, tx ...*gorm.DB) error
125+
DeleteModelConfig(ctx context.Context, id string) error
126+
117127
// Governance config CRUD
118128
GetGovernanceConfig(ctx context.Context) (*GovernanceConfig, error)
119129

0 commit comments

Comments
 (0)