423 lines
12 KiB
Go
423 lines
12 KiB
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
|
||
}
|