agentu/internal/agent/agent.go

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"sort"
	"strings"
	"time"

	"agentu/internal/llm"
	"agentu/internal/tools"
)

const defaultMaxToolIterations = 8

type Agent struct {
	provider          llm.Provider
	providerName      string
	model             string
	thinking          string
	thinkingParam     string
	thinkingEnabled   bool
	systemPrompt      string
	toolRegistry      *tools.Registry
	toolTimeout       time.Duration
	maxToolIterations int
	messages          []llm.Message
}

type Options struct {
	Provider          llm.Provider
	ProviderName      string
	Model             string
	Thinking          string
	ThinkingParam     string
	ThinkingEnabled   bool
	SystemPrompt      string
	ToolRegistry      *tools.Registry
	ToolTimeout       time.Duration
	MaxToolIterations int
}

func New(opts Options) *Agent {
	maxToolIterations := opts.MaxToolIterations
	if maxToolIterations <= 0 {
		maxToolIterations = defaultMaxToolIterations
	}
	toolTimeout := opts.ToolTimeout
	if toolTimeout <= 0 {
		toolTimeout = 30 * time.Second
	}
	a := &Agent{
		provider:          opts.Provider,
		providerName:      opts.ProviderName,
		model:             opts.Model,
		thinking:          opts.Thinking,
		thinkingParam:     opts.ThinkingParam,
		thinkingEnabled:   opts.ThinkingEnabled,
		systemPrompt:      opts.SystemPrompt,
		toolRegistry:      opts.ToolRegistry,
		toolTimeout:       toolTimeout,
		maxToolIterations: maxToolIterations,
	}
	a.Clear()
	return a
}

type RuntimeOptions struct {
	Provider        llm.Provider
	ProviderName    string
	Model           string
	Thinking        string
	ThinkingParam   string
	ThinkingEnabled bool
}

type RuntimeInfo struct {
	ProviderName    string
	Model           string
	Thinking        string
	ThinkingEnabled bool
}

func (a *Agent) SetRuntime(opts RuntimeOptions) {
	if opts.Provider != nil {
		a.provider = opts.Provider
	}
	if opts.ProviderName != "" {
		a.providerName = opts.ProviderName
	}
	if opts.Model != "" {
		a.model = opts.Model
	}
	a.thinking = opts.Thinking
	a.thinkingParam = opts.ThinkingParam
	a.thinkingEnabled = opts.ThinkingEnabled
}

func (a *Agent) RuntimeInfo() RuntimeInfo {
	return RuntimeInfo{
		ProviderName:    a.providerName,
		Model:           a.model,
		Thinking:        a.thinking,
		ThinkingEnabled: a.thinkingEnabled,
	}
}

func (a *Agent) Clear() {
	a.messages = nil
	if strings.TrimSpace(a.systemPrompt) != "" {
		a.messages = append(a.messages, llm.Message{
			Role:    llm.RoleSystem,
			Content: a.systemPrompt,
		})
	}
}

func (a *Agent) RunTurn(ctx context.Context, input string, out io.Writer, logs io.Writer) error {
	input = strings.TrimSpace(input)
	if input == "" {
		return nil
	}
	a.messages = append(a.messages, llm.Message{Role: llm.RoleUser, Content: input})

	for iteration := 0; iteration <= a.maxToolIterations; iteration++ {
		content, calls, err := a.streamAssistant(ctx, out)
		if err != nil {
			return err
		}
		if len(calls) == 0 {
			a.messages = append(a.messages, llm.Message{
				Role:    llm.RoleAssistant,
				Content: content,
			})
			return nil
		}

		a.messages = append(a.messages, llm.Message{
			Role:      llm.RoleAssistant,
			Content:   content,
			ToolCalls: calls,
		})
		if content != "" {
			fmt.Fprintln(out)
		}

		for _, call := range calls {
			result := a.executeTool(ctx, call, logs)
			a.messages = append(a.messages, llm.Message{
				Role:       llm.RoleTool,
				ToolCallID: call.ID,
				Content:    result,
			})
		}
	}

	return fmt.Errorf("tool iteration limit reached: %d", a.maxToolIterations)
}

func (a *Agent) streamAssistant(ctx context.Context, out io.Writer) (string, []llm.ToolCall, error) {
	var content strings.Builder
	calls := make(map[int]*toolCallBuilder)

	req := llm.ChatRequest{
		Model:    a.model,
		Messages: append([]llm.Message(nil), a.messages...),
		Tools:    a.toolDefinitions(),
	}
	if a.thinkingEnabled {
		thinkingParam := a.thinkingParam
		if thinkingParam == "" {
			thinkingParam = "thinking"
		}
		req.Extra = map[string]any{thinkingParam: a.thinking}
	}
	if len(req.Tools) > 0 {
		req.ToolChoice = "auto"
	}

	err := a.provider.ChatStream(ctx, req, func(event llm.StreamEvent) error {
		if event.Content != "" {
			content.WriteString(event.Content)
			if out != nil {
				if _, err := io.WriteString(out, event.Content); err != nil {
					return err
				}
			}
		}
		for _, delta := range event.ToolCalls {
			builder := calls[delta.Index]
			if builder == nil {
				builder = &toolCallBuilder{index: delta.Index}
				calls[delta.Index] = builder
			}
			builder.apply(delta)
		}
		return nil
	})
	if err != nil {
		return "", nil, err
	}

	return content.String(), finalizeToolCalls(calls), nil
}

func (a *Agent) toolDefinitions() []llm.Tool {
	if a.toolRegistry == nil {
		return nil
	}
	return a.toolRegistry.Definitions()
}

func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) string {
	name := call.Function.Name
	if logs != nil {
		fmt.Fprintf(logs, "\n[tool] %s %s\n", name, compact(call.Function.Arguments, 400))
	}

	if a.toolRegistry == nil {
		return tools.FormatError(fmt.Errorf("tool registry is disabled"))
	}
	tool, ok := a.toolRegistry.Get(name)
	if !ok {
		return tools.FormatError(fmt.Errorf("unknown tool: %s", name))
	}

	args := strings.TrimSpace(call.Function.Arguments)
	if args == "" {
		args = "{}"
	}

	toolCtx, cancel := context.WithTimeout(ctx, a.toolTimeout)
	defer cancel()
	output, err := tool.Execute(toolCtx, json.RawMessage(args))
	if err != nil {
		if output != "" {
			return output + "\n" + tools.FormatError(err)
		}
		return tools.FormatError(err)
	}
	if logs != nil {
		fmt.Fprintf(logs, "[tool] %s done\n", name)
	}
	return output
}

type toolCallBuilder struct {
	index     int
	id        string
	callType  string
	name      string
	arguments strings.Builder
}

func (b *toolCallBuilder) apply(delta llm.ToolCallDelta) {
	if delta.ID != "" {
		b.id = delta.ID
	}
	if delta.Type != "" {
		b.callType = delta.Type
	}
	if delta.Name != "" {
		b.name = delta.Name
	}
	if delta.Arguments != "" {
		b.arguments.WriteString(delta.Arguments)
	}
}

func finalizeToolCalls(builders map[int]*toolCallBuilder) []llm.ToolCall {
	if len(builders) == 0 {
		return nil
	}
	indexes := make([]int, 0, len(builders))
	for index := range builders {
		indexes = append(indexes, index)
	}
	sort.Ints(indexes)

	calls := make([]llm.ToolCall, 0, len(indexes))
	for _, index := range indexes {
		builder := builders[index]
		id := builder.id
		if id == "" {
			id = fmt.Sprintf("call_%d", index)
		}
		callType := builder.callType
		if callType == "" {
			callType = "function"
		}
		calls = append(calls, llm.ToolCall{
			ID:   id,
			Type: callType,
			Function: llm.FunctionCall{
				Name:      builder.name,
				Arguments: builder.arguments.String(),
			},
		})
	}
	return calls
}

func compact(value string, limit int) string {
	value = strings.TrimSpace(value)
	if len(value) <= limit {
		return value
	}
	return value[:limit] + "...[truncated]"
}