init.go

package azure

import (
	"fmt"
	"github.com/spf13/viper"
	"github.com/stulzq/azure-openai-proxy/constant"
	"github.com/stulzq/azure-openai-proxy/util"
	"log"
	"net/url"
	"path/filepath"
	"strings"
)

const (
	AuthHeaderKey = "api-key"
)

var (
	C                     Config
	ModelDeploymentConfig = map[string]DeploymentConfig{}
)

func Init() error {
	var (
		apiVersion        string
		endpoint          string
		openaiModelMapper string
		err               error
	)

	apiVersion = viper.GetString(constant.ENV_AZURE_OPENAI_API_VER)
	endpoint = viper.GetString(constant.ENV_AZURE_OPENAI_ENDPOINT)
	openaiModelMapper = viper.GetString(constant.ENV_AZURE_OPENAI_MODEL_MAPPER)
	if endpoint != "" {
		if apiVersion == "" {
			apiVersion = "2024-02-01"
		}
		InitFromEnvironmentVariables(apiVersion, endpoint, openaiModelMapper)
	} else {
		if err = InitFromConfigFile(); err != nil {
			return err
		}
	}

	// Initialize MiniMax from environment variables
	InitMiniMaxFromEnv()

	// ensure apiBase likes /v1
	viper.SetDefault("api_base", "/v1")
	apiBase := viper.GetString("api_base")
	if !strings.HasPrefix(apiBase, "/") {
		apiBase = "/" + apiBase
	}
	if strings.HasSuffix(apiBase, "/") {
		apiBase = apiBase[:len(apiBase)-1]
	}
	viper.Set("api_base", apiBase)
	log.Printf("apiBase is: %s", apiBase)
	for _, itemConfig := range C.DeploymentConfig {
		u, err := url.Parse(itemConfig.Endpoint)
		if err != nil {
			return fmt.Errorf("parse endpoint error: %w", err)
		}
		itemConfig.EndpointUrl = u
		ModelDeploymentConfig[itemConfig.ModelName] = itemConfig
	}
	return err
}

func InitFromEnvironmentVariables(apiVersion, endpoint, openaiModelMapper string) {
	log.Println("Init from environment variables")
	if openaiModelMapper != "" {
		// openaiModelMapper example:
		// gpt-3.5-turbo=deployment_name_for_gpt_model,text-davinci-003=deployment_name_for_davinci_model
		for _, pair := range strings.Split(openaiModelMapper, ",") {
			info := strings.Split(pair, "=")
			if len(info) != 2 {
				log.Fatalf("error parsing %s, invalid value %s", constant.ENV_AZURE_OPENAI_MODEL_MAPPER, pair)
			}
			modelName, deploymentName := info[0], info[1]
			u, err := url.Parse(endpoint)
			if err != nil {
				log.Fatalf("parse endpoint error: %s", err.Error())
			}
			ModelDeploymentConfig[modelName] = DeploymentConfig{
				DeploymentName: deploymentName,
				ModelName:      modelName,
				Endpoint:       endpoint,
				EndpointUrl:    u,
				ApiKey:         "",
				ApiVersion:     apiVersion,
			}
		}
	}
}

// InitMiniMaxFromEnv initializes MiniMax provider from environment variables.
// Set MINIMAX_API_KEY to enable MiniMax support.
// Optionally set MINIMAX_ENDPOINT (default: https://api.minimax.io/v1) and
// MINIMAX_MODEL_MAPPER (e.g. MiniMax-M2.7=MiniMax-M2.7,MiniMax-M2.5-highspeed=MiniMax-M2.5-highspeed).
// If MINIMAX_MODEL_MAPPER is not set, default models (MiniMax-M2.7, MiniMax-M2.5, MiniMax-M2.5-highspeed) are registered.
func InitMiniMaxFromEnv() {
	apiKey := viper.GetString(constant.ENV_MINIMAX_API_KEY)
	if apiKey == "" {
		return
	}

	log.Println("Init MiniMax from environment variables")

	endpoint := viper.GetString(constant.ENV_MINIMAX_ENDPOINT)
	if endpoint == "" {
		endpoint = "https://api.minimax.io/v1"
	}

	modelMapper := viper.GetString(constant.ENV_MINIMAX_MODEL_MAPPER)

	u, err := url.Parse(endpoint)
	if err != nil {
		log.Fatalf("parse MiniMax endpoint error: %s", err.Error())
	}

	if modelMapper != "" {
		for _, pair := range strings.Split(modelMapper, ",") {
			info := strings.Split(pair, "=")
			modelName := strings.TrimSpace(info[0])
			ModelDeploymentConfig[modelName] = DeploymentConfig{
				DeploymentName: modelName,
				ModelName:      modelName,
				Endpoint:       endpoint,
				EndpointUrl:    u,
				ApiKey:         apiKey,
				ProviderType:   "minimax",
			}
		}
	} else {
		// Register default MiniMax models
		defaultModels := []string{"MiniMax-M2.7", "MiniMax-M2.5", "MiniMax-M2.5-highspeed"}
		for _, model := range defaultModels {
			ModelDeploymentConfig[model] = DeploymentConfig{
				DeploymentName: model,
				ModelName:      model,
				Endpoint:       endpoint,
				EndpointUrl:    u,
				ApiKey:         apiKey,
				ProviderType:   "minimax",
			}
		}
	}
}

func InitFromConfigFile() error {
	log.Println("Init from config file")

	configFile := viper.GetString("configFile")
	if configFile == "" {
		configFile = filepath.Join(util.GetWorkdir(), "config.yaml")
	} else if !filepath.IsAbs(configFile) {
		configFile = filepath.Join(util.GetWorkdir(), configFile)
	}

	viper.SetConfigType("yaml")
	viper.SetConfigFile(configFile)
	if err := viper.ReadInConfig(); err != nil {
		log.Printf("read config file error: %+v\n", err)
		return err
	}

	if err := viper.Unmarshal(&C); err != nil {
		log.Printf("unmarshal config file error: %+v\n", err)
		return err
	}
	for _, configItem := range C.DeploymentConfig {
		ModelDeploymentConfig[configItem.ModelName] = configItem
	}

	log.Println("read config file success")
	return nil
}