feat: vk provider routing added

Author: Pratham-Mishra04

framework/configstore/migrations.go

@@ -19,6 +19,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
 	if err := migrationAddCustomProviderConfigJSONColumn(ctx, db); err != nil {
 		return err
 	}
+	if err := migrationAddVirtualKeyProviderConfigTable(db); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -247,3 +250,33 @@ func migrationAddCustomProviderConfigJSONColumn(ctx context.Context, db *gorm.DB
 	}
 	return nil
 }
+
+func migrationAddVirtualKeyProviderConfigTable(db *gorm.DB) error {
+	m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
+		ID: "addvirtualkeyproviderconfig",
+		Migrate: func(tx *gorm.DB) error {
+			migrator := tx.Migrator()
+
+			if !migrator.HasTable(&TableVirtualKeyProviderConfig{}) {
+				if err := migrator.CreateTable(&TableVirtualKeyProviderConfig{}); err != nil {
+					return err
+				}
+			}
+
+			return nil
+		},
+		Rollback: func(tx *gorm.DB) error {
+			migrator := tx.Migrator()
+
+			if err := migrator.DropTable(&TableVirtualKeyProviderConfig{}); err != nil {
+				return err
+			}
+			return nil
+		},
+	}})
+	err := m.Migrate()
+	if err != nil {
+		return fmt.Errorf("error while running db migration: %s", err.Error())
+	}
+	return nil
+}

framework/configstore/sqlite.go

@@ -845,6 +845,7 @@ func (s *SQLiteConfigStore) GetVirtualKeys(ctx context.Context) ([]TableVirtualK
 		Preload("Customer").
 		Preload("Budget").
 		Preload("RateLimit").
+		Preload("ProviderConfigs").
 		Preload("Keys", func(db *gorm.DB) *gorm.DB {
 			return db.Select("id, key_id, models_json")
 		}).Find(&virtualKeys).Error; err != nil {
@@ -861,6 +862,7 @@ func (s *SQLiteConfigStore) GetVirtualKey(ctx context.Context, id string) (*Tabl
 		Preload("Customer").
 		Preload("Budget").
 		Preload("RateLimit").
+		Preload("ProviderConfigs").
 		Preload("Keys", func(db *gorm.DB) *gorm.DB {
 			return db.Select("id, key_id, models_json")
 		}).First(&virtualKey, "id = ?", id).Error; err != nil {
@@ -923,6 +925,47 @@ func (s *SQLiteConfigStore) UpdateVirtualKey(ctx context.Context, virtualKey *Ta
 	return nil
 }
 
+func (s *SQLiteConfigStore) GetVirtualKeyProviderConfigsFromValue(virtualKeyValue string) ([]TableVirtualKeyProviderConfig, error) {
+	var virtualKey TableVirtualKey
+	if err := s.db.First(&virtualKey, "value = ?", virtualKeyValue).Error; err != nil {
+		return nil, err
+	}
+
+	if virtualKey.ID == "" {
+		return nil, nil
+	}
+
+	var providerConfigs []TableVirtualKeyProviderConfig
+	if err := s.db.Where("virtual_key_id = ?", virtualKey.ID).Find(&providerConfigs).Error; err != nil {
+		return nil, err
+	}
+	return providerConfigs, nil
+}
+
+func (s *SQLiteConfigStore) CreateVirtualKeyProviderConfig(virtualKeyProviderConfig *TableVirtualKeyProviderConfig, tx ...*gorm.DB) error {
+	var txDB *gorm.DB
+	if len(tx) > 0 {
+		txDB = tx[0]
+	} else {
+		txDB = s.db
+	}
+	return txDB.Create(virtualKeyProviderConfig).Error
+}
+
+func (s *SQLiteConfigStore) UpdateVirtualKeyProviderConfig(virtualKeyProviderConfig *TableVirtualKeyProviderConfig, tx ...*gorm.DB) error {
+	var txDB *gorm.DB
+	if len(tx) > 0 {
+		txDB = tx[0]
+	} else {
+		txDB = s.db
+	}
+	return txDB.Save(virtualKeyProviderConfig).Error
+}
+
+func (s *SQLiteConfigStore) DeleteVirtualKeyProviderConfig(id uint) error {
+	return s.db.Delete(&TableVirtualKeyProviderConfig{}, "id = ?", id).Error
+}
+
 // GetKeysByIDs retrieves multiple keys by their IDs
 func (s *SQLiteConfigStore) GetKeysByIDs(ctx context.Context, ids []string) ([]TableKey, error) {
 	if len(ids) == 0 {

framework/configstore/store.go

@@ -59,6 +59,12 @@ type ConfigStore interface {
 	UpdateVirtualKey(ctx context.Context, virtualKey *TableVirtualKey, tx ...*gorm.DB) error
 	DeleteVirtualKey(ctx context.Context, id string) error
 
+	// Virtual key provider config CRUD
+	GetVirtualKeyProviderConfigsFromValue(virtualKeyValue string) ([]TableVirtualKeyProviderConfig, error)
+	CreateVirtualKeyProviderConfig(virtualKeyProviderConfig *TableVirtualKeyProviderConfig, tx ...*gorm.DB) error
+	UpdateVirtualKeyProviderConfig(virtualKeyProviderConfig *TableVirtualKeyProviderConfig, tx ...*gorm.DB) error
+	DeleteVirtualKeyProviderConfig(id uint) error
+
 	// Team CRUD
 	GetTeams(ctx context.Context, customerID string) ([]TableTeam, error)
 	GetTeam(ctx context.Context, id string) (*TableTeam, error)

framework/configstore/tables.go

@@ -616,13 +616,13 @@ type TableTeam struct {
 
 // TableVirtualKey represents a virtual key with budget, rate limits, and team/customer association
 type TableVirtualKey struct {
-	ID               string   `gorm:"primaryKey;type:varchar(255)" json:"id"`
-	Name             string   `gorm:"uniqueIndex:idx_virtual_key_name;type:varchar(255);not null" json:"name"`
-	Description      string   `gorm:"type:text" json:"description,omitempty"`
-	Value            string   `gorm:"uniqueIndex:idx_virtual_key_value;type:varchar(255);not null" json:"value"` // The virtual key value
-	IsActive         bool     `gorm:"default:true" json:"is_active"`
-	AllowedModels    []string `gorm:"type:text;serializer:json" json:"allowed_models"`    // Empty means all models allowed
-	AllowedProviders []string `gorm:"type:text;serializer:json" json:"allowed_providers"` // Empty means all providers allowed
+	ID              string                          `gorm:"primaryKey;type:varchar(255)" json:"id"`
+	Name            string                          `gorm:"uniqueIndex:idx_virtual_key_name;type:varchar(255);not null" json:"name"`
+	Description     string                          `gorm:"type:text" json:"description,omitempty"`
+	Value           string                          `gorm:"uniqueIndex:idx_virtual_key_value;type:varchar(255);not null" json:"value"` // The virtual key value
+	IsActive        bool                            `gorm:"default:true" json:"is_active"`
+	AllowedModels   []string                        `gorm:"type:text;serializer:json" json:"allowed_models"`                             // Empty means all models allowed
+	ProviderConfigs []TableVirtualKeyProviderConfig `gorm:"foreignKey:VirtualKeyID;constraint:OnDelete:CASCADE" json:"provider_configs"` // Empty means all providers allowed
 
 	// Foreign key relationships (mutually exclusive: either TeamID or CustomerID, not both)
 	TeamID      *string    `gorm:"type:varchar(255);index" json:"team_id,omitempty"`
@@ -641,6 +641,14 @@ type TableVirtualKey struct {
 	UpdatedAt time.Time `gorm:"index;not null" json:"updated_at"`
 }
 
+// TableVirtualKeyProviderConfig represents a provider configuration for a virtual key
+type TableVirtualKeyProviderConfig struct {
+	ID           uint    `gorm:"primaryKey;autoIncrement" json:"id"`
+	VirtualKeyID string  `gorm:"type:varchar(255);not null" json:"virtual_key_id"`
+	Provider     string  `gorm:"type:varchar(50);not null" json:"provider"`
+	Weight       float64 `gorm:"default:1.0" json:"weight"`
+}
+
 // TableModelPricing represents pricing information for AI models
 type TableModelPricing struct {
 	ID                 uint    `gorm:"primaryKey;autoIncrement" json:"id"`
@@ -675,11 +683,14 @@ type TableModelPricing struct {
 }
 
 // Table names
-func (TableBudget) TableName() string       { return "governance_budgets" }
-func (TableRateLimit) TableName() string    { return "governance_rate_limits" }
-func (TableCustomer) TableName() string     { return "governance_customers" }
-func (TableTeam) TableName() string         { return "governance_teams" }
-func (TableVirtualKey) TableName() string   { return "governance_virtual_keys" }
+func (TableBudget) TableName() string     { return "governance_budgets" }
+func (TableRateLimit) TableName() string  { return "governance_rate_limits" }
+func (TableCustomer) TableName() string   { return "governance_customers" }
+func (TableTeam) TableName() string       { return "governance_teams" }
+func (TableVirtualKey) TableName() string { return "governance_virtual_keys" }
+func (TableVirtualKeyProviderConfig) TableName() string {
+	return "governance_virtual_key_provider_configs"
+}
 func (TableConfig) TableName() string       { return "governance_config" }
 func (TableModelPricing) TableName() string { return "governance_model_pricing" }
 

framework/pricing/main.go

@@ -3,6 +3,7 @@ package pricing
 import (
 	"context"
 	"fmt"
+	"slices"
 	"sync"
 	"time"
 
@@ -26,6 +27,8 @@ type PricingManager struct {
 	pricingData map[string]configstore.TableModelPricing
 	mu          sync.RWMutex
 
+	modelPool map[schemas.ModelProvider][]string
+
 	// Background sync worker
 	syncTicker *time.Ticker
 	done       chan struct{}
@@ -73,9 +76,12 @@ func Init(ctx context.Context, configStore configstore.ConfigStore, logger schem
 		configStore: configStore,
 		logger:      logger,
 		pricingData: make(map[string]configstore.TableModelPricing),
+		modelPool:   make(map[schemas.ModelProvider][]string),
 		done:        make(chan struct{}),
 	}
 
+	logger.Info("initializing pricing manager...")
+
 	if configStore != nil {
 		// Load initial pricing data
 		if err := pm.loadPricingFromDatabase(ctx); err != nil {
@@ -86,14 +92,16 @@ func Init(ctx context.Context, configStore configstore.ConfigStore, logger schem
 		if err := pm.syncPricing(ctx); err != nil {
 			return nil, fmt.Errorf("failed to sync pricing data: %w", err)
 		}
-		
 	} else {
 		// Load pricing data from config memory
 		if err := pm.loadPricingIntoMemory(); err != nil {
 			return nil, fmt.Errorf("failed to load pricing data from config memory: %w", err)
 		}
 	}
 
+	// Populate model pool with normalized providers
+	pm.populateModelPool()
+
 	// Start background sync worker
 	pm.startSyncWorker(ctx)
 	pm.configStore = configStore
@@ -327,6 +335,78 @@ func (pm *PricingManager) CalculateCostFromUsage(provider string, model string,
 	return totalCost
 }
 
+// populateModelPool populates the model pool with all available models per provider (thread-safe)
+func (pm *PricingManager) populateModelPool() {
+	// Clear existing model pool
+	pm.modelPool = make(map[schemas.ModelProvider][]string)
+
+	// Map to track unique models per provider
+	providerModels := make(map[schemas.ModelProvider]map[string]bool)
+
+	// Iterate through all pricing data to collect models per provider
+	for _, pricing := range pm.pricingData {
+		// Normalize provider before adding to model pool
+		normalizedProvider := schemas.ModelProvider(normalizeProvider(pricing.Provider))
+
+		// Initialize map for this provider if not exists
+		if providerModels[normalizedProvider] == nil {
+			providerModels[normalizedProvider] = make(map[string]bool)
+		}
+
+		// Add model to the provider's model set (using map for deduplication)
+		providerModels[normalizedProvider][pricing.Model] = true
+	}
+
+	// Convert sets to slices
+	for provider, modelSet := range providerModels {
+		models := make([]string, 0, len(modelSet))
+		for model := range modelSet {
+			models = append(models, model)
+		}
+		pm.modelPool[provider] = models
+	}
+
+	// Log the populated model pool for debugging
+	totalModels := 0
+	for provider, models := range pm.modelPool {
+		totalModels += len(models)
+		pm.logger.Debug("populated %d models for provider %s", len(models), string(provider))
+	}
+	pm.logger.Info("populated model pool with %d models across %d providers", totalModels, len(pm.modelPool))
+}
+
+// GetModelsForProvider returns all available models for a given provider (thread-safe)
+func (pm *PricingManager) GetModelsForProvider(provider schemas.ModelProvider) []string {
+	pm.mu.RLock()
+	defer pm.mu.RUnlock()
+
+	models, exists := pm.modelPool[provider]
+	if !exists {
+		return []string{}
+	}
+
+	// Return a copy to prevent external modification
+	result := make([]string, len(models))
+	copy(result, models)
+	return result
+}
+
+// GetProvidersForModel returns all providers for a given model (thread-safe)
+func (pm *PricingManager) GetProvidersForModel(model string) []schemas.ModelProvider {
+	pm.mu.RLock()
+	defer pm.mu.RUnlock()
+
+	providers := make([]schemas.ModelProvider, 0)
+	for provider, models := range pm.modelPool {
+		if slices.Contains(models, model) {
+			providers = append(providers, provider)
+		} else if slices.Contains(models, model) {
+			providers = append(providers, provider)
+		}
+	}
+	return providers
+}
+
 // getPricing returns pricing information for a model (thread-safe)
 func (pm *PricingManager) getPricing(model, provider string, requestType schemas.RequestType) (*configstore.TableModelPricing, bool) {
 	pm.mu.RLock()

framework/pricing/sync.go

@@ -62,8 +62,8 @@ func (pm *PricingManager) syncPricing(ctx context.Context) error {
 	pricingData, err := pm.loadPricingFromURL()
 	if err != nil {
 		// Check if we have existing data in database
-		pricingRecords, err := pm.configStore.GetModelPrices(ctx)
-		if err != nil {
+		pricingRecords, pricingErr := pm.configStore.GetModelPrices(ctx)
+		if pricingErr != nil {
 			return fmt.Errorf("failed to get pricing records: %w", err)
 		}
 		if len(pricingRecords) > 0 {

plugins/governance/resolver.go

@@ -153,11 +153,17 @@ func (r *BudgetResolver) isModelAllowed(vk *configstore.TableVirtualKey, model s
 // isProviderAllowed checks if the requested provider is allowed for this VK
 func (r *BudgetResolver) isProviderAllowed(vk *configstore.TableVirtualKey, provider schemas.ModelProvider) bool {
 	// Empty AllowedProviders means all providers are allowed
-	if len(vk.AllowedProviders) == 0 {
+	if len(vk.ProviderConfigs) == 0 {
 		return true
 	}
 
-	return slices.Contains(vk.AllowedProviders, string(provider))
+	for _, pc := range vk.ProviderConfigs {
+		if pc.Provider == string(provider) {
+			return true
+		}
+	}
+
+	return false
 }
 
 // checkRateLimits checks the VK's rate limits using flexible approach

plugins/governance/tracker.go

@@ -220,7 +220,6 @@ func (t *UsageTracker) PerformStartupResets(ctx context.Context) error {
 			}
 		}
 	}
-	t.logger.Info("startup reset summary: VKs with RL=%d, without RL=%d, RL resets=%d", vksWithRateLimits, vksWithoutRateLimits, len(resetRateLimits))
 	if len(errs) > 0 {
 		t.logger.Error("startup reset encountered %d errors: %v", len(errs), errs)
 		return fmt.Errorf("startup reset completed with %d errors", len(errs))