kit/internal/agent/factory.go

package agent

import (
	"context"
	"fmt"

	"charm.land/fantasy"

	"github.com/mark3labs/kit/internal/config"
	"github.com/mark3labs/kit/internal/models"
	"github.com/mark3labs/kit/internal/tools"
)

// SpinnerFunc is a function type for showing spinners during agent creation.
// It executes the provided function while displaying an animated spinner.
type SpinnerFunc func(fn func() error) error

// AgentCreationOptions contains options for creating an agent.
// It extends AgentConfig with UI-related options for showing progress during creation.
type AgentCreationOptions struct {
	// ModelConfig specifies the LLM provider and model to use
	ModelConfig *models.ProviderConfig
	// MCPConfig contains MCP server configurations
	MCPConfig *config.Config
	// SystemPrompt is the initial system message for the agent
	SystemPrompt string
	// MaxSteps limits the number of tool calls (0 for unlimited)
	MaxSteps int
	// StreamingEnabled controls whether responses are streamed
	StreamingEnabled bool
	// ShowSpinner indicates whether to show a spinner for Ollama models during loading
	ShowSpinner bool // For Ollama models
	// Quiet suppresses the spinner even if ShowSpinner is true
	Quiet bool // Skip spinner if quiet
	// SpinnerFunc is the function to show spinner, provided by the caller
	SpinnerFunc SpinnerFunc // Function to show spinner (provided by caller)
	// DebugLogger is an optional logger for debugging MCP communications
	DebugLogger tools.DebugLogger // Optional debug logger
	// AuthHandler handles OAuth authorization for remote MCP servers
	AuthHandler tools.MCPAuthHandler
	// TokenStoreFactory, if non-nil, creates a custom token store for each
	// remote MCP server's OAuth tokens. When nil, the default file-based
	// token store is used.
	TokenStoreFactory tools.TokenStoreFactory
	// CoreTools overrides the default core tool set. If empty, core.AllTools()
	// is used.
	CoreTools []fantasy.AgentTool
	// DisableCoreTools, when true, prevents loading any core tools.
	// If both DisableCoreTools is true and CoreTools is empty, the agent
	// will have no tools (useful for simple chat completions).
	DisableCoreTools bool
	// ToolWrapper wraps the combined tool list before agent creation.
	ToolWrapper func([]fantasy.AgentTool) []fantasy.AgentTool
	// ExtraTools are additional tools to include (e.g. from extensions).
	ExtraTools []fantasy.AgentTool
	// OnMCPServerLoaded, if non-nil, is called when each MCP server finishes
	// loading (successfully or with error). Called from the background goroutine.
	OnMCPServerLoaded func(serverName string, toolCount int, err error)
	// MCPTaskConfig configures task-augmented tools/call execution.
	MCPTaskConfig tools.MCPTaskConfig
}

// CreateAgent creates an agent with optional spinner for Ollama models.
// It shows a loading spinner for Ollama models if ShowSpinner is true and not in quiet mode.
// Returns the created agent or an error if creation fails.
func CreateAgent(ctx context.Context, opts *AgentCreationOptions) (*Agent, error) {
	agentConfig := &AgentConfig{
		ModelConfig:       opts.ModelConfig,
		MCPConfig:         opts.MCPConfig,
		SystemPrompt:      opts.SystemPrompt,
		MaxSteps:          opts.MaxSteps,
		StreamingEnabled:  opts.StreamingEnabled,
		DebugLogger:       opts.DebugLogger,
		AuthHandler:       opts.AuthHandler,
		TokenStoreFactory: opts.TokenStoreFactory,
		CoreTools:         opts.CoreTools,
		DisableCoreTools:  opts.DisableCoreTools,
		ToolWrapper:       opts.ToolWrapper,
		ExtraTools:        opts.ExtraTools,
		OnMCPServerLoaded: opts.OnMCPServerLoaded,
		MCPTaskConfig:     opts.MCPTaskConfig,
	}

	var agent *Agent
	var err error

	// Show spinner for Ollama models if requested and not quiet
	parsedProvider, _, _ := models.ParseModelString(opts.ModelConfig.ModelString)
	if opts.ShowSpinner && parsedProvider == "ollama" && !opts.Quiet && opts.SpinnerFunc != nil {
		err = opts.SpinnerFunc(func() error {
			agent, err = NewAgent(ctx, agentConfig)
			return err
		})
	} else {
		agent, err = NewAgent(ctx, agentConfig)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to create agent: %v", err)
	}

	return agent, nil
}