first commit
This commit is contained in:
Beyhan Oğur
294
core/schemas/account.go
Normal file
294
core/schemas/account.go
Normal file
@@ -0,0 +1,294 @@
|
||||
// Package schemas defines the core schemas and types used by the Bifrost system.
|
||||
package schemas
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"slices"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type KeyStatusType string
|
||||
|
||||
const (
|
||||
KeyStatusSuccess KeyStatusType = "success"
|
||||
KeyStatusListModelsFailed KeyStatusType = "list_models_failed"
|
||||
)
|
||||
|
||||
// WhiteList is a list of values that are allowed to be used.
|
||||
// Semantics:
|
||||
// - "*" (alone) means all values are allowed.
|
||||
// - Empty list means nothing is allowed.
|
||||
// - Non-empty list (without "*") means only the listed values are allowed.
|
||||
//
|
||||
// This type is used generically for any field that needs whitelist behavior
|
||||
// (e.g., allowed models, allowed tools).
|
||||
type WhiteList []string
|
||||
|
||||
// Contains reports whether value is in the whitelist.
|
||||
// Returns true if value is in the list.
|
||||
func (wl WhiteList) Contains(value string) bool {
|
||||
return slices.ContainsFunc(wl, func(s string) bool {
|
||||
return strings.EqualFold(s, value)
|
||||
})
|
||||
}
|
||||
|
||||
// IsAllowed reports whether value is in the whitelist.
|
||||
// Returns true if value is in the list.
|
||||
func (wl WhiteList) IsAllowed(value string) bool {
|
||||
return wl.IsUnrestricted() || wl.Contains(value)
|
||||
}
|
||||
|
||||
// IsEmpty reports whether the whitelist has no entries.
|
||||
func (wl WhiteList) IsEmpty() bool {
|
||||
return len(wl) == 0
|
||||
}
|
||||
|
||||
// IsUnrestricted reports whether the whitelist contains only "*",
|
||||
// meaning all values are allowed.
|
||||
func (wl WhiteList) IsUnrestricted() bool {
|
||||
return len(wl) == 1 && wl[0] == "*"
|
||||
}
|
||||
|
||||
// IsRestricted reports whether the whitelist contains entries other than "*",
|
||||
// meaning only the listed values are allowed.
|
||||
func (wl WhiteList) IsRestricted() bool {
|
||||
return !wl.IsUnrestricted()
|
||||
}
|
||||
|
||||
// Validate checks that the whitelist is well-formed.
|
||||
// Returns an error if "*" is present alongside other values, or if there are duplicate entries.
|
||||
func (wl WhiteList) Validate() error {
|
||||
if wl.Contains("*") && len(wl) > 1 {
|
||||
return fmt.Errorf("wildcard '*' cannot be used with other values in the whitelist")
|
||||
}
|
||||
seen := make(map[string]struct{}, len(wl))
|
||||
for _, v := range wl {
|
||||
normalized := strings.ToLower(v)
|
||||
if _, ok := seen[normalized]; ok {
|
||||
return fmt.Errorf("duplicate value '%s' in whitelist", v)
|
||||
}
|
||||
seen[normalized] = struct{}{}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// BlackList is a list of values that are denied.
|
||||
// Semantics:
|
||||
// - "*" (alone) means all values are blocked.
|
||||
// - Empty list means nothing is blocked.
|
||||
// - Non-empty list (without "*") means only the listed values are blocked.
|
||||
type BlackList []string
|
||||
|
||||
func (bl BlackList) Contains(value string) bool {
|
||||
return slices.ContainsFunc(bl, func(s string) bool {
|
||||
return strings.EqualFold(s, value)
|
||||
})
|
||||
}
|
||||
|
||||
// IsBlocked reports whether value is blocked.
|
||||
func (bl BlackList) IsBlocked(value string) bool {
|
||||
return bl.IsBlockAll() || bl.Contains(value)
|
||||
}
|
||||
|
||||
// IsEmpty reports whether the blacklist has no entries (nothing is blocked).
|
||||
func (bl BlackList) IsEmpty() bool {
|
||||
return len(bl) == 0
|
||||
}
|
||||
|
||||
// IsBlockAll reports whether the blacklist contains "*", meaning all values are blocked.
|
||||
func (bl BlackList) IsBlockAll() bool {
|
||||
return len(bl) == 1 && bl[0] == "*"
|
||||
}
|
||||
|
||||
// Validate checks that the blacklist is well-formed.
|
||||
func (bl BlackList) Validate() error {
|
||||
if bl.Contains("*") && len(bl) > 1 {
|
||||
return fmt.Errorf("wildcard '*' cannot be used with other values in the blacklist")
|
||||
}
|
||||
seen := make(map[string]struct{}, len(bl))
|
||||
for _, v := range bl {
|
||||
normalized := strings.ToLower(v)
|
||||
if _, ok := seen[normalized]; ok {
|
||||
return fmt.Errorf("duplicate value '%s' in blacklist", v)
|
||||
}
|
||||
seen[normalized] = struct{}{}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Key represents an API key and its associated configuration for a provider.
|
||||
// It contains the key value, supported models, and a weight for load balancing.
|
||||
type Key struct {
|
||||
ID string `json:"id"` // The unique identifier for the key (used by bifrost to identify the key)
|
||||
Name string `json:"name"` // The name of the key (used by users to identify the key, not used by bifrost)
|
||||
Value EnvVar `json:"value"` // The actual API key value
|
||||
Models WhiteList `json:"models"` // List of models this key can access
|
||||
BlacklistedModels BlackList `json:"blacklisted_models"` // List of models this key cannot access
|
||||
Weight float64 `json:"weight"` // Weight for load balancing between multiple keys
|
||||
Aliases KeyAliases `json:"aliases,omitempty"` // Mapping of model identifiers to inference profiles
|
||||
AzureKeyConfig *AzureKeyConfig `json:"azure_key_config,omitempty"` // Azure-specific key configuration
|
||||
VertexKeyConfig *VertexKeyConfig `json:"vertex_key_config,omitempty"` // Vertex-specific key configuration
|
||||
BedrockKeyConfig *BedrockKeyConfig `json:"bedrock_key_config,omitempty"` // AWS Bedrock-specific key configuration
|
||||
VLLMKeyConfig *VLLMKeyConfig `json:"vllm_key_config,omitempty"` // vLLM-specific key configuration
|
||||
ReplicateKeyConfig *ReplicateKeyConfig `json:"replicate_key_config,omitempty"` // Replicate-specific key configuration
|
||||
OllamaKeyConfig *OllamaKeyConfig `json:"ollama_key_config,omitempty"` // Ollama-specific key configuration
|
||||
SGLKeyConfig *SGLKeyConfig `json:"sgl_key_config,omitempty"` // SGLang-specific key configuration
|
||||
Enabled *bool `json:"enabled,omitempty"` // Whether the key is active (default:true)
|
||||
UseForBatchAPI *bool `json:"use_for_batch_api,omitempty"` // Whether this key can be used for batch API operations (default:false for new keys, migrated keys default to true)
|
||||
ConfigHash string `json:"config_hash,omitempty"` // Hash of config.json version, used for change detection
|
||||
Status KeyStatusType `json:"status,omitempty"` // Status of key
|
||||
Description string `json:"description,omitempty"` // Description of key
|
||||
}
|
||||
|
||||
type KeyAliases map[string]string
|
||||
|
||||
func (ka KeyAliases) Validate() error {
|
||||
seen := make(map[string]struct{}, len(ka))
|
||||
for from, to := range ka {
|
||||
if strings.TrimSpace(from) == "" {
|
||||
return fmt.Errorf("alias source cannot be empty")
|
||||
}
|
||||
if strings.TrimSpace(to) == "" {
|
||||
return fmt.Errorf("alias target for %q cannot be empty", from)
|
||||
}
|
||||
if strings.TrimSpace(from) != from {
|
||||
return fmt.Errorf("alias source %q cannot have leading or trailing whitespace", from)
|
||||
}
|
||||
if strings.TrimSpace(to) != to {
|
||||
return fmt.Errorf("alias target for %q cannot have leading or trailing whitespace", from)
|
||||
}
|
||||
normalized := strings.ToLower(from)
|
||||
if _, ok := seen[normalized]; ok {
|
||||
return fmt.Errorf("duplicate alias source %q (case-insensitive)", from)
|
||||
}
|
||||
seen[normalized] = struct{}{}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (ka KeyAliases) Resolve(model string) string {
|
||||
if ka == nil {
|
||||
return model
|
||||
}
|
||||
if alias, ok := ka[model]; ok {
|
||||
return alias
|
||||
}
|
||||
// Fall back to case-insensitive lookup for consistency with WhiteList.Contains
|
||||
for k, v := range ka {
|
||||
if strings.EqualFold(k, model) {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return model
|
||||
}
|
||||
|
||||
type AzureAuthType string
|
||||
|
||||
const (
|
||||
AzureAuthTypeClientSecret AzureAuthType = "client_secret"
|
||||
AzureAuthTypeManagedIdentity AzureAuthType = "managed_identity"
|
||||
)
|
||||
|
||||
// AzureKeyConfig represents the Azure-specific configuration.
|
||||
// It contains Azure-specific settings required for service access and deployment management.
|
||||
type AzureKeyConfig struct {
|
||||
Endpoint EnvVar `json:"endpoint"` // Azure service endpoint URL
|
||||
APIVersion *EnvVar `json:"api_version,omitempty"` // Azure API version to use; defaults to "2024-10-21"
|
||||
|
||||
ClientID *EnvVar `json:"client_id,omitempty"` // Azure client ID for authentication
|
||||
ClientSecret *EnvVar `json:"client_secret,omitempty"` // Azure client secret for authentication
|
||||
TenantID *EnvVar `json:"tenant_id,omitempty"` // Azure tenant ID for authentication
|
||||
Scopes []string `json:"scopes,omitempty"`
|
||||
}
|
||||
|
||||
// VertexKeyConfig represents the Vertex-specific configuration.
|
||||
// It contains Vertex-specific settings required for authentication and service access.
|
||||
type VertexKeyConfig struct {
|
||||
ProjectID EnvVar `json:"project_id"`
|
||||
ProjectNumber EnvVar `json:"project_number"`
|
||||
Region EnvVar `json:"region"`
|
||||
AuthCredentials EnvVar `json:"auth_credentials"`
|
||||
}
|
||||
|
||||
// NOTE: To use Vertex IAM role authentication, set AuthCredentials to empty string.
|
||||
|
||||
// S3BucketConfig represents a single S3 bucket configuration for batch operations.
|
||||
type S3BucketConfig struct {
|
||||
BucketName string `json:"bucket_name"` // S3 bucket name
|
||||
Prefix string `json:"prefix,omitempty"` // S3 key prefix for batch files
|
||||
IsDefault bool `json:"is_default,omitempty"` // Whether this is the default bucket for batch operations
|
||||
}
|
||||
|
||||
// BatchS3Config holds S3 bucket configurations for Bedrock batch operations.
|
||||
// Supports multiple buckets to allow flexible batch job routing.
|
||||
type BatchS3Config struct {
|
||||
Buckets []S3BucketConfig `json:"buckets,omitempty"` // List of S3 bucket configurations
|
||||
}
|
||||
|
||||
// BedrockKeyConfig represents the AWS Bedrock-specific configuration.
|
||||
// It contains AWS-specific settings required for authentication and service access.
|
||||
type BedrockKeyConfig struct {
|
||||
AccessKey EnvVar `json:"access_key,omitempty"` // AWS access key for authentication
|
||||
SecretKey EnvVar `json:"secret_key,omitempty"` // AWS secret access key for authentication
|
||||
SessionToken *EnvVar `json:"session_token,omitempty"` // AWS session token for temporary credentials
|
||||
Region *EnvVar `json:"region,omitempty"` // AWS region for service access
|
||||
ARN *EnvVar `json:"arn,omitempty"` // Amazon Resource Name for resource identification
|
||||
// IAM role for STS AssumeRole
|
||||
RoleARN *EnvVar `json:"role_arn,omitempty"`
|
||||
ExternalID *EnvVar `json:"external_id,omitempty"`
|
||||
RoleSessionName *EnvVar `json:"session_name,omitempty"`
|
||||
|
||||
BatchS3Config *BatchS3Config `json:"batch_s3_config,omitempty"` // S3 bucket configuration for batch operations
|
||||
}
|
||||
|
||||
// NOTE: To use Bedrock IAM role authentication, set both AccessKey and SecretKey to empty strings.
|
||||
// To use Bedrock API Key authentication, set Value in Key struct instead.
|
||||
|
||||
// VLLMKeyConfig represents the vLLM-specific key configuration.
|
||||
// It allows each key to target a different vLLM server URL and model name,
|
||||
// enabling per-key routing and round-robin load balancing across multiple vLLM instances.
|
||||
type VLLMKeyConfig struct {
|
||||
URL EnvVar `json:"url"` // VLLM server base URL (required, supports env. prefix)
|
||||
ModelName string `json:"model_name"` // Exact model name served on this VLLM instance (used for key selection)
|
||||
}
|
||||
|
||||
// ReplicateKeyConfig represents the Replicate-specific key configuration.
|
||||
// It contains Replicate-specific settings required for authentication and service access.
|
||||
type ReplicateKeyConfig struct {
|
||||
UseDeploymentsEndpoint bool `json:"use_deployments_endpoint"` // Whether to use the deployments endpoint instead of the models endpoint
|
||||
}
|
||||
|
||||
// OllamaKeyConfig represents the Ollama-specific key configuration.
|
||||
// It allows each key to target a different Ollama server URL,
|
||||
// enabling per-key routing and round-robin load balancing across multiple Ollama instances.
|
||||
type OllamaKeyConfig struct {
|
||||
URL EnvVar `json:"url"` // Ollama server base URL (required, supports env. prefix)
|
||||
}
|
||||
|
||||
// SGLKeyConfig represents the SGLang-specific key configuration.
|
||||
// It allows each key to target a different SGLang server URL,
|
||||
// enabling per-key routing and round-robin load balancing across multiple SGLang instances.
|
||||
type SGLKeyConfig struct {
|
||||
URL EnvVar `json:"url"` // SGLang server base URL (required, supports env. prefix)
|
||||
}
|
||||
|
||||
// Account defines the interface for managing provider accounts and their configurations.
|
||||
// It provides methods to access provider-specific settings, API keys, and configurations.
|
||||
type Account interface {
|
||||
// GetConfiguredProviders returns a list of providers that are configured
|
||||
// in the account. This is used to determine which providers are available for use.
|
||||
GetConfiguredProviders() ([]ModelProvider, error)
|
||||
|
||||
// GetKeysForProvider returns the API keys configured for a specific provider.
|
||||
// The keys include their values, supported models, and weights for load balancing.
|
||||
// The context can carry data from any source that sets values before the Bifrost request,
|
||||
// including but not limited to plugin pre-hooks, application logic, or any in app middleware sharing the context.
|
||||
// This enables dynamic key selection based on any context values present during the request.
|
||||
GetKeysForProvider(ctx context.Context, providerKey ModelProvider) ([]Key, error)
|
||||
|
||||
// GetConfigForProvider returns the configuration for a specific provider.
|
||||
// This includes network settings, authentication details, and other provider-specific
|
||||
// configurations.
|
||||
GetConfigForProvider(providerKey ModelProvider) (*ProviderConfig, error)
|
||||
}
|
||||
Reference in New Issue
Block a user