135 lines
3.4 KiB
Go
135 lines
3.4 KiB
Go
package proxy
|
|
|
|
import (
|
|
"errors"
|
|
"os"
|
|
"sync"
|
|
)
|
|
|
|
// Factory manages the proxy server and user configurations
|
|
type Factory struct {
|
|
// Config is the proxy server configuration
|
|
Config ProxyConfig
|
|
|
|
// userConfigs is a map of API keys to user configurations
|
|
userConfigs map[string]UserConfig
|
|
|
|
// Lock for concurrent access to userConfigs
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
// NewFactory creates a new proxy factory with the given configuration
|
|
func NewFactory(config ProxyConfig) *Factory {
|
|
// Check for OPENAIKEY environment variable and use it if available
|
|
if envKey := os.Getenv("OPENAIKEY"); envKey != "" {
|
|
config.DefaultOpenAIKey = envKey
|
|
}
|
|
|
|
return &Factory{
|
|
Config: config,
|
|
userConfigs: make(map[string]UserConfig),
|
|
}
|
|
}
|
|
|
|
// AddUserConfig adds or updates a user configuration with the associated API key
|
|
func (f *Factory) AddUserConfig(apiKey string, config UserConfig) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
f.userConfigs[apiKey] = config
|
|
}
|
|
|
|
// GetUserConfig retrieves a user configuration by API key
|
|
func (f *Factory) GetUserConfig(apiKey string) (UserConfig, error) {
|
|
f.mu.RLock()
|
|
defer f.mu.RUnlock()
|
|
|
|
config, exists := f.userConfigs[apiKey]
|
|
if !exists {
|
|
return UserConfig{}, errors.New("invalid API key")
|
|
}
|
|
|
|
return config, nil
|
|
}
|
|
|
|
// RemoveUserConfig removes a user configuration by API key
|
|
func (f *Factory) RemoveUserConfig(apiKey string) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
delete(f.userConfigs, apiKey)
|
|
}
|
|
|
|
// GetOpenAIKey returns the OpenAI API key to use for a given proxy API key
|
|
// Always returns the default OpenAI key from environment variable
|
|
func (f *Factory) GetOpenAIKey(proxyAPIKey string) string {
|
|
// Always use the default OpenAI key from environment variable
|
|
// This ensures that all requests to OpenAI use our key, not the user's key
|
|
return f.Config.DefaultOpenAIKey
|
|
}
|
|
|
|
// DecreaseBudget decreases a user's budget by the specified amount
|
|
// Returns error if the user doesn't have enough budget
|
|
func (f *Factory) DecreaseBudget(apiKey string, amount uint32) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
|
|
config, exists := f.userConfigs[apiKey]
|
|
if !exists {
|
|
return errors.New("invalid API key")
|
|
}
|
|
|
|
if config.Budget < amount {
|
|
return errors.New("insufficient budget")
|
|
}
|
|
|
|
config.Budget -= amount
|
|
f.userConfigs[apiKey] = config
|
|
return nil
|
|
}
|
|
|
|
// IncreaseBudget increases a user's budget by the specified amount
|
|
func (f *Factory) IncreaseBudget(apiKey string, amount uint32) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
|
|
config, exists := f.userConfigs[apiKey]
|
|
if !exists {
|
|
return errors.New("invalid API key")
|
|
}
|
|
|
|
config.Budget += amount
|
|
f.userConfigs[apiKey] = config
|
|
return nil
|
|
}
|
|
|
|
// CanAccessModel checks if a user can access a specific model
|
|
func (f *Factory) CanAccessModel(apiKey string, model string) bool {
|
|
f.mu.RLock()
|
|
defer f.mu.RUnlock()
|
|
|
|
config, exists := f.userConfigs[apiKey]
|
|
if !exists {
|
|
return false
|
|
}
|
|
|
|
// If no model groups are specified, allow access to all models
|
|
if len(config.ModelGroups) == 0 {
|
|
return true
|
|
}
|
|
|
|
// Check if the model is in any of the allowed model groups
|
|
// This is a placeholder - the actual implementation would depend on
|
|
// how model groups are defined and mapped to specific models
|
|
for _, group := range config.ModelGroups {
|
|
if group == "all" {
|
|
return true
|
|
}
|
|
// Add logic to check if model is in group
|
|
// For now we'll just check if the model contains the group name
|
|
if group == model {
|
|
return true
|
|
}
|
|
}
|
|
|
|
return false
|
|
}
|