feat: vk provider routing added

Author: Pratham-Mishra04

Makefile

@@ -3,7 +3,7 @@
 # Variables
 HOST ?= localhost
 PORT ?= 8080
-APP_DIR ?=
+APP_DIR ?= 
 PROMETHEUS_LABELS ?=
 LOG_STYLE ?= json
 LOG_LEVEL ?= info

core/bifrost.go

@@ -1583,7 +1583,7 @@ func (bifrost *Bifrost) selectKeyFromProviderForModel(ctx *context.Context, prov
 	}
 
 	if len(keys) == 0 {
-		return schemas.Key{}, fmt.Errorf("no keys found for provider: %v", providerKey)
+		return schemas.Key{}, fmt.Errorf("no keys found for provider: %v and model: %s", providerKey, model)
 	}
 
 	// filter out keys which dont support the model, if the key has no models, it is supported for all models

core/providers/bedrock.go

@@ -88,7 +88,7 @@ type BedrockMistralContent struct {
 type BedrockMistralChatMessage struct {
 	Role       schemas.ModelChatMessageRole `json:"role"`                   // Role of the message sender
 	Content    []BedrockMistralContent      `json:"content"`                // Array of message content
-	ToolCalls  []BedrockAnthropicToolCall  `json:"tool_calls,omitempty"`   // Optional tool calls
+	ToolCalls  []BedrockAnthropicToolCall   `json:"tool_calls,omitempty"`   // Optional tool calls
 	ToolCallID *string                      `json:"tool_call_id,omitempty"` // Optional tool call ID
 }
 

core/schemas/bifrost.go

@@ -118,7 +118,7 @@ const (
 // a text completion, a chat completion, an embedding request, a speech request, or a transcription request.
 type RequestInput struct {
 	TextCompletionInput *string             `json:"text_completion_input,omitempty"`
-	ChatCompletionInput []BifrostMessage   `json:"chat_completion_input,omitempty"`
+	ChatCompletionInput []BifrostMessage    `json:"chat_completion_input,omitempty"`
 	EmbeddingInput      *EmbeddingInput     `json:"embedding_input,omitempty"`
 	SpeechInput         *SpeechInput        `json:"speech_input,omitempty"`
 	TranscriptionInput  *TranscriptionInput `json:"transcription_input,omitempty"`
@@ -295,12 +295,12 @@ type Fallback struct {
 // mapped to the provider's parameters.
 type ModelParameters struct {
 	ToolChoice        *ToolChoice `json:"tool_choice,omitempty"`         // Whether to call a tool
-	Tools             []Tool     `json:"tools,omitempty"`               // Tools to use
+	Tools             []Tool      `json:"tools,omitempty"`               // Tools to use
 	Temperature       *float64    `json:"temperature,omitempty"`         // Controls randomness in the output
 	TopP              *float64    `json:"top_p,omitempty"`               // Controls diversity via nucleus sampling
 	TopK              *int        `json:"top_k,omitempty"`               // Controls diversity via top-k sampling
 	MaxTokens         *int        `json:"max_tokens,omitempty"`          // Maximum number of tokens to generate
-	StopSequences     []string   `json:"stop_sequences,omitempty"`      // Sequences that stop generation
+	StopSequences     []string    `json:"stop_sequences,omitempty"`      // Sequences that stop generation
 	PresencePenalty   *float64    `json:"presence_penalty,omitempty"`    // Penalizes repeated tokens
 	FrequencyPenalty  *float64    `json:"frequency_penalty,omitempty"`   // Penalizes frequent tokens
 	ParallelToolCalls *bool       `json:"parallel_tool_calls,omitempty"` // Enables parallel tool calls
@@ -318,7 +318,7 @@ type FunctionParameters struct {
 	Description *string                `json:"description,omitempty"` // Description of the parameters
 	Required    []string               `json:"required,omitempty"`    // Required parameter names
 	Properties  map[string]interface{} `json:"properties,omitempty"`  // Parameter properties
-	Enum        []string              `json:"enum,omitempty"`        // Enum values for the parameters
+	Enum        []string               `json:"enum,omitempty"`        // Enum values for the parameters
 }
 
 // Function represents a function that can be called by the model.
@@ -492,7 +492,7 @@ type ToolMessage struct {
 type AssistantMessage struct {
 	Refusal     *string      `json:"refusal,omitempty"`
 	Annotations []Annotation `json:"annotations,omitempty"`
-	ToolCalls   []ToolCall  `json:"tool_calls,omitempty"`
+	ToolCalls   []ToolCall   `json:"tool_calls,omitempty"`
 	Thought     *string      `json:"thought,omitempty"`
 }
 
@@ -795,7 +795,7 @@ type BifrostResponseExtraFields struct {
 	Provider    ModelProvider      `json:"provider"`
 	Params      ModelParameters    `json:"model_params"`
 	Latency     *int64             `json:"latency,omitempty"`
-	ChatHistory []BifrostMessage  `json:"chat_history,omitempty"`
+	ChatHistory []BifrostMessage   `json:"chat_history,omitempty"`
 	BilledUsage *BilledLLMUsage    `json:"billed_usage,omitempty"`
 	ChunkIndex  int                `json:"chunk_index"` // used for streaming responses to identify the chunk index, will be 0 for non-streaming responses
 	RawResponse interface{}        `json:"raw_response,omitempty"`

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(ctx, db); err != nil {
+		return err
+	}
 	return nil
 }
 
@@ -247,3 +250,35 @@ func migrationAddCustomProviderConfigJSONColumn(ctx context.Context, db *gorm.DB
 	}
 	return nil
 }
+
+func migrationAddVirtualKeyProviderConfigTable(ctx context.Context, db *gorm.DB) error {
+	m := migration.New(db, migration.DefaultOptions, []*migration.Migration{{
+		ID: "addvirtualkeyproviderconfig",
+		Migrate: func(tx *gorm.DB) error {
+			tx = tx.WithContext(ctx)
+			migrator := tx.Migrator()
+
+			if !migrator.HasTable(&TableVirtualKeyProviderConfig{}) {
+				if err := migrator.CreateTable(&TableVirtualKeyProviderConfig{}); err != nil {
+					return err
+				}
+			}
+
+			return nil
+		},
+		Rollback: func(tx *gorm.DB) error {
+			tx = tx.WithContext(ctx)
+			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 {
@@ -869,6 +871,21 @@ func (s *SQLiteConfigStore) GetVirtualKey(ctx context.Context, id string) (*Tabl
 	return &virtualKey, nil
 }
 
+func (s *SQLiteConfigStore) GetVirtualKeyByValue(ctx context.Context, value string) (*TableVirtualKey, error) {
+	var virtualKey TableVirtualKey
+	if err := s.db.WithContext(ctx).Preload("Team").
+		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, "value = ?", value).Error; err != nil {
+		return nil, err
+	}
+	return &virtualKey, nil
+}
+
 func (s *SQLiteConfigStore) CreateVirtualKey(ctx context.Context, virtualKey *TableVirtualKey, tx ...*gorm.DB) error {
 	var txDB *gorm.DB
 	if len(tx) > 0 {
@@ -923,6 +940,53 @@ func (s *SQLiteConfigStore) UpdateVirtualKey(ctx context.Context, virtualKey *Ta
 	return nil
 }
 
+func (s *SQLiteConfigStore) GetVirtualKeyProviderConfigs(ctx context.Context, virtualKeyID string) ([]TableVirtualKeyProviderConfig, error) {
+	var virtualKey TableVirtualKey
+	if err := s.db.WithContext(ctx).First(&virtualKey, "id = ?", virtualKeyID).Error; err != nil {
+		return nil, err
+	}
+
+	if virtualKey.ID == "" {
+		return nil, nil
+	}
+
+	var providerConfigs []TableVirtualKeyProviderConfig
+	if err := s.db.WithContext(ctx).Where("virtual_key_id = ?", virtualKey.ID).Find(&providerConfigs).Error; err != nil {
+		return nil, err
+	}
+	return providerConfigs, nil
+}
+
+func (s *SQLiteConfigStore) CreateVirtualKeyProviderConfig(ctx context.Context, virtualKeyProviderConfig *TableVirtualKeyProviderConfig, tx ...*gorm.DB) error {
+	var txDB *gorm.DB
+	if len(tx) > 0 {
+		txDB = tx[0]
+	} else {
+		txDB = s.db
+	}
+	return txDB.WithContext(ctx).Create(virtualKeyProviderConfig).Error
+}
+
+func (s *SQLiteConfigStore) UpdateVirtualKeyProviderConfig(ctx context.Context, virtualKeyProviderConfig *TableVirtualKeyProviderConfig, tx ...*gorm.DB) error {
+	var txDB *gorm.DB
+	if len(tx) > 0 {
+		txDB = tx[0]
+	} else {
+		txDB = s.db
+	}
+	return txDB.WithContext(ctx).Save(virtualKeyProviderConfig).Error
+}
+
+func (s *SQLiteConfigStore) DeleteVirtualKeyProviderConfig(ctx context.Context, id uint, tx ...*gorm.DB) error {
+	var txDB *gorm.DB
+	if len(tx) > 0 {
+		txDB = tx[0]
+	} else {
+		txDB = s.db
+	}
+	return txDB.WithContext(ctx).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 {
@@ -1246,10 +1310,10 @@ func (s *SQLiteConfigStore) removeDuplicateKeysAndNullKeys(ctx context.Context)
 // Close closes the SQLite config store.
 func (s *SQLiteConfigStore) Close(ctx context.Context) error {
 	sqlDB, err := s.db.DB()
-    if err != nil {
-        return err
-    }
-    return sqlDB.Close()
+	if err != nil {
+		return err
+	}
+	return sqlDB.Close()
 }
 
 // newSqliteConfigStore creates a new SQLite config store.

framework/configstore/store.go

@@ -55,10 +55,17 @@ type ConfigStore interface {
 	// Governance config CRUD
 	GetVirtualKeys(ctx context.Context) ([]TableVirtualKey, error)
 	GetVirtualKey(ctx context.Context, id string) (*TableVirtualKey, error)
+	GetVirtualKeyByValue(ctx context.Context, value string) (*TableVirtualKey, error)
 	CreateVirtualKey(ctx context.Context, virtualKey *TableVirtualKey, tx ...*gorm.DB) error
 	UpdateVirtualKey(ctx context.Context, virtualKey *TableVirtualKey, tx ...*gorm.DB) error
 	DeleteVirtualKey(ctx context.Context, id string) error
 
+	// Virtual key provider config CRUD
+	GetVirtualKeyProviderConfigs(ctx context.Context, virtualKeyID string) ([]TableVirtualKeyProviderConfig, error)
+	CreateVirtualKeyProviderConfig(ctx context.Context, virtualKeyProviderConfig *TableVirtualKeyProviderConfig, tx ...*gorm.DB) error
+	UpdateVirtualKeyProviderConfig(ctx context.Context, virtualKeyProviderConfig *TableVirtualKeyProviderConfig, tx ...*gorm.DB) error
+	DeleteVirtualKeyProviderConfig(ctx context.Context, id uint, tx ...*gorm.DB) 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,12 @@ 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"`
+	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 +640,15 @@ 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"`
+	AllowedModels []string `gorm:"type:text;serializer:json" json:"allowed_models"` // Empty means all models allowed
+}
+
 // 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/logstore/tables.go

@@ -102,10 +102,10 @@ type Log struct {
 	// Virtual fields for JSON output - these will be populated when needed
 	InputHistoryParsed        []schemas.BifrostMessage    `gorm:"-" json:"input_history,omitempty"`
 	OutputMessageParsed       *schemas.BifrostMessage     `gorm:"-" json:"output_message,omitempty"`
-	EmbeddingOutputParsed     []schemas.BifrostEmbedding `gorm:"-" json:"embedding_output,omitempty"`
+	EmbeddingOutputParsed     []schemas.BifrostEmbedding  `gorm:"-" json:"embedding_output,omitempty"`
 	ParamsParsed              *schemas.ModelParameters    `gorm:"-" json:"params,omitempty"`
-	ToolsParsed               []schemas.Tool             `gorm:"-" json:"tools,omitempty"`
-	ToolCallsParsed           []schemas.ToolCall         `gorm:"-" json:"tool_calls,omitempty"`
+	ToolsParsed               []schemas.Tool              `gorm:"-" json:"tools,omitempty"`
+	ToolCallsParsed           []schemas.ToolCall          `gorm:"-" json:"tool_calls,omitempty"`
 	TokenUsageParsed          *schemas.LLMUsage           `gorm:"-" json:"token_usage,omitempty"`
 	ErrorDetailsParsed        *schemas.BifrostError       `gorm:"-" json:"error_details,omitempty"`
 	SpeechInputParsed         *schemas.SpeechInput        `gorm:"-" json:"speech_input,omitempty"`