Merge pull request #437 from Deepam02/feature/model-override-flag

Feature/model override flag

Author: dgageot

cmd/root/exec.go

@@ -14,6 +14,7 @@ func NewExecCmd() *cobra.Command {
 	cmd.PersistentFlags().StringVar(&workingDir, "working-dir", "", "Set the working directory for the session (applies to tools and relative paths)")
 	cmd.PersistentFlags().BoolVar(&autoApprove, "yolo", false, "Automatically approve all tool calls without prompting")
 	cmd.PersistentFlags().StringVar(&attachmentPath, "attach", "", "Attach an image file to the message")
+	cmd.PersistentFlags().StringArrayVar(&modelOverrides, "model", nil, "Override agent model: [agent=]provider/model (repeatable)")
 	cmd.PersistentFlags().BoolVar(&dryRun, "dry-run", false, "Initialize the agent without executing anything")
 	_ = cmd.PersistentFlags().MarkHidden("dry-run")
 

cmd/root/run.go

@@ -43,6 +43,7 @@ var (
 	remoteAddress  string
 	dryRun         bool
 	commandName    string
+	modelOverrides []string
 )
 
 const commandListSentinel = "__LIST__"
@@ -68,6 +69,7 @@ func NewRunCmd() *cobra.Command {
 	cmd.PersistentFlags().BoolVar(&useTUI, "tui", true, "Run the agent with a Terminal User Interface (TUI)")
 	cmd.PersistentFlags().StringVar(&remoteAddress, "remote", "", "Use remote runtime with specified address (only supported with TUI)")
 	cmd.PersistentFlags().StringVarP(&commandName, "command", "c", "", "Run a named command from the agent's commands section")
+	cmd.PersistentFlags().StringArrayVar(&modelOverrides, "model", nil, "Override agent model: [agent=]provider/model (repeatable)")
 	if f := cmd.PersistentFlags().Lookup("command"); f != nil {
 		// Allow `-c` without value to list available commands
 		f.NoOptDefVal = commandListSentinel
@@ -194,7 +196,7 @@ func doRunCommand(ctx context.Context, args []string, exec bool) error {
 			runConfig.RedirectURI = "http://localhost:8083/oauth-callback"
 		}
 
-		agents, err = teamloader.Load(ctx, agentFilename, runConfig)
+		agents, err = teamloader.LoadWithOverrides(ctx, agentFilename, runConfig, modelOverrides)
 		if err != nil {
 			return err
 		}

docs/USAGE.md

@@ -43,6 +43,11 @@ $ cagent run config.yaml --yolo           # Auto-accept all the tool calls
 $ cagent run config.yaml "First message"  # Start the conversation with the agent with a first message
 $ cagent run config.yaml -c df            # Run with a named command from YAML
 
+# Model Override Examples
+$ cagent run config.yaml --model anthropic/claude-sonnet-4-0    # Override all agents to use Claude
+$ cagent run config.yaml --model "agent1=openai/gpt-4o"         # Override specific agent
+$ cagent run config.yaml --model "agent1=openai/gpt-4o,agent2=anthropic/claude-sonnet-4-0"  # Multiple overrides
+
 # One off without TUI
 $ cagent exec config.yaml                 # Run the agent once, with default instructions
 $ cagent exec config.yaml "First message" # Run the agent once with instructions

pkg/config/overrides.go

@@ -0,0 +1,106 @@
+package config
+
+import (
+	"fmt"
+	"strings"
+
+	v2 "github.com/docker/cagent/pkg/config/v2"
+)
+
+// ApplyModelOverrides applies CLI model overrides to the configuration
+func ApplyModelOverrides(cfg *v2.Config, overrides []string) error {
+	for _, override := range overrides {
+		if err := applySingleOverride(cfg, override); err != nil {
+			return err
+		}
+	}
+
+	// After applying overrides, ensure new models are added to cfg.Models
+	return ensureModelsExist(cfg)
+}
+
+// applySingleOverride processes a single model override string
+func applySingleOverride(cfg *v2.Config, override string) error {
+	override = strings.TrimSpace(override)
+	if override == "" {
+		return nil // Skip empty overrides
+	}
+
+	// Handle comma-separated format: "agent1=model1,agent2=model2"
+	if strings.Contains(override, ",") {
+		for part := range strings.SplitSeq(override, ",") {
+			if err := applySingleOverride(cfg, part); err != nil {
+				return err
+			}
+		}
+		return nil
+	}
+
+	// Check if this is an agent-specific override (contains '=')
+	agentName, modelSpec, ok := strings.Cut(override, "=")
+	if ok {
+		agentName = strings.TrimSpace(agentName)
+		if agentName == "" {
+			return fmt.Errorf("empty agent name in override: %s", override)
+		}
+
+		modelSpec = strings.TrimSpace(modelSpec)
+		if modelSpec == "" {
+			return fmt.Errorf("empty model specification in override: %s", override)
+		}
+
+		// Apply to specific agent
+		agentConfig, exists := cfg.Agents[agentName]
+		if !exists {
+			return fmt.Errorf("unknown agent '%s'", agentName)
+		}
+
+		agentConfig.Model = modelSpec
+		cfg.Agents[agentName] = agentConfig
+	} else {
+		// Global override: apply to all agents
+		modelSpec := strings.TrimSpace(override)
+		if modelSpec == "" {
+			return fmt.Errorf("empty model specification")
+		}
+
+		for name := range cfg.Agents {
+			agentConfig := cfg.Agents[name]
+			agentConfig.Model = modelSpec
+			cfg.Agents[name] = agentConfig
+		}
+	}
+
+	return nil
+}
+
+// ensureModelsExist ensures that all models referenced by agents exist in cfg.Models
+// This handles inline model specs that may have been added via CLI overrides
+func ensureModelsExist(cfg *v2.Config) error {
+	if cfg.Models == nil {
+		cfg.Models = map[string]v2.ModelConfig{}
+	}
+
+	for agentName := range cfg.Agents {
+		agentConfig := cfg.Agents[agentName]
+
+		modelNames := strings.SplitSeq(agentConfig.Model, ",")
+		for modelName := range modelNames {
+			if _, exists := cfg.Models[modelName]; exists {
+				continue
+			}
+
+			providerName, model, ok := strings.Cut(modelName, "/")
+			if !ok {
+				return fmt.Errorf("agent '%s' references non-existent model '%s'", agentName, modelName)
+			}
+
+			cfg.Models[modelName] = v2.ModelConfig{
+				Provider: providerName,
+				Model:    model,
+			}
+		}
+	}
+
+	return nil
+}

pkg/teamloader/teamloader.go

@@ -97,6 +97,10 @@ func checkRequiredEnvVars(ctx context.Context, cfg *latest.Config, env environme
 }
 
 func Load(ctx context.Context, path string, runtimeConfig config.RuntimeConfig) (*team.Team, error) {
+	return LoadWithOverrides(ctx, path, runtimeConfig, nil)
+}
+
+func LoadWithOverrides(ctx context.Context, path string, runtimeConfig config.RuntimeConfig, modelOverrides []string) (*team.Team, error) {
 	fileName := filepath.Base(path)
 	parentDir := filepath.Dir(path)
 
@@ -123,6 +127,11 @@ func Load(ctx context.Context, path string, runtimeConfig config.RuntimeConfig)
 		return nil, err
 	}
 
+	// Apply model overrides from CLI flags before checking required env vars
+	if err := config.ApplyModelOverrides(cfg, modelOverrides); err != nil {
+		return nil, err
+	}
+
 	// Early check for required env vars before loading models and tools.
 	if err := checkRequiredEnvVars(ctx, cfg, env, runtimeConfig); err != nil {
 		return nil, err

pkg/teamloader/teamloader_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/docker/cagent/pkg/config"
+	latest "github.com/docker/cagent/pkg/config/v2"
 	"github.com/docker/cagent/pkg/environment"
 )
 
@@ -160,3 +161,153 @@ func collectExamples(t *testing.T) []string {
 
 	return files
 }
+
+func TestApplyModelOverrides(t *testing.T) {
+	tests := []struct {
+		name        string
+		agents      map[string]latest.AgentConfig
+		overrides   []string
+		expected    map[string]string // agent name -> expected model
+		expectError bool
+		errorMsg    string
+	}{
+		{
+			name: "global override",
+			agents: map[string]latest.AgentConfig{
+				"root":  {Model: "openai/gpt-4"},
+				"other": {Model: "anthropic/claude-3"},
+			},
+			overrides: []string{"google/gemini-pro"},
+			expected: map[string]string{
+				"root":  "google/gemini-pro",
+				"other": "google/gemini-pro",
+			},
+		},
+		{
+			name: "single per-agent override",
+			agents: map[string]latest.AgentConfig{
+				"root":  {Model: "openai/gpt-4"},
+				"other": {Model: "anthropic/claude-3"},
+			},
+			overrides: []string{"other=google/gemini-pro"},
+			expected: map[string]string{
+				"root":  "openai/gpt-4",
+				"other": "google/gemini-pro",
+			},
+		},
+		{
+			name: "multiple separate flags",
+			agents: map[string]latest.AgentConfig{
+				"root":  {Model: "openai/gpt-4"},
+				"other": {Model: "anthropic/claude-3"},
+			},
+			overrides: []string{"root=openai/gpt-5", "other=anthropic/claude-sonnet-4-0"},
+			expected: map[string]string{
+				"root":  "openai/gpt-5",
+				"other": "anthropic/claude-sonnet-4-0",
+			},
+		},
+		{
+			name: "comma-separated format",
+			agents: map[string]latest.AgentConfig{
+				"root":  {Model: "openai/gpt-4"},
+				"other": {Model: "anthropic/claude-3"},
+				"third": {Model: "google/gemini-pro"},
+			},
+			overrides: []string{"root=openai/gpt-5,other=anthropic/claude-sonnet-4-0"},
+			expected: map[string]string{
+				"root":  "openai/gpt-5",
+				"other": "anthropic/claude-sonnet-4-0",
+				"third": "google/gemini-pro",
+			},
+		},
+		{
+			name: "mixed formats",
+			agents: map[string]latest.AgentConfig{
+				"root":     {Model: "openai/gpt-4"},
+				"other":    {Model: "anthropic/claude-3"},
+				"third":    {Model: "google/gemini-pro"},
+				"reviewer": {Model: "openai/gpt-3.5-turbo"},
+			},
+			overrides: []string{"root=openai/gpt-5,other=anthropic/claude-4", "reviewer=google/gemini-1.5-pro"},
+			expected: map[string]string{
+				"root":     "openai/gpt-5",
+				"other":    "anthropic/claude-4",
+				"third":    "google/gemini-pro",
+				"reviewer": "google/gemini-1.5-pro",
+			},
+		},
+		{
+			name: "last override wins",
+			agents: map[string]latest.AgentConfig{
+				"root": {Model: "openai/gpt-4"},
+			},
+			overrides: []string{"root=openai/gpt-5", "root=anthropic/claude-4"},
+			expected: map[string]string{
+				"root": "anthropic/claude-4",
+			},
+		},
+		{
+			name: "unknown agent error",
+			agents: map[string]latest.AgentConfig{
+				"root": {Model: "openai/gpt-4"},
+			},
+			overrides:   []string{"nonexistent=openai/gpt-5"},
+			expectError: true,
+			errorMsg:    "unknown agent 'nonexistent'",
+		},
+		{
+			name: "empty model spec error",
+			agents: map[string]latest.AgentConfig{
+				"root": {Model: "openai/gpt-4"},
+			},
+			overrides:   []string{"root="},
+			expectError: true,
+			errorMsg:    "empty model specification in override: root=",
+		},
+		{
+			name: "empty global model spec is skipped",
+			agents: map[string]latest.AgentConfig{
+				"root": {Model: "openai/gpt-4"},
+			},
+			overrides: []string{""},
+			expected: map[string]string{
+				"root": "openai/gpt-4",
+			},
+		},
+		{
+			name: "whitespace handling",
+			agents: map[string]latest.AgentConfig{
+				"root":  {Model: "openai/gpt-4"},
+				"other": {Model: "anthropic/claude-3"},
+			},
+			overrides: []string{" root = openai/gpt-5 , other = anthropic/claude-4 "},
+			expected: map[string]string{
+				"root":  "openai/gpt-5",
+				"other": "anthropic/claude-4",
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			cfg := &latest.Config{
+				Agents: tt.agents,
+				Models: make(map[string]latest.ModelConfig),
+			}
+
+			err := config.ApplyModelOverrides(cfg, tt.overrides)
+
+			if tt.expectError {
+				require.ErrorContains(t, err, tt.errorMsg)
+			} else {
+				require.NoError(t, err)
+				for agentName, expectedModel := range tt.expected {
+					assert.Equal(t, expectedModel, cfg.Agents[agentName].Model)
+				}
+			}
+		})
+	}
+}