cogito/options.go

package cogito

import (
	"context"

	"github.com/modelcontextprotocol/go-sdk/mcp"
	"github.com/mudler/cogito/prompt"
	"github.com/mudler/cogito/structures"
	"github.com/mudler/xlog"
	"github.com/sashabaranov/go-openai"
)

// MessageInjectionResult provides feedback about injected messages
type MessageInjectionResult struct {
	Count    int // Number of messages successfully injected
	Position int // Position in fragment where messages were added
}

// Options contains all configuration options for the Cogito agent
// It allows customization of behavior, tools, prompts, and execution parameters
type Options struct {
	prompts                           prompt.PromptMap
	maxIterations                     int
	tools                             Tools
	deepContext                       bool
	toolReasoner                      bool
	autoPlan                          bool
	planReEvaluator                   bool
	statusCallback, reasoningCallback func(string)
	gaps                              []string
	context                           context.Context
	infiniteExecution                 bool
	maxAttempts                       int
	feedbackCallback                  func() *Fragment
	toolCallCallback                  func(*ToolChoice, *SessionState) ToolCallDecision
	maxAdjustmentAttempts             int
	toolCallResultCallback            func(ToolStatus)
	strictGuidelines                  bool
	mcpSessions                       []*mcp.ClientSession
	guidelines                        Guidelines
	mcpPrompts                        bool
	mcpArgs                           map[string]string
	mcpToolFilter                     MCPToolFilter
	maxRetries                        int
	loopDetectionSteps                int
	forceReasoning                    bool
	forceReasoningTool                bool
	guidedTools                       bool
	parallelToolExecution             bool

	startWithAction []*ToolChoice

	sinkState bool

	sinkStateTool ToolDefinitionInterface

	// Message injection for concurrent conversation updates
	messageInjectionChan       chan openai.ChatCompletionMessage
	messageInjectionResultChan chan MessageInjectionResult

	// pendingWork, when set, keeps the loop parked (waiting on the
	// message-injection channel) while it returns true — for embedder-owned
	// background work that cogito's AgentManager knows nothing about.
	pendingWork func() bool

	// onPark, when set, fires immediately before the loop blocks on the
	// message-injection channel at a park gate; it receives the assistant
	// reply text that preceded the park ("" when the model produced none).
	// onResume, when set, fires immediately after an injected message wakes
	// the loop at a park gate.
	onPark   func(reply string)
	onResume func()

	// TODO-based iterative execution options
	reviewerLLMs        []LLM
	todoPersistencePath string
	todos               *structures.TODOList

	messagesManipulator func([]openai.ChatCompletionMessage) []openai.ChatCompletionMessage

	// Streaming callback for live token delivery
	streamCallback StreamCallback

	// Compaction options - automatic conversation compaction based on token count
	compactionThreshold    int // Token count threshold that triggers compaction (0 = disabled)
	compactionKeepMessages int // Number of recent messages to keep after compaction

	// AutoImprove options
	autoImproveState       *AutoImproveState
	autoImproveReviewerLLM LLM

	// Sub-agent spawning options
	enableAgentSpawning      bool
	agentManager             *AgentManager
	agentLLM                 LLM
	agentCompletionCallback  func(*AgentState)
	agentSpawnCallback       func(*AgentState)
	agentCompletionFormatter func(*AgentState) string
	agentDefinitions         []AgentDefinition
	agentLLMFactory          func(model string, temperature float32, metadata map[string]string) LLM
	agentDispatcher          AgentDispatcher
}

type Option func(*Options)

func defaultOptions() *Options {
	return &Options{
		maxIterations:          1,
		maxAttempts:            1,
		maxRetries:             5,
		loopDetectionSteps:     0,
		forceReasoning:         false,
		maxAdjustmentAttempts:  5,
		sinkStateTool:          &defaultSinkStateTool{},
		sinkState:              true,
		context:                context.Background(),
		statusCallback:         func(s string) {},
		reasoningCallback:      func(s string) {},
		compactionThreshold:    0,  // Disabled by default
		compactionKeepMessages: 10, // Keep 10 recent messages by default
	}
}

func (o *Options) Apply(opts ...Option) {
	for _, opt := range opts {
		opt(o)
	}
}

var (
	// EnableDeepContext enables full context to the LLM when chaining conversations
	// It might yield to better results to the cost of bigger context use.
	EnableDeepContext Option = func(o *Options) {
		o.deepContext = true
	}

	// EnableToolReasoner enables the reasoning about the need to call other tools
	// before each tool call, preventing calling more tools than necessary.
	EnableToolReasoner Option = func(o *Options) {
		o.toolReasoner = true
	}

	// DisableSinkState disables the use of a sink state
	// when the LLM decides that no tool is needed
	DisableSinkState Option = func(o *Options) {
		o.sinkState = false
	}

	// EnableInfiniteExecution enables infinite, long-term execution on Plans
	EnableInfiniteExecution Option = func(o *Options) {
		o.infiniteExecution = true
	}

	// EnableStrictGuidelines enforces cogito to pick tools only from the guidelines
	EnableStrictGuidelines Option = func(o *Options) {
		o.strictGuidelines = true
	}

	// EnableAutoPlan enables cogito to automatically use planning if needed
	EnableAutoPlan Option = func(o *Options) {
		o.autoPlan = true
	}

	// EnableAutoPlanReEvaluator enables cogito to automatically re-evaluate the need to use planning
	EnableAutoPlanReEvaluator Option = func(o *Options) {
		o.planReEvaluator = true
	}

	// EnableMCPPrompts enables the use of MCP prompts
	EnableMCPPrompts Option = func(o *Options) {
		o.mcpPrompts = true
	}

	// EnableGuidedTools enables filtering tools through guidance using their descriptions.
	// When no guidelines exist, creates virtual guidelines for all tools using their descriptions.
	// When guidelines exist, creates virtual guidelines for tools not in any guideline.
	EnableGuidedTools Option = func(o *Options) {
		o.guidedTools = true
	}

	// EnableParallelToolExecution enables parallel execution of multiple tool calls.
	// When enabled, the LLM can select multiple tools and they will be executed concurrently.
	EnableParallelToolExecution Option = func(o *Options) {
		o.parallelToolExecution = true
	}
)

// WithIterations allows to set the number of refinement iterations
func WithIterations(i int) func(o *Options) {
	return func(o *Options) {
		o.maxIterations = i
	}
}

func WithSinkState(tool ToolDefinitionInterface) func(o *Options) {
	return func(o *Options) {
		o.sinkState = true
		o.sinkStateTool = tool
	}
}

// WithPrompt allows to set a custom prompt for a given PromptType
func WithPrompt(t prompt.PromptType, p prompt.StaticPrompt) func(o *Options) {
	return func(o *Options) {
		if o.prompts == nil {
			o.prompts = make(prompt.PromptMap)
		}

		o.prompts[t] = p
	}
}

// WithTools allows to set the tools available to the Agent.
// Pass *ToolDefinition[T] instances - they will automatically generate openai.Tool via their Tool() method.
// Example: WithTools(&ToolDefinition[SearchArgs]{...}, &ToolDefinition[WeatherArgs]{...})
func WithTools(tools ...ToolDefinitionInterface) func(o *Options) {
	return func(o *Options) {
		o.tools = append(o.tools, tools...)
	}
}

// WithStatusCallback sets a callback function to receive status updates during execution
func WithStatusCallback(fn func(string)) func(o *Options) {
	return func(o *Options) {
		o.statusCallback = fn
	}
}

// WithGaps adds knowledge gaps that the agent should address
func WithGaps(gaps ...string) func(o *Options) {
	return func(o *Options) {
		o.gaps = append(o.gaps, gaps...)
	}
}

// WithContext sets the execution context for the agent
func WithContext(ctx context.Context) func(o *Options) {
	return func(o *Options) {
		o.context = ctx
	}
}

// WithMaxAttempts sets the maximum number of execution attempts
func WithMaxAttempts(i int) func(o *Options) {
	return func(o *Options) {
		o.maxAttempts = i
	}
}

// WithFeedbackCallback sets a callback to get continous feedback during execution of plans
func WithFeedbackCallback(fn func() *Fragment) func(o *Options) {
	return func(o *Options) {
		o.feedbackCallback = fn
	}
}

// WithToolCallBack allows to set a callback to intercept and modify tool calls before execution
// The callback receives the proposed tool choice and session state, and returns a ToolCallDecision
// that can approve, reject, provide adjustment feedback, or directly modify the tool choice
func WithToolCallBack(fn func(*ToolChoice, *SessionState) ToolCallDecision) func(o *Options) {
	return func(o *Options) {
		o.toolCallCallback = fn
	}
}

// WithMaxAdjustmentAttempts sets the maximum number of adjustment attempts when using tool call callbacks
// This prevents infinite loops when the user provides adjustment feedback
// Default is 5 attempts
func WithMaxAdjustmentAttempts(attempts int) func(o *Options) {
	return func(o *Options) {
		o.maxAdjustmentAttempts = attempts
	}
}

// WithToolCallResultCallback runs the callback on every tool result
func WithToolCallResultCallback(fn func(ToolStatus)) func(o *Options) {
	return func(o *Options) {
		o.toolCallResultCallback = fn
	}
}

// WithGuidelines adds behavioral guidelines for the agent to follow.
// The guildelines allows a more curated selection of the tool to use and only relevant are shown to the LLM during tool selection.
func WithGuidelines(guidelines ...Guideline) func(o *Options) {
	return func(o *Options) {
		o.guidelines = append(o.guidelines, guidelines...)
	}
}

// WithMCPs adds Model Context Protocol client sessions for external tool integration.
// When specified, the tools available in the MCPs will be available to the cogito pipelines
func WithMCPs(sessions ...*mcp.ClientSession) func(o *Options) {
	return func(o *Options) {
		o.mcpSessions = append(o.mcpSessions, sessions...)
	}
}

// WithMCPArgs sets the arguments for the MCP prompts
func WithMCPArgs(args map[string]string) func(o *Options) {
	return func(o *Options) {
		o.mcpArgs = args
	}
}

// WithMCPToolFilter installs a per-tool gate applied during the initial
// tool-discovery pass over each MCP session. fn is invoked once per
// (session, tool) pair after ListTools returns; tools for which fn
// returns false are dropped from the agent's discovered set and the LLM
// never sees them.
//
// The filter is not invoked on subsequent CallTool requests — the LLM
// can only request tools it learned about during discovery, so dropping
// at discovery time is sufficient. A nil fn (or no option set) means
// every tool from every session is exposed.
//
// Example: per-user enable/disable of remote MCP server tools.
//
//	enabled := map[*mcp.ClientSession]map[string]bool{ ... }
//	cogito.WithMCPToolFilter(func(s *mcp.ClientSession, tool string) bool {
//	    if e, ok := enabled[s]; ok {
//	        return e[tool]
//	    }
//	    return true // sessions not in the map are unfiltered
//	})
func WithMCPToolFilter(fn MCPToolFilter) func(o *Options) {
	return func(o *Options) {
		o.mcpToolFilter = fn
	}
}

// WithMessagesManipulator allows to manipulate the messages before they are sent to the LLM
// This is useful to add additional system messages or other context to the messages that needs to change during execution
func WithMessagesManipulator(fn func([]openai.ChatCompletionMessage) []openai.ChatCompletionMessage) func(o *Options) {
	return func(o *Options) {
		o.messagesManipulator = fn
	}
}

// WithMaxRetries sets the maximum number of retries for LLM calls
func WithMaxRetries(retries int) func(o *Options) {
	return func(o *Options) {
		o.maxRetries = retries
	}
}

// WithLoopDetection enables loop detection to prevent repeated tool calls
// If the same tool with the same parameters is called more than 'steps' times, it will be detected
func WithLoopDetection(steps int) func(o *Options) {
	return func(o *Options) {
		o.loopDetectionSteps = steps
	}
}

// WithForceReasoning enables forcing the LLM to reason before selecting tools
func WithForceReasoning() func(o *Options) {
	return func(o *Options) {
		o.forceReasoning = true
		o.sinkState = true
	}
}

// WithForceReasoningTool enables forcing the LLM to use the reasoning tool before selecting tools.
// This ensures structured output from the LLM instead of free text that might accidentally
// contain tool call JSON.
func WithForceReasoningTool() func(o *Options) {
	return func(o *Options) {
		o.forceReasoningTool = true
		o.forceReasoning = true
		o.sinkState = true
	}
}

// WithStartWithAction sets the initial tool choice to start with
func WithStartWithAction(tool ...*ToolChoice) func(o *Options) {
	return func(o *Options) {
		o.startWithAction = append(o.startWithAction, tool...)
	}
}

// WithReasoningCallback sets a callback function to receive reasoning updates during execution
func WithReasoningCallback(fn func(string)) func(o *Options) {
	return func(o *Options) {
		o.reasoningCallback = fn
	}
}

// WithReviewerLLM specifies a judge LLM for Planning with TODOs.
// When provided along with a plan, enables Planning with TODOs where the judge LLM
// reviews work after each iteration and decides whether goal execution is completed or needs rework.
func WithReviewerLLM(reviewerLLMs ...LLM) func(o *Options) {
	return func(o *Options) {
		o.reviewerLLMs = append(o.reviewerLLMs, reviewerLLMs...)
	}
}

// WithTODOPersistence enables file-based TODO persistence.
// TODOs will be saved to and loaded from the specified file path.
func WithTODOPersistence(path string) func(o *Options) {
	return func(o *Options) {
		o.todoPersistencePath = path
	}
}

// WithTODOs provides an in-memory TODO list for TODO-based iterative execution.
// If not provided, TODOs will be automatically generated from plan subtasks.
func WithTODOs(todoList *structures.TODOList) func(o *Options) {
	return func(o *Options) {
		o.todos = todoList
	}
}

// WithMessageInjectionChan sets a channel for injecting new messages during tool loop execution.
// Users can send messages through this channel, and they will be appended to the fragment after each iteration.
// Passing nil (default) disables this feature.
// When the channel is closed, no more messages will be accepted.
func WithMessageInjectionChan(ch chan openai.ChatCompletionMessage) func(o *Options) {
	return func(o *Options) {
		o.messageInjectionChan = ch
	}
}

// WithPendingWork makes the loop park (waiting on the message-injection
// channel) while fn returns true, even when cogito's own AgentManager has
// no running agents. For embedder-owned background work whose completion
// is delivered by injecting a message (see WithMessageInjectionChan).
// Pair it with WithMessageInjectionChan so there is a channel to wake on.
func WithPendingWork(fn func() bool) Option { return func(o *Options) { o.pendingWork = fn } }

// WithOnPark registers a callback fired immediately BEFORE the loop blocks on
// the message-injection channel at a park gate (i.e. when background work —
// cogito's own running agents or an embedder's WithPendingWork predicate — is
// still pending). The callback receives the assistant reply text that preceded
// the park — the no-tool text reply recorded in the fragment just before the
// loop blocked — or "" when the model produced none (e.g. a sink-state park).
// An embedder can use this to surface the parked reply and finalize the
// current assistant turn the instant the loop parks.
//
// Across a single run the loop may park and resume multiple times (e.g. several
// injected messages), so onPark may fire multiple times — that is expected.
func WithOnPark(fn func(reply string)) Option { return func(o *Options) { o.onPark = fn } }

// WithOnResume registers a callback fired immediately AFTER an injected message
// wakes the loop at a park gate (the resume path). It does NOT fire when the
// loop unblocks because the injection channel was closed or the context was
// cancelled. An embedder can use this to start a fresh assistant turn the
// instant the loop resumes.
//
// Across a single run the loop may park and resume multiple times (e.g. several
// injected messages), so onResume may fire multiple times — that is expected.
func WithOnResume(fn func()) Option { return func(o *Options) { o.onResume = fn } }

// WithMessageInjectionResultChan sets a channel to receive feedback about injected messages.
// For each message injection attempt, a MessageInjectionResult is sent back indicating:
// - Count: number of messages successfully added
// - Position: where in the fragment they were added
// - Err: error if validation or injection failed
// Passing nil (default) disables feedback.
func WithMessageInjectionResultChan(ch chan MessageInjectionResult) func(o *Options) {
	return func(o *Options) {
		o.messageInjectionResultChan = ch
	}
}

// WithCompactionThreshold sets the token count threshold that triggers automatic
// conversation compaction. When total tokens in the response >= threshold,
// the conversation will be compacted to stay within the limit.
// Set to 0 (default) to disable automatic compaction.
func WithCompactionThreshold(threshold int) func(o *Options) {
	return func(o *Options) {
		o.compactionThreshold = threshold
	}
}

// WithCompactionKeepMessages sets the number of recent messages to keep after
// compaction. Default is 10. This only applies when WithCompactionThreshold is set.
func WithCompactionKeepMessages(count int) func(o *Options) {
	return func(o *Options) {
		o.compactionKeepMessages = count
	}
}

// WithStreamCallback sets a callback to receive streaming events during execution.
// When set alongside a StreamingLLM, final answer generation will stream token-by-token.
func WithStreamCallback(fn StreamCallback) func(o *Options) {
	return func(o *Options) {
		o.streamCallback = fn
	}
}

// WithAutoImproveState enables the autoimproving feature.
// The provided state is mutated in-place: after ExecuteTools returns, state.SystemPrompt
// may have been updated by the review step. Persist and reuse the same pointer across calls.
func WithAutoImproveState(state *AutoImproveState) Option {
	return func(o *Options) {
		o.autoImproveState = state
	}
}

// WithAutoImproveReviewerLLM sets a separate LLM for the autoimprove review step.
// If not set, the same LLM passed to ExecuteTools is used.
func WithAutoImproveReviewerLLM(llm LLM) Option {
	return func(o *Options) {
		o.autoImproveReviewerLLM = llm
	}
}

// EnableAgentSpawning enables sub-agent spawning tools (spawn_agent, check_agent, get_agent_result).
// When enabled, the LLM can delegate tasks to sub-agents that run in foreground (blocking) or background (non-blocking).
var EnableAgentSpawning Option = func(o *Options) {
	o.enableAgentSpawning = true
}

// WithAgentManager provides an existing AgentManager for sharing across multiple ExecuteTools calls.
// If not provided and EnableAgentSpawning is set, a new AgentManager is created automatically.
func WithAgentManager(m *AgentManager) Option {
	return func(o *Options) {
		o.agentManager = m
	}
}

// WithAgentLLM sets a separate LLM for sub-agents to use.
// If not set, sub-agents share the parent's LLM.
func WithAgentLLM(llm LLM) Option {
	return func(o *Options) {
		o.agentLLM = llm
	}
}

// WithAgentDefinitions registers named sub-agent types (personas). spawn_agent
// can select one via its agent_type argument; the chosen definition supplies the
// system prompt, tool allow-list, model, temperature, and per-type execution limits.
func WithAgentDefinitions(defs ...AgentDefinition) Option {
	return func(o *Options) {
		o.agentDefinitions = defs
	}
}

// WithAgentLLMFactory sets a factory that builds an LLM for a sub-agent from a
// model name, temperature, and per-request metadata. Used to resolve
// per-agent-type or per-spawn model/metadata overrides while reusing the
// parent's endpoint/credentials.
func WithAgentLLMFactory(fn func(model string, temperature float32, metadata map[string]string) LLM) Option {
	return func(o *Options) {
		o.agentLLMFactory = fn
	}
}

// WithAgentCompletionCallback sets a callback that fires when any background sub-agent finishes.
// Useful for external monitoring or UI updates outside the LLM loop.
func WithAgentCompletionCallback(fn func(*AgentState)) Option {
	return func(o *Options) {
		o.agentCompletionCallback = fn
	}
}

// WithAgentSpawnCallback sets a callback that fires when a sub-agent starts
// (is registered and about to run), for both foreground and background spawns.
// Useful for UIs that show running agents. The AgentState has Status=running.
func WithAgentSpawnCallback(fn func(*AgentState)) Option {
	return func(o *Options) { o.agentSpawnCallback = fn }
}

// WithAgentCompletionFormatter overrides the message a finished background
// sub-agent injects into the parent loop. By default cogito injects a
// fixed prose notification ("Background agent <id> has completed…"); set
// this to control exactly what the parent LLM sees on wake — e.g. a clean
// structured summary, or a marker a host application can intercept and
// render itself rather than leaving the model to re-parse prose. The
// returned string is injected verbatim as a user-role message; returning
// "" injects an empty message (the default prose is only used when no
// formatter is set).
func WithAgentCompletionFormatter(fn func(*AgentState) string) Option {
	return func(o *Options) {
		o.agentCompletionFormatter = fn
	}
}

// WithAgentDispatcher installs an execution seam for spawned sub-agents. When
// set, cogito calls the dispatcher instead of running the sub-agent in-process
// (e.g. to dispatch the run to a remote worker), while still owning every part
// of the sub-agent lifecycle: registration, status transitions, the done
// channel, completion/spawn callbacks, completion-message injection, and
// foreground detach. A nil dispatcher (the default) preserves the in-process
// behavior exactly. Returning ErrDispatchFallback from the dispatcher makes
// cogito run the sub-agent in-process for that call.
func WithAgentDispatcher(d AgentDispatcher) Option {
	return func(o *Options) {
		o.agentDispatcher = d
	}
}

type defaultSinkStateTool struct{}

func (d *defaultSinkStateTool) Execute(args map[string]any) (string, any, error) {
	reasoning, ok := args["reasoning"].(string)
	if !ok {
		return "", nil, nil
	}
	xlog.Debug("[defaultSinkStateTool] Running default sink state tool", "reasoning", reasoning)
	return reasoning, reasoning, nil
}

func (d *defaultSinkStateTool) Tool() openai.Tool {
	return openai.Tool{
		Type: openai.ToolTypeFunction,
		Function: &openai.FunctionDefinition{
			Name:        "reply",
			Description: "This tool is used to reply to the user",
		},
	}
}