Merge pull request #6 from mostlygeek/multiline-config

Support multiline cmds in YAML configuration

config.example.yaml

@@ -4,36 +4,38 @@ healthCheckTimeout: 60
 
 models:
   "llama":
-    cmd: "models/llama-server-osx --port 8999 -m models/Llama-3.2-1B-Instruct-Q4_K_M.gguf"
-    proxy: "http://127.0.0.1:8999"
+    cmd: >
+      models/llama-server-osx
+      --port 8999
+      -m models/Llama-3.2-1B-Instruct-Q4_K_M.gguf
+    proxy: http://127.0.0.1:8999
 
     # list of model name aliases this llama.cpp instance can serve
     aliases:
-    - "gpt-4o-mini"
+    - gpt-4o-mini
 
     # check this path for a HTTP 200 response for the server to be ready
-    checkEndpoint: "/health"
+    checkEndpoint: /health
 
   "qwen":
-    cmd: "models/llama-server-osx --port 8999 -m models/Qwen2.5-1.5B-Instruct-Q4_K_M.gguf"
-    proxy: "http://127.0.0.1:8999"
+    cmd: models/llama-server-osx --port 8999 -m models/Qwen2.5-1.5B-Instruct-Q4_K_M.gguf
+    proxy: http://127.0.0.1:8999
     aliases:
-    - "gpt-3.5-turbo"
+    - gpt-3.5-turbo
 
   "simple":
     # example of setting environment variables
     env:
-      - "CUDA_VISIBLE_DEVICES=0,1"
-      - "env1=hello"
-    cmd: "build/simple-responder --port 8999"
-    proxy: "http://127.0.0.1:8999"
+      - CUDA_VISIBLE_DEVICES=0,1
+      - env1=hello
+    cmd: build/simple-responder --port 8999
+    proxy: http://127.0.0.1:8999
 
     # use "none" to skip check. Caution this may cause some requests to fail
     # until the upstream server is ready for traffic
-    checkEndpoint: "none"
+    checkEndpoint: none
 
   # don't use this, just for testing if things are broken
   "broken":
-    cmd: "models/llama-server-osx --port 8999 -m models/doesnotexist.gguf"
-    proxy: "http://127.0.0.1:8999"
-
+    cmd: models/llama-server-osx --port 8999 -m models/doesnotexist.gguf
+    proxy: http://127.0.0.1:8999

go.mod

@@ -2,4 +2,12 @@ module github.com/mostlygeek/llama-swap
 
 go 1.23.0
 
-require gopkg.in/yaml.v3 v3.0.1
+require (
+	github.com/stretchr/testify v1.9.0
+	gopkg.in/yaml.v3 v3.0.1
+)
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+)

go.sum

@@ -1,3 +1,9 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

proxy/config.go

@@ -1,7 +1,9 @@
 package proxy
 
 import (
+	"fmt"
 	"os"
+	"strings"
 
 	"gopkg.in/yaml.v3"
 )
@@ -14,6 +16,10 @@ type ModelConfig struct {
 	CheckEndpoint string   `yaml:"checkEndpoint"`
 }
 
+func (m *ModelConfig) SanitizedCommand() ([]string, error) {
+	return SanitizeCommand(m.Cmd)
+}
+
 type Config struct {
 	Models             map[string]ModelConfig `yaml:"models"`
 	HealthCheckTimeout int                    `yaml:"healthCheckTimeout"`
@@ -55,3 +61,19 @@ func LoadConfig(path string) (*Config, error) {
 
 	return &config, nil
 }
+
+func SanitizeCommand(cmdStr string) ([]string, error) {
+	// Remove trailing backslashes
+	cmdStr = strings.ReplaceAll(cmdStr, "\\ \n", " ")
+	cmdStr = strings.ReplaceAll(cmdStr, "\\\n", " ")
+
+	// Split the command into arguments
+	args := strings.Fields(cmdStr)
+
+	// Ensure the command is not empty
+	if len(args) == 0 {
+		return nil, fmt.Errorf("empty command")
+	}
+
+	return args, nil
+}

proxy/config_test.go

@@ -0,0 +1,126 @@
+package proxy
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLoadConfig(t *testing.T) {
+	// Create a temporary YAML file for testing
+	tempDir, err := os.MkdirTemp("", "test-config")
+	if err != nil {
+		t.Fatalf("Failed to create temporary directory: %v", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	tempFile := filepath.Join(tempDir, "config.yaml")
+	content := `models:
+  model1:
+    cmd: path/to/cmd --arg1 one
+    proxy: "http://localhost:8080"
+    aliases:
+      - "m1"
+      - "model-one"
+    env:
+      - "VAR1=value1"
+      - "VAR2=value2"
+    checkEndpoint: "/health"
+healthCheckTimeout: 15
+`
+
+	if err := os.WriteFile(tempFile, []byte(content), 0644); err != nil {
+		t.Fatalf("Failed to write temporary file: %v", err)
+	}
+
+	// Load the config and verify
+	config, err := LoadConfig(tempFile)
+	if err != nil {
+		t.Fatalf("Failed to load config: %v", err)
+	}
+
+	expected := &Config{
+		Models: map[string]ModelConfig{
+			"model1": {
+				Cmd:           "path/to/cmd --arg1 one",
+				Proxy:         "http://localhost:8080",
+				Aliases:       []string{"m1", "model-one"},
+				Env:           []string{"VAR1=value1", "VAR2=value2"},
+				CheckEndpoint: "/health",
+			},
+		},
+		HealthCheckTimeout: 15,
+	}
+
+	assert.Equal(t, expected, config)
+}
+
+func TestModelConfigSanitizedCommand(t *testing.T) {
+	config := &ModelConfig{
+		Cmd: `python model1.py \
+    --arg1 value1 \
+    --arg2 value2`,
+	}
+
+	args, err := config.SanitizedCommand()
+	assert.NoError(t, err)
+	assert.Equal(t, []string{"python", "model1.py", "--arg1", "value1", "--arg2", "value2"}, args)
+}
+
+func TestFindConfig(t *testing.T) {
+	config := &Config{
+		Models: map[string]ModelConfig{
+			"model1": {
+				Cmd:           "python model1.py",
+				Proxy:         "http://localhost:8080",
+				Aliases:       []string{"m1", "model-one"},
+				Env:           []string{"VAR1=value1", "VAR2=value2"},
+				CheckEndpoint: "/health",
+			},
+			"model2": {
+				Cmd:           "python model2.py",
+				Proxy:         "http://localhost:8081",
+				Aliases:       []string{"m2", "model-two"},
+				Env:           []string{"VAR3=value3", "VAR4=value4"},
+				CheckEndpoint: "/status",
+			},
+		},
+		HealthCheckTimeout: 10,
+	}
+
+	// Test finding a model by its name
+	modelConfig, found := config.FindConfig("model1")
+	assert.True(t, found)
+	assert.Equal(t, config.Models["model1"], modelConfig)
+
+	// Test finding a model by its alias
+	modelConfig, found = config.FindConfig("m1")
+	assert.True(t, found)
+	assert.Equal(t, config.Models["model1"], modelConfig)
+
+	// Test finding a model that does not exist
+	modelConfig, found = config.FindConfig("model3")
+	assert.False(t, found)
+	assert.Equal(t, ModelConfig{}, modelConfig)
+}
+
+func TestSanitizeCommand(t *testing.T) {
+	// Test a simple command
+	args, err := SanitizeCommand("python model1.py")
+	assert.NoError(t, err)
+	assert.Equal(t, []string{"python", "model1.py"}, args)
+
+	// Test a command with spaces and newlines
+	args, err = SanitizeCommand(`python model1.py \
+    --arg1 value1 \
+    --arg2 value2`)
+	assert.NoError(t, err)
+	assert.Equal(t, []string{"python", "model1.py", "--arg1", "value1", "--arg2", "value2"}, args)
+
+	// Test an empty command
+	args, err = SanitizeCommand("")
+	assert.Error(t, err)
+	assert.Nil(t, args)
+}

proxy/manager.go

@@ -65,13 +65,16 @@ func (pm *ProxyManager) swapModel(requestedModel string) error {
 
 	pm.currentConfig = modelConfig
 
-	args := strings.Fields(modelConfig.Cmd)
+	args, err := modelConfig.SanitizedCommand()
+	if err != nil {
+		return fmt.Errorf("unable to get sanitized command: %v", err)
+	}
 	cmd := exec.Command(args[0], args[1:]...)
 	cmd.Stdout = os.Stdout
 	cmd.Stderr = os.Stderr
 	cmd.Env = modelConfig.Env
 
-	err := cmd.Start()
+	err = cmd.Start()
 	if err != nil {
 		return err
 	}