GoLoom/internal/service/agent.go

package service

import (
	"context"
	"fmt"
	"strings"
	"sync/atomic"

	"ai-agent-scaffold-go/internal/model"
)

const maxToolCallIterations = 4

// ChatModelWithTools 扩展接口，同时具备 ChatModel 和工具调用能力
type ChatModelWithTools interface {
	model.ChatModel
	CallTool(ctx context.Context, name, arguments string) (string, error)
}

// ============================================================
// LLMAgent — 基础 LLM Agent
// ============================================================

// LLMAgent 基于 LLM 的智能体，支持多轮工具调用
type LLMAgent struct {
	name        string
	description string
	instruction string
	outputKey   string
	chatModel   ChatModelWithTools
}

// NewLLMAgent 创建 LLM Agent
func NewLLMAgent(name, instruction, description, outputKey string, chatModel ChatModelWithTools) *LLMAgent {
	return &LLMAgent{
		name:        name,
		instruction: instruction,
		description: description,
		outputKey:   outputKey,
		chatModel:   chatModel,
	}
}

func (a *LLMAgent) Name() string      { return a.name }
func (a *LLMAgent) OutputKey() string { return a.outputKey }

// Run 同步执行
func (a *LLMAgent) Run(ctx context.Context, content model.ChatContent) (string, error) {
	return a.runWithVars(ctx, content, map[string]string{})
}

// Stream 流式执行
func (a *LLMAgent) Stream(ctx context.Context, content model.ChatContent, out chan<- string) error {
	return a.streamWithVars(ctx, content, out, map[string]string{})
}

// runWithVars 同步执行，支持变量替换
func (a *LLMAgent) runWithVars(ctx context.Context, content model.ChatContent, vars map[string]string) (string, error) {
	messages := initialMessages(applyVars(a.instruction, vars), firstText(content))

	for iter := 0; iter < maxToolCallIterations; iter++ {
		reply, err := a.chatModel.Generate(ctx, messages)
		if err != nil {
			return "", err
		}
		if len(reply.ToolCalls) == 0 {
			return reply.Content, nil
		}
		// 将 assistant 回复（含工具调用）加入消息历史
		messages = append(messages, model.ChatMessage{
			Role:      model.ChatRoleAssistant,
			Content:   reply.Content,
			ToolCalls: reply.ToolCalls,
		})
		// 执行工具调用，将结果加入消息历史
		toolMessages, err := a.executeToolCalls(ctx, reply.ToolCalls)
		if err != nil {
			return "", err
		}
		messages = append(messages, toolMessages...)
	}
	return "", fmt.Errorf("agent %q exceeded tool-call iteration limit %d", a.name, maxToolCallIterations)
}

// streamWithVars 流式执行，支持变量替换
func (a *LLMAgent) streamWithVars(ctx context.Context, content model.ChatContent, out chan<- string, vars map[string]string) error {
	messages := initialMessages(applyVars(a.instruction, vars), firstText(content))

	for iter := 0; iter < maxToolCallIterations; iter++ {
		events, errs := a.chatModel.Stream(ctx, messages)

		var (
			finalText strings.Builder
			toolCalls []model.ChatToolCall
			done      bool
		)
		streamErr := error(nil)

	streamLoop:
		for {
			select {
			case <-ctx.Done():
				streamErr = ctx.Err()
				break streamLoop
			case ev, ok := <-events:
				if !ok {
					break streamLoop
				}
				if ev.Done {
					done = true
				}
				if ev.Delta != "" {
					finalText.WriteString(ev.Delta)
					select {
					case out <- ev.Delta:
					case <-ctx.Done():
						streamErr = ctx.Err()
						break streamLoop
					}
				}
				if len(ev.ToolCalls) > 0 {
					toolCalls = append(toolCalls, ev.ToolCalls...)
				}
			case err, ok := <-errs:
				if ok && err != nil {
					streamErr = err
				}
				break streamLoop
			}
		}

		if streamErr != nil {
			return streamErr
		}

		// 没有工具调用，说明模型已完成回复
		if len(toolCalls) == 0 {
			if !done {
				return fmt.Errorf("agent %q stream closed without completion", a.name)
			}
			return nil
		}

		// 处理工具调用
		messages = append(messages, model.ChatMessage{
			Role:      model.ChatRoleAssistant,
			Content:   finalText.String(),
			ToolCalls: toolCalls,
		})
		toolMessages, err := a.executeToolCalls(ctx, toolCalls)
		if err != nil {
			return err
		}
		messages = append(messages, toolMessages...)
	}

	return fmt.Errorf("agent %q exceeded tool-call iteration limit %d", a.name, maxToolCallIterations)
}

// executeToolCalls 执行一组工具调用
func (a *LLMAgent) executeToolCalls(ctx context.Context, calls []model.ChatToolCall) ([]model.ChatMessage, error) {
	out := make([]model.ChatMessage, 0, len(calls))
	for _, call := range calls {
		result, err := a.chatModel.CallTool(ctx, call.Name, call.Arguments)
		if err != nil {
			return nil, fmt.Errorf("tool %q: %w", call.Name, err)
		}
		out = append(out, model.ChatMessage{
			Role:       model.ChatRoleTool,
			Content:    result,
			ToolCallID: call.ID,
			Name:       call.Name,
		})
	}
	return out, nil
}

// ============================================================
// 工作流子 Agent 接口
// ============================================================

// workflowSubAgent 工作流内部使用的 Agent 扩展接口
type workflowSubAgent interface {
	model.Agent
	OutputKey() string
	runWithVars(ctx context.Context, content model.ChatContent, vars map[string]string) (string, error)
	streamWithVars(ctx context.Context, content model.ChatContent, out chan<- string, vars map[string]string) error
}

// ============================================================
// SequentialAgent — 顺序工作流
// ============================================================

// SequentialAgent 顺序执行子 Agent，前一个的输出可通过 OutputKey 传递给后一个
type SequentialAgent struct {
	name        string
	description string
	subAgents   []workflowSubAgent
}

// NewSequentialAgent 创建顺序工作流 Agent
func NewSequentialAgent(name, description string, subs []model.Agent) *SequentialAgent {
	wrapped := make([]workflowSubAgent, 0, len(subs))
	for _, s := range subs {
		wrapped = append(wrapped, s.(workflowSubAgent))
	}
	return &SequentialAgent{name: name, description: description, subAgents: wrapped}
}

func (a *SequentialAgent) Name() string      { return a.name }
func (a *SequentialAgent) OutputKey() string { return "" }

func (a *SequentialAgent) Run(ctx context.Context, content model.ChatContent) (string, error) {
	return a.runWithVars(ctx, content, map[string]string{})
}

func (a *SequentialAgent) Stream(ctx context.Context, content model.ChatContent, out chan<- string) error {
	return a.streamWithVars(ctx, content, out, map[string]string{})
}

func (a *SequentialAgent) runWithVars(ctx context.Context, content model.ChatContent, vars map[string]string) (string, error) {
	scope := cloneVars(vars)
	var last string
	for _, sub := range a.subAgents {
		text, err := sub.runWithVars(ctx, content, scope)
		if err != nil {
			return "", err
		}
		last = text
		if key := sub.OutputKey(); key != "" {
			scope[key] = text
		}
	}
	return last, nil
}

func (a *SequentialAgent) streamWithVars(ctx context.Context, content model.ChatContent, out chan<- string, vars map[string]string) error {
	scope := cloneVars(vars)
	for i, sub := range a.subAgents {
		if i == len(a.subAgents)-1 {
			return sub.streamWithVars(ctx, content, out, scope)
		}
		text, err := sub.runWithVars(ctx, content, scope)
		if err != nil {
			return err
		}
		if key := sub.OutputKey(); key != "" {
			scope[key] = text
		}
	}
	return nil
}

// ============================================================
// ParallelAgent — 并行工作流
// ============================================================

// ParallelAgent 并行执行所有子 Agent 并汇总结果
type ParallelAgent struct {
	name        string
	description string
	subAgents   []workflowSubAgent
}

// NewParallelAgent 创建并行工作流 Agent
func NewParallelAgent(name, description string, subs []model.Agent) *ParallelAgent {
	wrapped := make([]workflowSubAgent, 0, len(subs))
	for _, s := range subs {
		wrapped = append(wrapped, s.(workflowSubAgent))
	}
	return &ParallelAgent{name: name, description: description, subAgents: wrapped}
}

func (a *ParallelAgent) Name() string      { return a.name }
func (a *ParallelAgent) OutputKey() string { return "" }

func (a *ParallelAgent) Run(ctx context.Context, content model.ChatContent) (string, error) {
	return a.runWithVars(ctx, content, map[string]string{})
}

func (a *ParallelAgent) Stream(ctx context.Context, content model.ChatContent, out chan<- string) error {
	return a.streamWithVars(ctx, content, out, map[string]string{})
}

func (a *ParallelAgent) runWithVars(ctx context.Context, content model.ChatContent, vars map[string]string) (string, error) {
	parts := make([]string, 0, len(a.subAgents))
	for _, sub := range a.subAgents {
		text, err := sub.runWithVars(ctx, content, vars)
		if err != nil {
			return "", err
		}
		parts = append(parts, fmt.Sprintf("[%s] %s", sub.Name(), text))
	}
	return strings.Join(parts, "\n"), nil
}

func (a *ParallelAgent) streamWithVars(ctx context.Context, content model.ChatContent, out chan<- string, vars map[string]string) error {
	text, err := a.runWithVars(ctx, content, vars)
	if err != nil {
		return err
	}
	select {
	case out <- text:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// ============================================================
// LoopAgent — 循环工作流
// ============================================================

// LoopAgent 循环执行子 Agent，最多执行 maxIterations 次
type LoopAgent struct {
	name          string
	description   string
	subAgents     []workflowSubAgent
	maxIterations int
}

// NewLoopAgent 创建循环工作流 Agent
func NewLoopAgent(name, description string, subs []model.Agent, maxIterations int) *LoopAgent {
	if maxIterations <= 0 {
		maxIterations = 3
	}
	wrapped := make([]workflowSubAgent, 0, len(subs))
	for _, s := range subs {
		wrapped = append(wrapped, s.(workflowSubAgent))
	}
	return &LoopAgent{
		name:          name,
		description:   description,
		subAgents:     wrapped,
		maxIterations: maxIterations,
	}
}

func (a *LoopAgent) Name() string      { return a.name }
func (a *LoopAgent) OutputKey() string { return "" }

func (a *LoopAgent) Run(ctx context.Context, content model.ChatContent) (string, error) {
	return a.runWithVars(ctx, content, map[string]string{})
}

func (a *LoopAgent) Stream(ctx context.Context, content model.ChatContent, out chan<- string) error {
	return a.streamWithVars(ctx, content, out, map[string]string{})
}

func (a *LoopAgent) runWithVars(ctx context.Context, content model.ChatContent, vars map[string]string) (string, error) {
	var last string
	for i := 0; i < a.maxIterations; i++ {
		parts := make([]string, 0, len(a.subAgents))
		for _, sub := range a.subAgents {
			text, err := sub.runWithVars(ctx, content, vars)
			if err != nil {
				return "", err
			}
			parts = append(parts, fmt.Sprintf("[%s] %s", sub.Name(), text))
		}
		last = strings.Join(parts, "\n")
	}
	return last, nil
}

func (a *LoopAgent) streamWithVars(ctx context.Context, content model.ChatContent, out chan<- string, vars map[string]string) error {
	text, err := a.runWithVars(ctx, content, vars)
	if err != nil {
		return err
	}
	select {
	case out <- text:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// ============================================================
// 工具函数
// ============================================================

var sessionCounter atomic.Uint64

// cloneVars 克隆变量映射
func cloneVars(vars map[string]string) map[string]string {
	out := make(map[string]string, len(vars)+4)
	for k, v := range vars {
		out[k] = v
	}
	return out
}

// applyVars 替换模板中的 {key} 占位符
func applyVars(template string, vars map[string]string) string {
	if template == "" || len(vars) == 0 {
		return template
	}
	out := template
	for k, v := range vars {
		out = strings.ReplaceAll(out, "{"+k+"}", v)
	}
	return out
}

// initialMessages 构建初始消息列表（system + user）
func initialMessages(instruction, userText string) []model.ChatMessage {
	messages := make([]model.ChatMessage, 0, 2)
	if strings.TrimSpace(instruction) != "" {
		messages = append(messages, model.ChatMessage{Role: model.ChatRoleSystem, Content: instruction})
	}
	messages = append(messages, model.ChatMessage{Role: model.ChatRoleUser, Content: userText})
	return messages
}

// firstText 从 ChatContent 中提取第一段文本
func firstText(content model.ChatContent) string {
	if len(content.Texts) == 0 {
		return ""
	}
	return content.Texts[0].Message
}