Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion core/bifrost.go
Original file line number Diff line number Diff line change
Expand Up @@ -6039,7 +6039,10 @@ func clearAnthropicPassthroughForNonNativeProvider(ctx *schemas.BifrostContext,
if integrationType, _ := ctx.Value(schemas.BifrostContextKeyIntegrationType).(string); integrationType != "anthropic" {
return
}
if baseProvider == schemas.Anthropic || baseProvider == schemas.Vertex || baseProvider == schemas.Azure {
if baseProvider == schemas.Anthropic ||
baseProvider == schemas.Vertex ||
baseProvider == schemas.Azure ||
baseProvider == schemas.BedrockMantle {
return
}
ctx.SetValue(schemas.BifrostContextKeyUseRawRequestBody, false)
Expand Down
23 changes: 23 additions & 0 deletions framework/configstore/clientconfig.go
Original file line number Diff line number Diff line change
Expand Up @@ -547,6 +547,29 @@ func (p *ProviderConfig) Redacted() *ProviderConfig {
redactedConfig.Keys[i].BedrockKeyConfig = bedrockConfig
}

// Redact Bedrock Mantle key config if present
if key.BedrockMantleKeyConfig != nil {
mantleConfig := &schemas.BedrockMantleKeyConfig{}
mantleConfig.AccessKey = *key.BedrockMantleKeyConfig.AccessKey.Redacted()
mantleConfig.SecretKey = *key.BedrockMantleKeyConfig.SecretKey.Redacted()
if key.BedrockMantleKeyConfig.SessionToken != nil {
mantleConfig.SessionToken = key.BedrockMantleKeyConfig.SessionToken.Redacted()
}
if key.BedrockMantleKeyConfig.Region != nil {
mantleConfig.Region = key.BedrockMantleKeyConfig.Region.Redacted()
}
if key.BedrockMantleKeyConfig.RoleARN != nil {
mantleConfig.RoleARN = key.BedrockMantleKeyConfig.RoleARN.Redacted()
}
if key.BedrockMantleKeyConfig.ExternalID != nil {
mantleConfig.ExternalID = key.BedrockMantleKeyConfig.ExternalID.Redacted()
}
if key.BedrockMantleKeyConfig.RoleSessionName != nil {
mantleConfig.RoleSessionName = key.BedrockMantleKeyConfig.RoleSessionName.Redacted()
}
redactedConfig.Keys[i].BedrockMantleKeyConfig = mantleConfig
}

if key.VLLMKeyConfig != nil {
vllmConfig := &schemas.VLLMKeyConfig{
ModelName: key.VLLMKeyConfig.ModelName,
Expand Down
43 changes: 43 additions & 0 deletions framework/configstore/migrations.go
Original file line number Diff line number Diff line change
Expand Up @@ -430,6 +430,7 @@ var configstoreMigrationSteps = []migrationStep{
{IDs: []string{"null_legacy_customer_budget_id_refs"}, run: migrationNullLegacyCustomerBudgetID},
{IDs: []string{"add_skills_repo_tables"}, run: migrationAddSkillsRepoTables},
{IDs: []string{"add_dump_errors_in_console_logs_column"}, run: migrationAddDumpErrorsInConsoleLogsColumn},
{IDs: []string{"add_bedrock_mantle_key_columns"}, run: migrationAddBedrockMantleKeyColumns},
}

// quoteSQLiteIdentifier quotes a SQLite identifier, escaping any double quotes.
Expand Down Expand Up @@ -1110,6 +1111,48 @@ func migrationAddVirtualKeyProviderConfigTable(ctx context.Context, db *gorm.DB,
}

// migrationAddAllowedOriginsJSONColumn adds the allowed_origins_json column to the client config table
// migrationAddBedrockMantleKeyColumns adds the bedrock_mantle_* SigV4 credential columns to the
// config_keys table for the standalone bedrock_mantle provider.
func migrationAddBedrockMantleKeyColumns(ctx context.Context, db *gorm.DB, logger schemas.Logger) error {
migrationName := "add_bedrock_mantle_key_columns"
logger.Info("[configstore] starting migration %s", migrationName)
defer logger.Info("[configstore] finished migration %s", migrationName)
cols := []string{
"bedrock_mantle_access_key",
"bedrock_mantle_secret_key",
"bedrock_mantle_session_token",
"bedrock_mantle_region",
"bedrock_mantle_role_arn",
"bedrock_mantle_external_id",
"bedrock_mantle_role_session_name",
}
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
ID: migrationName,
Migrate: func(tx *gorm.DB) error {
tx = tx.WithContext(ctx)
for _, col := range cols {
if err := addColumnIfNotExists(tx, logger, &tables.TableKey{}, col); err != nil {
return err
}
}
return nil
},
Rollback: func(tx *gorm.DB) error {
tx = tx.WithContext(ctx)
for _, col := range cols {
if err := dropColumnIfExists(tx, logger, &tables.TableKey{}, col); err != nil {
return err
}
}
return nil
},
}})
if err := m.Migrate(); err != nil {
return fmt.Errorf("error while running db migration: %s", err.Error())
}
return nil
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

func migrationAddAllowedOriginsJSONColumn(ctx context.Context, db *gorm.DB, logger schemas.Logger) error {
migrationName := "add_allowed_origins_json_column"
logger.Info("[configstore] starting migration %s", migrationName)
Expand Down
172 changes: 148 additions & 24 deletions framework/configstore/tables/key.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,18 +13,18 @@ import (

// TableKey represents an API key configuration in the database
type TableKey struct {
ID uint `gorm:"primaryKey;autoIncrement" json:"id"`
Name string `gorm:"type:varchar(255);uniqueIndex:idx_key_name;not null" json:"name"`
ProviderID uint `gorm:"index;not null" json:"provider_id"`
Provider string `gorm:"index;type:varchar(50)" json:"provider"` // ModelProvider as string
KeyID string `gorm:"type:varchar(255);uniqueIndex:idx_key_id;not null" json:"key_id"` // UUID from schemas.Key
ID uint `gorm:"primaryKey;autoIncrement" json:"id"`
Name string `gorm:"type:varchar(255);uniqueIndex:idx_key_name;not null" json:"name"`
ProviderID uint `gorm:"index;not null" json:"provider_id"`
Provider string `gorm:"index;type:varchar(50)" json:"provider"` // ModelProvider as string
KeyID string `gorm:"type:varchar(255);uniqueIndex:idx_key_id;not null" json:"key_id"` // UUID from schemas.Key
Value schemas.SecretVar `gorm:"type:text;not null" json:"value"`
ModelsJSON string `gorm:"type:text" json:"-"` // JSON serialized []string
BlacklistedModelsJSON string `gorm:"type:text" json:"-"` // JSON serialized []string
Weight *float64 `json:"weight"`
Enabled *bool `gorm:"default:true" json:"enabled,omitempty"`
CreatedAt time.Time `gorm:"index;not null" json:"created_at"`
UpdatedAt time.Time `gorm:"index;not null" json:"updated_at"`
ModelsJSON string `gorm:"type:text" json:"-"` // JSON serialized []string
BlacklistedModelsJSON string `gorm:"type:text" json:"-"` // JSON serialized []string
Weight *float64 `json:"weight"`
Enabled *bool `gorm:"default:true" json:"enabled,omitempty"`
CreatedAt time.Time `gorm:"index;not null" json:"created_at"`
UpdatedAt time.Time `gorm:"index;not null" json:"updated_at"`

// Config hash is used to detect changes synced from config.json file
ConfigHash string `gorm:"type:varchar(255);null" json:"config_hash"`
Expand All @@ -37,7 +37,7 @@ type TableKey struct {
AzureClientID *schemas.SecretVar `gorm:"type:text" json:"azure_client_id,omitempty"`
AzureClientSecret *schemas.SecretVar `gorm:"type:text" json:"azure_client_secret,omitempty"`
AzureTenantID *schemas.SecretVar `gorm:"type:text" json:"azure_tenant_id,omitempty"`
AzureScopesJSON *string `gorm:"column:azure_scopes;type:text" json:"-"` // JSON serialized []string
AzureScopesJSON *string `gorm:"column:azure_scopes;type:text" json:"-"` // JSON serialized []string

// Vertex config fields (embedded)
VertexProjectID *schemas.SecretVar `gorm:"type:text" json:"vertex_project_id,omitempty"`
Expand All @@ -54,11 +54,20 @@ type TableKey struct {
BedrockRoleARN *schemas.SecretVar `gorm:"type:text" json:"bedrock_role_arn,omitempty"`
BedrockExternalID *schemas.SecretVar `gorm:"type:text" json:"bedrock_external_id,omitempty"`
BedrockRoleSessionName *schemas.SecretVar `gorm:"type:text" json:"bedrock_role_session_name,omitempty"`
BedrockBatchS3ConfigJSON *string `gorm:"type:text" json:"-"` // JSON serialized schemas.BatchS3Config
BedrockBatchS3ConfigJSON *string `gorm:"type:text" json:"-"` // JSON serialized schemas.BatchS3Config

// Bedrock Mantle config fields (embedded)
BedrockMantleAccessKey *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_access_key,omitempty"`
BedrockMantleSecretKey *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_secret_key,omitempty"`
BedrockMantleSessionToken *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_session_token,omitempty"`
BedrockMantleRegion *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_region,omitempty"`
BedrockMantleRoleARN *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_role_arn,omitempty"`
BedrockMantleExternalID *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_external_id,omitempty"`
BedrockMantleRoleSessionName *schemas.SecretVar `gorm:"type:text" json:"bedrock_mantle_role_session_name,omitempty"`

// VLLM config fields (embedded)
VLLMUrl *schemas.SecretVar `gorm:"type:text" json:"vllm_url,omitempty"`
VLLMModelName *string `gorm:"type:varchar(255)" json:"vllm_model_name,omitempty"`
VLLMModelName *string `gorm:"type:varchar(255)" json:"vllm_model_name,omitempty"`

// Replicate config fields (embedded)
ReplicateUseDeploymentsEndpoint *bool `gorm:"column:replicate_use_deployments_endpoint" json:"replicate_use_deployments_endpoint,omitempty"`
Expand All @@ -78,16 +87,17 @@ type TableKey struct {
EncryptionStatus string `gorm:"type:varchar(20);default:'plain_text'" json:"-"`

// Virtual fields for runtime use (not stored in DB)
Models schemas.WhiteList `gorm:"-" json:"models"` // ["*"] allows all models; empty denies all (deny-by-default)
BlacklistedModels schemas.BlackList `gorm:"-" json:"blacklisted_models"`
Aliases schemas.KeyAliases `gorm:"-" json:"aliases,omitempty"`
AzureKeyConfig *schemas.AzureKeyConfig `gorm:"-" json:"azure_key_config,omitempty"`
VertexKeyConfig *schemas.VertexKeyConfig `gorm:"-" json:"vertex_key_config,omitempty"`
BedrockKeyConfig *schemas.BedrockKeyConfig `gorm:"-" json:"bedrock_key_config,omitempty"`
VLLMKeyConfig *schemas.VLLMKeyConfig `gorm:"-" json:"vllm_key_config,omitempty"`
ReplicateKeyConfig *schemas.ReplicateKeyConfig `gorm:"-" json:"replicate_key_config,omitempty"`
OllamaKeyConfig *schemas.OllamaKeyConfig `gorm:"-" json:"ollama_key_config,omitempty"`
SGLKeyConfig *schemas.SGLKeyConfig `gorm:"-" json:"sgl_key_config,omitempty"`
Models schemas.WhiteList `gorm:"-" json:"models"` // ["*"] allows all models; empty denies all (deny-by-default)
BlacklistedModels schemas.BlackList `gorm:"-" json:"blacklisted_models"`
Aliases schemas.KeyAliases `gorm:"-" json:"aliases,omitempty"`
AzureKeyConfig *schemas.AzureKeyConfig `gorm:"-" json:"azure_key_config,omitempty"`
VertexKeyConfig *schemas.VertexKeyConfig `gorm:"-" json:"vertex_key_config,omitempty"`
BedrockKeyConfig *schemas.BedrockKeyConfig `gorm:"-" json:"bedrock_key_config,omitempty"`
BedrockMantleKeyConfig *schemas.BedrockMantleKeyConfig `gorm:"-" json:"bedrock_mantle_key_config,omitempty"`
VLLMKeyConfig *schemas.VLLMKeyConfig `gorm:"-" json:"vllm_key_config,omitempty"`
ReplicateKeyConfig *schemas.ReplicateKeyConfig `gorm:"-" json:"replicate_key_config,omitempty"`
OllamaKeyConfig *schemas.OllamaKeyConfig `gorm:"-" json:"ollama_key_config,omitempty"`
SGLKeyConfig *schemas.SGLKeyConfig `gorm:"-" json:"sgl_key_config,omitempty"`
}

// TableName sets the table name for each model
Expand Down Expand Up @@ -275,6 +285,60 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error {
k.BedrockBatchS3ConfigJSON = nil
}

if k.BedrockMantleKeyConfig != nil {
// Copy to avoid encrypting the shared BedrockMantleKeyConfig through the pointer.
if k.BedrockMantleKeyConfig.AccessKey.IsSet() {
ak := k.BedrockMantleKeyConfig.AccessKey
k.BedrockMantleAccessKey = &ak
} else {
k.BedrockMantleAccessKey = nil
}
if k.BedrockMantleKeyConfig.SecretKey.IsSet() {
sk := k.BedrockMantleKeyConfig.SecretKey
k.BedrockMantleSecretKey = &sk
} else {
k.BedrockMantleSecretKey = nil
}
if k.BedrockMantleKeyConfig.SessionToken != nil {
st := *k.BedrockMantleKeyConfig.SessionToken
k.BedrockMantleSessionToken = &st
} else {
k.BedrockMantleSessionToken = nil
}
if k.BedrockMantleKeyConfig.Region != nil {
br := *k.BedrockMantleKeyConfig.Region
k.BedrockMantleRegion = &br
} else {
k.BedrockMantleRegion = nil
}
if k.BedrockMantleKeyConfig.RoleARN != nil {
bra := *k.BedrockMantleKeyConfig.RoleARN
k.BedrockMantleRoleARN = &bra
} else {
k.BedrockMantleRoleARN = nil
}
if k.BedrockMantleKeyConfig.ExternalID != nil {
ei := *k.BedrockMantleKeyConfig.ExternalID
k.BedrockMantleExternalID = &ei
} else {
k.BedrockMantleExternalID = nil
}
if k.BedrockMantleKeyConfig.RoleSessionName != nil {
rsn := *k.BedrockMantleKeyConfig.RoleSessionName
k.BedrockMantleRoleSessionName = &rsn
} else {
k.BedrockMantleRoleSessionName = nil
}
} else {
k.BedrockMantleAccessKey = nil
k.BedrockMantleSecretKey = nil
k.BedrockMantleSessionToken = nil
k.BedrockMantleRegion = nil
k.BedrockMantleRoleARN = nil
k.BedrockMantleExternalID = nil
k.BedrockMantleRoleSessionName = nil
}

Comment thread
coderabbitai[bot] marked this conversation as resolved.
if k.Aliases != nil {
data, err := sonic.Marshal(k.Aliases)
if err != nil {
Expand Down Expand Up @@ -396,6 +460,28 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error {
if err := encryptString(k.BedrockBatchS3ConfigJSON); err != nil {
return fmt.Errorf("failed to encrypt bedrock batch s3 config: %w", err)
}
// Bedrock Mantle
if err := encryptSecretVarPtr(&k.BedrockMantleAccessKey); err != nil {
return fmt.Errorf("failed to encrypt bedrock mantle access key: %w", err)
}
if err := encryptSecretVarPtr(&k.BedrockMantleSecretKey); err != nil {
return fmt.Errorf("failed to encrypt bedrock mantle secret key: %w", err)
}
if err := encryptSecretVarPtr(&k.BedrockMantleSessionToken); err != nil {
return fmt.Errorf("failed to encrypt bedrock mantle session token: %w", err)
}
if err := encryptSecretVarPtr(&k.BedrockMantleRegion); err != nil {
return fmt.Errorf("failed to encrypt bedrock mantle region: %w", err)
}
if err := encryptSecretVarPtr(&k.BedrockMantleRoleARN); err != nil {
return fmt.Errorf("failed to encrypt bedrock mantle role arn: %w", err)
}
if err := encryptSecretVarPtr(&k.BedrockMantleExternalID); err != nil {
return fmt.Errorf("failed to encrypt bedrock mantle external id: %w", err)
}
if err := encryptSecretVarPtr(&k.BedrockMantleRoleSessionName); err != nil {
return fmt.Errorf("failed to encrypt bedrock mantle role session name: %w", err)
}
// Aliases
if err := encryptString(k.AliasesJSON); err != nil {
return fmt.Errorf("failed to encrypt aliases: %w", err)
Expand Down Expand Up @@ -480,6 +566,28 @@ func (k *TableKey) AfterFind(tx *gorm.DB) error {
if err := decryptString(k.BedrockBatchS3ConfigJSON); err != nil {
return fmt.Errorf("failed to decrypt bedrock batch s3 config: %w", err)
}
// Bedrock Mantle
if err := decryptSecretVarPtr(&k.BedrockMantleAccessKey); err != nil {
return fmt.Errorf("failed to decrypt bedrock mantle access key: %w", err)
}
if err := decryptSecretVarPtr(&k.BedrockMantleSecretKey); err != nil {
return fmt.Errorf("failed to decrypt bedrock mantle secret key: %w", err)
}
if err := decryptSecretVarPtr(&k.BedrockMantleSessionToken); err != nil {
return fmt.Errorf("failed to decrypt bedrock mantle session token: %w", err)
}
if err := decryptSecretVarPtr(&k.BedrockMantleRegion); err != nil {
return fmt.Errorf("failed to decrypt bedrock mantle region: %w", err)
}
if err := decryptSecretVarPtr(&k.BedrockMantleRoleARN); err != nil {
return fmt.Errorf("failed to decrypt bedrock mantle role arn: %w", err)
}
if err := decryptSecretVarPtr(&k.BedrockMantleExternalID); err != nil {
return fmt.Errorf("failed to decrypt bedrock mantle external id: %w", err)
}
if err := decryptSecretVarPtr(&k.BedrockMantleRoleSessionName); err != nil {
return fmt.Errorf("failed to decrypt bedrock mantle role session name: %w", err)
}
// Aliases
if err := decryptString(k.AliasesJSON); err != nil {
return fmt.Errorf("failed to decrypt aliases: %w", err)
Expand Down Expand Up @@ -587,6 +695,22 @@ func (k *TableKey) AfterFind(tx *gorm.DB) error {

k.BedrockKeyConfig = bedrockConfig
}
// Reconstruct Bedrock Mantle config if fields are present
if k.BedrockMantleAccessKey != nil || k.BedrockMantleSecretKey != nil || k.BedrockMantleSessionToken != nil || k.BedrockMantleRegion != nil || k.BedrockMantleRoleARN != nil || k.BedrockMantleExternalID != nil || k.BedrockMantleRoleSessionName != nil {
mantleConfig := &schemas.BedrockMantleKeyConfig{}
if k.BedrockMantleAccessKey != nil {
mantleConfig.AccessKey = *k.BedrockMantleAccessKey
}
if k.BedrockMantleSecretKey != nil {
mantleConfig.SecretKey = *k.BedrockMantleSecretKey
}
mantleConfig.SessionToken = k.BedrockMantleSessionToken
mantleConfig.Region = k.BedrockMantleRegion
mantleConfig.RoleARN = k.BedrockMantleRoleARN
mantleConfig.ExternalID = k.BedrockMantleExternalID
mantleConfig.RoleSessionName = k.BedrockMantleRoleSessionName
k.BedrockMantleKeyConfig = mantleConfig
}
// Reconstruct Aliases
if k.AliasesJSON != nil && *k.AliasesJSON != "" {
var aliases schemas.KeyAliases
Expand Down
Loading
Loading