Skip to content

Commit 86a7d30

Browse files
tejas ghattetejas ghatte
authored andcommitted
feat: add scopes field in azure key config for entra id auth
1 parent 447e726 commit 86a7d30

File tree

12 files changed

+128
-19
lines changed

12 files changed

+128
-19
lines changed

core/providers/azure/azure.go

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,8 +83,13 @@ func (provider *AzureProvider) getAzureAuthHeaders(ctx *schemas.BifrostContext,
8383
return nil, providerUtils.NewBifrostOperationError("failed to get or create Azure authentication", err, schemas.Azure)
8484
}
8585

86+
scopes := []string{DefaultAzureScope}
87+
if len(key.AzureKeyConfig.Scopes) > 0 {
88+
scopes = key.AzureKeyConfig.Scopes
89+
}
90+
8691
token, err := cred.GetToken(ctx, policy.TokenRequestOptions{
87-
Scopes: []string{DefaultAzureScope},
92+
Scopes: scopes,
8893
})
8994
if err != nil {
9095
return nil, providerUtils.NewBifrostOperationError("failed to get Azure access token", err, schemas.Azure)

core/providers/azure/files.go

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,13 @@ func (provider *AzureProvider) setAzureAuth(ctx context.Context, req *fasthttp.R
2626
return providerUtils.NewBifrostOperationError("failed to get or create Azure authentication", err, schemas.Azure)
2727
}
2828

29+
scopes := []string{DefaultAzureScope}
30+
if len(key.AzureKeyConfig.Scopes) > 0 {
31+
scopes = key.AzureKeyConfig.Scopes
32+
}
33+
2934
token, err := cred.GetToken(ctx, policy.TokenRequestOptions{
30-
Scopes: []string{DefaultAzureScope},
35+
Scopes: scopes,
3136
})
3237
if err != nil {
3338
return providerUtils.NewBifrostOperationError("failed to get Azure access token", err, schemas.Azure)

core/schemas/account.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,10 @@ type AzureKeyConfig struct {
2727
Deployments map[string]string `json:"deployments,omitempty"` // Mapping of model names to deployment names
2828
APIVersion *EnvVar `json:"api_version,omitempty"` // Azure API version to use; defaults to "2024-10-21"
2929

30-
ClientID *EnvVar `json:"client_id,omitempty"` // Azure client ID for authentication
31-
ClientSecret *EnvVar `json:"client_secret,omitempty"` // Azure client secret for authentication
32-
TenantID *EnvVar `json:"tenant_id,omitempty"` // Azure tenant ID for authentication
30+
ClientID *EnvVar `json:"client_id,omitempty"` // Azure client ID for authentication
31+
ClientSecret *EnvVar `json:"client_secret,omitempty"` // Azure client secret for authentication
32+
TenantID *EnvVar `json:"tenant_id,omitempty"` // Azure tenant ID for authentication
33+
Scopes []string `json:"scopes,omitempty"`
3334
}
3435

3536
// VertexKeyConfig represents the Vertex-specific configuration.

docs/providers/supported-providers/azure.mdx

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ If you set `client_id`, `client_secret`, and `tenant_id`, Azure Entra ID authent
104104
"client_id": "your-client-id",
105105
"client_secret": "your-client-secret",
106106
"tenant_id": "your-tenant-id",
107+
"scopes": ["https://cognitiveservices.azure.com/.default"],
107108
"api_version": "2024-10-21",
108109
"deployments": {
109110
"gpt-4": "my-gpt4-deployment",
@@ -139,6 +140,7 @@ If you set `client_id`, `client_secret`, and `tenant_id`, Azure Entra ID authent
139140
- `client_id` - Azure Entra ID client ID (optional, for Service Principal auth)
140141
- `client_secret` - Azure Entra ID client secret (optional, for Service Principal auth)
141142
- `tenant_id` - Azure Entra ID tenant ID (optional, for Service Principal auth)
143+
- `scopes` - OAuth scopes for token requests (default: `["https://cognitiveservices.azure.com/.default"]`)
142144
- `api_version` - API version to use (default: `2024-10-21`)
143145
- `deployments` - Map of model names to deployment IDs (optional, can be provided per-request)
144146
- `allowed_models` - List of allowed models to use from this key (optional)

docs/quickstart/gateway/provider-configuration.mdx

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1031,6 +1031,7 @@ curl --location 'http://localhost:8080/api/providers' \
10311031
"client_id": "env.AZURE_CLIENT_ID",
10321032
"client_secret": "env.AZURE_CLIENT_SECRET",
10331033
"tenant_id": "env.AZURE_TENANT_ID",
1034+
"scopes": ["https://cognitiveservices.azure.com/.default"],
10341035
"deployments": {
10351036
"gpt-4o": "gpt-4o-deployment",
10361037
"gpt-4o-mini": "gpt-4o-mini-deployment"
@@ -1061,6 +1062,7 @@ curl --location 'http://localhost:8080/api/providers' \
10611062
"client_id": "env.AZURE_CLIENT_ID",
10621063
"client_secret": "env.AZURE_CLIENT_SECRET",
10631064
"tenant_id": "env.AZURE_TENANT_ID",
1065+
"scopes": ["https://cognitiveservices.azure.com/.default"],
10641066
"deployments": {
10651067
"gpt-4o": "gpt-4o-deployment",
10661068
"gpt-4o-mini": "gpt-4o-mini-deployment"

framework/configstore/clientconfig.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,9 @@ func (p *ProviderConfig) Redacted() *ProviderConfig {
293293
if key.AzureKeyConfig.TenantID != nil {
294294
azureConfig.TenantID = key.AzureKeyConfig.TenantID.Redacted()
295295
}
296+
if len(key.AzureKeyConfig.Scopes) > 0 {
297+
azureConfig.Scopes = key.AzureKeyConfig.Scopes
298+
}
296299
redactedConfig.Keys[i].AzureKeyConfig = azureConfig
297300
}
298301

framework/configstore/migrations.go

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error {
179179
if err := migrationAddBaseModelPricingColumn(ctx, db); err != nil {
180180
return err
181181
}
182+
if err := migrationAddAzureScopesColumn(ctx, db); err != nil {
183+
return err
184+
}
182185
return nil
183186
}
184187

@@ -3227,3 +3230,34 @@ func migrationAddBaseModelPricingColumn(ctx context.Context, db *gorm.DB) error
32273230
}})
32283231
return m.Migrate()
32293232
}
3233+
3234+
// migrationAddAzureScopesColumn adds the azure_scopes column to the key table for Entra ID OAuth scopes
3235+
func migrationAddAzureScopesColumn(ctx context.Context, db *gorm.DB) error {
3236+
m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{
3237+
ID: "add_azure_scopes_column",
3238+
Migrate: func(tx *gorm.DB) error {
3239+
tx = tx.WithContext(ctx)
3240+
migrator := tx.Migrator()
3241+
if !migrator.HasColumn(&tables.TableKey{}, "azure_scopes") {
3242+
if err := migrator.AddColumn(&tables.TableKey{}, "azure_scopes"); err != nil {
3243+
return fmt.Errorf("failed to add azure_scopes column: %w", err)
3244+
}
3245+
}
3246+
return nil
3247+
},
3248+
Rollback: func(tx *gorm.DB) error {
3249+
tx = tx.WithContext(ctx)
3250+
migrator := tx.Migrator()
3251+
if migrator.HasColumn(&tables.TableKey{}, "azure_scopes") {
3252+
if err := migrator.DropColumn(&tables.TableKey{}, "azure_scopes"); err != nil {
3253+
return fmt.Errorf("failed to drop azure_scopes column: %w", err)
3254+
}
3255+
}
3256+
return nil
3257+
},
3258+
}})
3259+
if err := m.Migrate(); err != nil {
3260+
return fmt.Errorf("error running azure_scopes migration: %s", err.Error())
3261+
}
3262+
return nil
3263+
}

framework/configstore/tables/key.go

Lines changed: 33 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -11,44 +11,45 @@ import (
1111

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

2626
// Config hash is used to detect changes synced from config.json file
2727
ConfigHash string `gorm:"type:varchar(255);null" json:"config_hash"`
2828

2929
// Azure config fields (embedded instead of separate table for simplicity)
3030
AzureEndpoint *schemas.EnvVar `gorm:"type:text" json:"azure_endpoint,omitempty"`
3131
AzureAPIVersion *schemas.EnvVar `gorm:"type:varchar(50)" json:"azure_api_version,omitempty"`
32-
AzureDeploymentsJSON *string `gorm:"type:text" json:"-"` // JSON serialized map[string]string
32+
AzureDeploymentsJSON *string `gorm:"type:text" json:"-"` // JSON serialized map[string]string
3333
AzureClientID *schemas.EnvVar `gorm:"type:varchar(255)" json:"azure_client_id,omitempty"`
3434
AzureClientSecret *schemas.EnvVar `gorm:"type:text" json:"azure_client_secret,omitempty"`
3535
AzureTenantID *schemas.EnvVar `gorm:"type:varchar(255)" json:"azure_tenant_id,omitempty"`
36+
AzureScopesJSON *string `gorm:"column:azure_scopes;type:text" json:"-"` // JSON serialized []string
3637

3738
// Vertex config fields (embedded)
3839
VertexProjectID *schemas.EnvVar `gorm:"type:varchar(255)" json:"vertex_project_id,omitempty"`
3940
VertexProjectNumber *schemas.EnvVar `gorm:"type:varchar(255)" json:"vertex_project_number,omitempty"`
4041
VertexRegion *schemas.EnvVar `gorm:"type:varchar(100)" json:"vertex_region,omitempty"`
4142
VertexAuthCredentials *schemas.EnvVar `gorm:"type:text" json:"vertex_auth_credentials,omitempty"`
42-
VertexDeploymentsJSON *string `gorm:"type:text" json:"-"` // JSON serialized map[string]string
43+
VertexDeploymentsJSON *string `gorm:"type:text" json:"-"` // JSON serialized map[string]string
4344

4445
// Bedrock config fields (embedded)
4546
BedrockAccessKey *schemas.EnvVar `gorm:"type:varchar(255)" json:"bedrock_access_key,omitempty"`
4647
BedrockSecretKey *schemas.EnvVar `gorm:"type:text" json:"bedrock_secret_key,omitempty"`
4748
BedrockSessionToken *schemas.EnvVar `gorm:"type:text" json:"bedrock_session_token,omitempty"`
4849
BedrockRegion *schemas.EnvVar `gorm:"type:varchar(100)" json:"bedrock_region,omitempty"`
4950
BedrockARN *schemas.EnvVar `gorm:"type:text" json:"bedrock_arn,omitempty"`
50-
BedrockDeploymentsJSON *string `gorm:"type:text" json:"-"` // JSON serialized map[string]string
51-
BedrockBatchS3ConfigJSON *string `gorm:"type:text" json:"-"` // JSON serialized schemas.BatchS3Config
51+
BedrockDeploymentsJSON *string `gorm:"type:text" json:"-"` // JSON serialized map[string]string
52+
BedrockBatchS3ConfigJSON *string `gorm:"type:text" json:"-"` // JSON serialized schemas.BatchS3Config
5253

5354
// Batch API configuration
5455
UseForBatchAPI *bool `gorm:"default:false" json:"use_for_batch_api,omitempty"` // Whether this key can be used for batch API operations
@@ -95,6 +96,16 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error {
9596
k.AzureClientID = k.AzureKeyConfig.ClientID
9697
k.AzureClientSecret = k.AzureKeyConfig.ClientSecret
9798
k.AzureTenantID = k.AzureKeyConfig.TenantID
99+
if len(k.AzureKeyConfig.Scopes) > 0 {
100+
data, err := json.Marshal(k.AzureKeyConfig.Scopes)
101+
if err != nil {
102+
return err
103+
}
104+
s := string(data)
105+
k.AzureScopesJSON = &s
106+
} else {
107+
k.AzureScopesJSON = nil
108+
}
98109
if k.AzureKeyConfig.Deployments != nil {
99110
data, err := json.Marshal(k.AzureKeyConfig.Deployments)
100111
if err != nil {
@@ -112,6 +123,7 @@ func (k *TableKey) BeforeSave(tx *gorm.DB) error {
112123
k.AzureClientID = nil
113124
k.AzureClientSecret = nil
114125
k.AzureTenantID = nil
126+
k.AzureScopesJSON = nil
115127
}
116128
// BeforeSave is called before saving the key to the database
117129
if k.VertexKeyConfig != nil {
@@ -217,12 +229,19 @@ func (k *TableKey) AfterFind(tx *gorm.DB) error {
217229
}
218230
// Reconstruct Azure config if fields are present
219231
if k.AzureEndpoint != nil {
232+
var scopes []string
233+
if k.AzureScopesJSON != nil && *k.AzureScopesJSON != "" {
234+
if err := json.Unmarshal([]byte(*k.AzureScopesJSON), &scopes); err != nil {
235+
return err
236+
}
237+
}
220238
azureConfig := &schemas.AzureKeyConfig{
221239
Endpoint: *schemas.NewEnvVar(""),
222240
APIVersion: k.AzureAPIVersion,
223241
ClientID: k.AzureClientID,
224242
ClientSecret: k.AzureClientSecret,
225243
TenantID: k.AzureTenantID,
244+
Scopes: scopes,
226245
}
227246

228247
if k.AzureEndpoint != nil {

framework/configstore/tables/virtualkey.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ func (pc *TableVirtualKeyProviderConfig) AfterFind(tx *gorm.DB) error {
106106
key.AzureClientID = nil
107107
key.AzureClientSecret = nil
108108
key.AzureTenantID = nil
109+
key.AzureScopesJSON = nil
109110
key.AzureDeploymentsJSON = nil
110111
key.AzureKeyConfig = nil
111112

ui/app/workspace/providers/fragments/apiKeysFormFragment.tsx

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { EnvVarInput } from "@/components/ui/envVarInput";
66
import { FormControl, FormDescription, FormField, FormItem, FormLabel, FormMessage } from "@/components/ui/form";
77
import { Input } from "@/components/ui/input";
88
import { ModelMultiselect } from "@/components/ui/modelMultiselect";
9+
import { TagInput } from "@/components/ui/tagInput";
910
import { Separator } from "@/components/ui/separator";
1011
import { Switch } from "@/components/ui/switch";
1112
import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
@@ -220,6 +221,7 @@ export function ApiKeyFormFragment({ control, providerName, form }: Props) {
220221
form.setValue('key.azure_key_config.client_id', undefined)
221222
form.setValue('key.azure_key_config.client_secret', undefined)
222223
form.setValue('key.azure_key_config.tenant_id', undefined)
224+
form.setValue('key.azure_key_config.scopes', undefined)
223225
}
224226
}}>
225227
<TabsList className="grid w-full grid-cols-2">
@@ -297,6 +299,38 @@ export function ApiKeyFormFragment({ control, providerName, form }: Props) {
297299
</FormItem>
298300
)}
299301
/>
302+
<FormField
303+
control={control}
304+
name={`key.azure_key_config.scopes`}
305+
render={({ field }) => (
306+
<FormItem>
307+
<div className="flex items-center gap-2">
308+
<FormLabel>Scopes (Optional)</FormLabel>
309+
<TooltipProvider>
310+
<Tooltip>
311+
<TooltipTrigger asChild>
312+
<span>
313+
<Info className="text-muted-foreground h-3 w-3" />
314+
</span>
315+
</TooltipTrigger>
316+
<TooltipContent>
317+
<p>Optional OAuth scopes for token requests. By default we use
318+
https://cognitiveservices.azure.com/.default — add additional scopes here if your setup requires extra permissions.</p>
319+
</TooltipContent>
320+
</Tooltip>
321+
</TooltipProvider>
322+
</div>
323+
<FormControl>
324+
<TagInput
325+
placeholder="Add scope (Enter or comma)"
326+
value={field.value ?? []}
327+
onValueChange={field.onChange}
328+
/>
329+
</FormControl>
330+
<FormMessage />
331+
</FormItem>
332+
)}
333+
/>
300334
</>
301335
)}
302336

0 commit comments

Comments
 (0)