rony-llm-agent/pkg/agent/loop.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"iter"
	"strings"
	"time"

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
	"github.com/VictorVargas/rony-llm-agent/pkg/persona"
	"github.com/VictorVargas/rony-llm-agent/pkg/tools"
)

// MaxIterations is the default maximum number of agent iterations.
const DefaultMaxIterations = 50

// MaxToolOutputBytes is the default maximum size for tool output in bytes.
const DefaultMaxToolOutputBytes = 50 * 1024 // 50KB

// Sandbox defines the interface for filesystem sandbox operations.
type Sandbox interface {
	ValidateToolCall(tool tools.Tool, call llm.ToolCall) error
}

// Approver is called before executing tools with Ask permission.
// Return true to allow, false to deny.
type Approver func(tool tools.Tool, call llm.ToolCall) bool

// OnIterationHook is called after each iteration.
type OnIterationHook func(Iteration)

// Config holds the dependencies and settings for the agent loop.
type Config struct {
	LLM                llm.LLMClient
	Persona            persona.Persona
	Tools              tools.Registry
	Sandbox            Sandbox
	MaxIters           int
	Approver           Approver
	OnIteration        OnIterationHook
	ToolTimeout        time.Duration
	ChatTemplateKwargs map[string]any // passed to the LLM provider (e.g. Qwen enable_thinking)
}

// Iteration represents a single cycle of the agent loop.
type Iteration struct {
	Number     int
	ToolCalls  []llm.ToolCall
	ToolsUsed  int
	Duration   time.Duration
}

// Response is the final output of the agent loop.
type Response struct {
	Content     string
	ToolCalls   []llm.ToolCall
	Iterations  int
	Duration    time.Duration
	TokenUsage  llm.TokenUsage
}

// Loop is the main agent loop that orchestrates LLM calls and tool execution.
type Loop struct {
	cfg Config
}

// New creates a new Loop with the given configuration.
func New(cfg Config) *Loop {
	if cfg.MaxIters == 0 {
		cfg.MaxIters = DefaultMaxIterations
	}
	return &Loop{cfg: cfg}
}

// Run executes the agent loop and returns the final response.
// Optional history messages are appended after the system prompt and before
// the new user input.
func (l *Loop) Run(ctx context.Context, input string, history ...llm.Message) (Response, error) {
	start := time.Now()

	messages := l.buildInitialMessages(input, history)
	var finalContent string
	var allToolCalls []llm.ToolCall
	var totalUsage llm.TokenUsage
	iterations := 0

	for iterations < l.cfg.MaxIters {
		iterations++

		resp, err := l.cfg.LLM.Generate(ctx, llm.CompletionRequest{
			Messages:           messages,
			Tools:              l.getToolSchemas(),
			ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
		})
		if err != nil {
			return Response{}, fmt.Errorf("LLM generate failed: %w", err)
		}

		totalUsage.InputTokens += resp.Usage.InputTokens
		totalUsage.OutputTokens += resp.Usage.OutputTokens
		totalUsage.TotalTokens += resp.Usage.TotalTokens

		if len(resp.ToolCalls) == 0 {
			finalContent = resp.Content
			break
		}

		for _, call := range resp.ToolCalls {
			result, err := l.executeTool(ctx, call)
			if err != nil {
				messages = append(messages, llm.Message{
					Role:    llm.RoleTool,
					Content: fmt.Sprintf("Error: %v", err),
				})
				continue
			}
			messages = append(messages, llm.Message{
				Role:    llm.RoleTool,
				Content: result.Content,
			})
			allToolCalls = append(allToolCalls, call)
		}
	}

	duration := time.Since(start)
	if iterations >= l.cfg.MaxIters {
		return Response{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters)
	}

	return Response{
		Content:     finalContent,
		ToolCalls:   allToolCalls,
		Iterations:  iterations,
		Duration:    duration,
		TokenUsage:  totalUsage,
	}, nil
}

// RunStream executes the agent loop with streaming output.
// Optional history messages are appended after the system prompt and before
// the new user input.
func (l *Loop) RunStream(ctx context.Context, input string, history ...llm.Message) iter.Seq2[llm.StreamChunk, error] {
	return func(yield func(llm.StreamChunk, error) bool) {
		messages := l.buildInitialMessages(input, history)
		iterations := 0

		for iterations < l.cfg.MaxIters {
			iterations++

			stream := l.cfg.LLM.Stream(ctx, llm.CompletionRequest{
				Messages:           messages,
				Tools:              l.getToolSchemas(),
				ChatTemplateKwargs: l.cfg.ChatTemplateKwargs,
			})

			var hasToolCalls bool
			var responseBuilder strings.Builder
			for chunk, err := range stream {
				if err != nil {
					yield(llm.StreamChunk{}, err)
					return
				}

				if len(chunk.ToolCalls) > 0 {
					hasToolCalls = true
					for _, tc := range chunk.ToolCalls {
						result, err := l.executeTool(ctx, tc)
						if err != nil {
							messages = append(messages, llm.Message{
								Role:    llm.RoleTool,
								Content: fmt.Sprintf("Error: %v", err),
							})
							continue
						}
						messages = append(messages, llm.Message{
							Role:    llm.RoleTool,
							Content: result.Content,
						})
					}
				}

				// A trailing usage-only chunk (no Delta/ReasoningDelta, per
				// providers that report token usage in a separate final
				// event) must still be forwarded, or callers can never see
				// real token counts.
				hasContent := chunk.Delta != "" || chunk.ReasoningDelta != ""
				hasUsage := chunk.Usage.TotalTokens > 0
				if !hasToolCalls && (hasContent || hasUsage) {
					responseBuilder.WriteString(chunk.Delta)
					if !yield(chunk, nil) {
						return
					}
				}
			}

			if !hasToolCalls {
				return
			}
		}

		yield(llm.StreamChunk{}, fmt.Errorf("max iterations (%d) reached", l.cfg.MaxIters))
	}
}

func (l *Loop) buildInitialMessages(input string, history []llm.Message) []llm.Message {
	systemPrompt := persona.AssembleSystemPrompt(l.cfg.Persona, "")
	messages := make([]llm.Message, 0, len(history)+2)
	messages = append(messages, llm.Message{Role: llm.RoleSystem, Content: systemPrompt})
	messages = append(messages, history...)
	messages = append(messages, llm.Message{Role: llm.RoleUser, Content: input})
	return messages
}

func (l *Loop) getToolSchemas() []json.RawMessage {
	var schemas []json.RawMessage
	for _, tool := range l.cfg.Tools.List() {
		def := map[string]interface{}{
			"type": "function",
			"function": map[string]interface{}{
				"name":        tool.Name,
				"description": tool.Description,
				"parameters":  json.RawMessage(tool.InputSchema),
			},
		}
		data, _ := json.Marshal(def)
		schemas = append(schemas, data)
	}
	return schemas
}

func (l *Loop) executeTool(ctx context.Context, call llm.ToolCall) (tools.ToolResult, error) {
	tool, found := l.cfg.Tools.Get(call.Name)
	if !found {
		return tools.ToolResult{IsError: true}, fmt.Errorf("tool %q not found", call.Name)
	}

	if tool.Permission == tools.Ask && l.cfg.Approver != nil && !l.cfg.Approver(tool, call) {
		return tools.ToolResult{Content: "Tool execution denied by user"}, nil
	}

	if l.cfg.Sandbox != nil {
		if err := l.cfg.Sandbox.ValidateToolCall(tool, call); err != nil {
			return tools.ToolResult{IsError: true}, fmt.Errorf("sandbox violation: %w", err)
		}
	}

	if l.cfg.OnIteration != nil {
		l.cfg.OnIteration(Iteration{
			Number:    1,
			ToolCalls: []llm.ToolCall{call},
			ToolsUsed: 1,
		})
	}

	result, err := tool.Handler(ctx, call.Arguments)
	if err != nil {
		return tools.ToolResult{IsError: true, Content: fmt.Sprintf("Tool error: %v", err)}, nil
	}

	return result, nil
}