...

pkg/heroservices/openaiproxy/factory.go

@@ -0,0 +1,134 @@
+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
+}