envconfig.go

package envconfig

import (
	"fmt"
	"log/slog"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/docker/model-runner/pkg/logging"
)

// Var returns an environment variable stripped of leading/trailing quotes and spaces.
func Var(key string) string {
	return strings.Trim(strings.TrimSpace(os.Getenv(key)), "\"'")
}

// String returns a lazy string accessor for the given environment variable.
func String(key string) func() string {
	return func() string {
		return Var(key)
	}
}

// BoolWithDefault returns a lazy bool accessor for the given environment variable,
// allowing a caller-specified default. If the variable is set but cannot be parsed
// as a bool, the defaultValue is returned.
func BoolWithDefault(key string) func(defaultValue bool) bool {
	return func(defaultValue bool) bool {
		if s := Var(key); s != "" {
			b, err := strconv.ParseBool(s)
			if err != nil {
				return defaultValue
			}
			return b
		}
		return defaultValue
	}
}

// Bool returns a lazy bool accessor that defaults to false when the variable is unset.
func Bool(key string) func() bool {
	withDefault := BoolWithDefault(key)
	return func() bool {
		return withDefault(false)
	}
}

// LogLevel reads LOG_LEVEL and returns the corresponding slog.Level.
func LogLevel() slog.Level {
	return logging.ParseLevel(Var("LOG_LEVEL"))
}

// AllowedOrigins returns a list of CORS-allowed origins. It reads DMR_ORIGINS
// and always appends default localhost/127.0.0.1/0.0.0.0 entries on http and
// https with wildcard ports.
func AllowedOrigins() (origins []string) {
	if s := Var("DMR_ORIGINS"); s != "" {
		for _, o := range strings.Split(s, ",") {
			if trimmed := strings.TrimSpace(o); trimmed != "" {
				origins = append(origins, trimmed)
			}
		}
	}

	for _, host := range []string{"localhost", "127.0.0.1", "0.0.0.0"} {
		origins = append(origins,
			fmt.Sprintf("http://%s", host),
			fmt.Sprintf("https://%s", host),
			fmt.Sprintf("http://%s", net.JoinHostPort(host, "*")),
			fmt.Sprintf("https://%s", net.JoinHostPort(host, "*")),
		)
	}

	return origins
}

// SocketPath returns the Unix socket path for the model runner.
// Configured via MODEL_RUNNER_SOCK; defaults to "model-runner.sock".
func SocketPath() string {
	if s := Var("MODEL_RUNNER_SOCK"); s != "" {
		return s
	}
	return "model-runner.sock"
}

// ModelsPath returns the directory where models are stored.
// Configured via MODELS_PATH; defaults to ~/.docker/models.
func ModelsPath() (string, error) {
	if s := Var("MODELS_PATH"); s != "" {
		return s, nil
	}
	home, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}
	return filepath.Join(home, ".docker", "models"), nil
}

// TCPPort returns the optional TCP port for the model runner HTTP server.
// Configured via MODEL_RUNNER_PORT; empty string means use Unix socket.
func TCPPort() string {
	return Var("MODEL_RUNNER_PORT")
}

// LlamaServerPath returns the path to the directory containing the llama.cpp server binary.
// Configured via LLAMA_SERVER_PATH; defaults to ~/.docker/model-runner/llama.cpp/bin.
func LlamaServerPath() (string, error) {
	if s := Var("LLAMA_SERVER_PATH"); s != "" {
		return s, nil
	}
	home, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}
	return filepath.Join(home, ".docker", "model-runner", "llama.cpp", "bin"), nil
}

// LlamaArgs returns custom arguments to pass to the llama.cpp server.
// Configured via LLAMA_ARGS.
func LlamaArgs() string {
	return Var("LLAMA_ARGS")
}

// DisableServerUpdate is true when DISABLE_SERVER_UPDATE is set to a truthy value.
var DisableServerUpdate = Bool("DISABLE_SERVER_UPDATE")

// LlamaServerVersion returns a specific llama.cpp server version to pin.
// Configured via LLAMA_SERVER_VERSION; empty string means use the bundled version.
func LlamaServerVersion() string {
	return Var("LLAMA_SERVER_VERSION")
}

// VLLMServerPath returns the optional path to the vLLM server binary.
// Configured via VLLM_SERVER_PATH.
func VLLMServerPath() string {
	return Var("VLLM_SERVER_PATH")
}

// SGLangServerPath returns the optional path to the SGLang server binary.
// Configured via SGLANG_SERVER_PATH.
func SGLangServerPath() string {
	return Var("SGLANG_SERVER_PATH")
}

// MLXServerPath returns the optional path to the MLX server binary.
// Configured via MLX_SERVER_PATH.
func MLXServerPath() string {
	return Var("MLX_SERVER_PATH")
}

// DiffusersServerPath returns the optional path to the Diffusers server binary.
// Configured via DIFFUSERS_SERVER_PATH.
func DiffusersServerPath() string {
	return Var("DIFFUSERS_SERVER_PATH")
}

// VLLMMetalServerPath returns the optional path to the vLLM Metal server binary.
// Configured via VLLM_METAL_SERVER_PATH.
func VLLMMetalServerPath() string {
	return Var("VLLM_METAL_SERVER_PATH")
}

// DisableMetrics is true when DISABLE_METRICS is set to a truthy value (e.g. "1").
var DisableMetrics = Bool("DISABLE_METRICS")

// TLSEnabled is true when MODEL_RUNNER_TLS_ENABLED is set to a truthy value.
var TLSEnabled = Bool("MODEL_RUNNER_TLS_ENABLED")

// TLSPort returns the TLS listener port.
// Configured via MODEL_RUNNER_TLS_PORT; defaults to "12444".
func TLSPort() string {
	if s := Var("MODEL_RUNNER_TLS_PORT"); s != "" {
		return s
	}
	return "12444"
}

// TLSCert returns the path to the TLS certificate file.
// Configured via MODEL_RUNNER_TLS_CERT.
func TLSCert() string {
	return Var("MODEL_RUNNER_TLS_CERT")
}

// TLSKey returns the path to the TLS private key file.
// Configured via MODEL_RUNNER_TLS_KEY.
func TLSKey() string {
	return Var("MODEL_RUNNER_TLS_KEY")
}

// TLSAutoCert is true (default) unless MODEL_RUNNER_TLS_AUTO_CERT is set to a falsy value.
// Call as TLSAutoCert(true) to get the default-true behaviour.
var TLSAutoCert = BoolWithDefault("MODEL_RUNNER_TLS_AUTO_CERT")