233 lines
8.5 KiB
Go
233 lines
8.5 KiB
Go
package flow
|
||
|
||
import (
|
||
"ai-agent/gateway"
|
||
sessionDao "ai-agent/workflow/dao/session"
|
||
sessionDto "ai-agent/workflow/model/dto/session"
|
||
"ai-agent/workflow/model/entity"
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"time"
|
||
|
||
"ai-agent/tools/runner"
|
||
|
||
"gitea.redpowerfuture.com/red-future/common/tools"
|
||
wsCommon "gitea.redpowerfuture.com/red-future/common/websocket"
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
|
||
"github.com/gogf/gf/v2/os/glog"
|
||
"github.com/gogf/gf/v2/util/gconv"
|
||
)
|
||
|
||
func init() {
|
||
// 普通对话消息处理器（会话服务器 SessionWsService 见 ws_server.go）
|
||
SessionWsService.OnMessage("agent", handleToolAgent)
|
||
SessionWsService.OnMessage("agent_cancel", handleToolAgentCancel)
|
||
}
|
||
|
||
// 工具对话默认系统提示词
|
||
const defaultAgentSystemPrompt = "你是一个智能助手，可以调用工具完成任务。请根据任务需要选择合适的工具，参考工具返回结果，最终给出完整回答。"
|
||
|
||
// 工具对话 ReAct 最大循环步数
|
||
const defaultAgentMaxStep = 15
|
||
|
||
// handleToolAgent 处理工具对话消息：解析 payload 后异步运行 ReAct 循环，逐步推送过程事件。
|
||
// 首条消息惰性建会话（复用已存在会话），跑完后把问答/token 写入 exec_chat 落库。
|
||
func handleToolAgent(ctx context.Context, conn *wsCommon.WsConnection, payload interface{}) {
|
||
var p sessionDto.WebSocketExecChatReq
|
||
if err := gconv.Struct(payload, &p); err != nil {
|
||
pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "参数解析失败", Error: err.Error()})
|
||
return
|
||
}
|
||
if p.Question == "" {
|
||
pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "内容不能为空", Error: "提问内容不能为空"})
|
||
return
|
||
}
|
||
|
||
saveCtx := context.WithoutCancel(ctx)
|
||
|
||
id := p.Id
|
||
if g.IsEmpty(id) {
|
||
// 会话落库：前端临时 sessionId 对应已存在会话则复用，否则新建
|
||
err := ensureSession(saveCtx, conn.SessionId, p.Question)
|
||
if err != nil {
|
||
glog.Errorf(ctx, "会话创建失败: %v", err)
|
||
pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "会话创建失败", Error: err.Error()})
|
||
return
|
||
}
|
||
|
||
// 问答落库：
|
||
chatId, err := sessionDao.ExecChatDao.Insert(ctx, &sessionDto.CreateExecChatReq{
|
||
SessionId: conn.SessionId,
|
||
RequestParams: entity.ExecChatRequestParams{Question: p.Question},
|
||
})
|
||
if err != nil {
|
||
glog.Errorf(ctx, "问答创建失败: %v", err)
|
||
pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "问答创建失败", Error: err.Error()})
|
||
return
|
||
}
|
||
id = chatId
|
||
pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventRoundStart, Id: chatId})
|
||
}
|
||
|
||
modelTools, err := tools.Default.List(ctx)
|
||
if err != nil {
|
||
pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "工具列表获取失败", Error: err.Error()})
|
||
|
||
errChat := recordChat(saveCtx, id, "", "工具列表获取失败", err, 0, 0, 0)
|
||
if errChat != nil {
|
||
glog.Errorf(ctx, "普通对话落库失败: %v", errChat)
|
||
pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "对话落库失败", Error: errChat.Error()})
|
||
}
|
||
return
|
||
}
|
||
systemPrompt := p.SystemPrompt
|
||
if systemPrompt == "" {
|
||
systemPrompt = defaultAgentSystemPrompt
|
||
}
|
||
|
||
// 支持前端终止：agent 上下文可取消；落库用不带取消的 ctx（保留 request 值），保证终止后 token 仍能记录
|
||
agentCtx, agentCancel := context.WithCancel(ctx)
|
||
if oldCancel := getToolCancel(conn); oldCancel != nil {
|
||
oldCancel()
|
||
}
|
||
conn.SetMeta("toolCancel", agentCancel)
|
||
defer conn.SetMeta("toolCancel", nil)
|
||
defer agentCancel()
|
||
|
||
agent := runner.NewReActAgent(p.ModelId, conn.SessionId, modelTools, systemPrompt, defaultAgentMaxStep)
|
||
agent.OnEvent = func(ev runner.ReActEvent) {
|
||
pushAgentEvent(conn, ev)
|
||
}
|
||
|
||
start := time.Now()
|
||
answer, runErr := agent.Run(agentCtx, p.Question)
|
||
duration := int64(time.Since(start).Seconds())
|
||
|
||
// 前端终止：结果/错误不推前端，仅把已产生的 token 正常落库（友好提示记「用户已终止对话」，不记原始错误）
|
||
var errMsg string
|
||
terminated := runErr != nil && errors.Is(runErr, context.Canceled)
|
||
if terminated {
|
||
errMsg = errChatTerminated.Error()
|
||
runErr = nil
|
||
} else if runErr != nil {
|
||
errMsg = "对话运行失败"
|
||
pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "对话运行失败", Error: runErr.Error()})
|
||
}
|
||
err = recordChat(saveCtx, id, answer, errMsg, runErr, agent.TotalTokens, agent.TotalCost, duration)
|
||
if err != nil {
|
||
glog.Errorf(ctx, "普通对话落库失败: %v", err)
|
||
pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventError, Message: "对话落库失败", Error: err.Error()})
|
||
}
|
||
pushAgentEvent(conn, runner.ReActEvent{Type: runner.ReActEventAnswer, Answer: answer})
|
||
}
|
||
|
||
// handleToolAgentCancel 终止正在运行的对话（前端停止按钮发送 agent_cancel）
|
||
func handleToolAgentCancel(ctx context.Context, conn *wsCommon.WsConnection, _ interface{}) {
|
||
if cancel := getToolCancel(conn); cancel != nil {
|
||
cancel()
|
||
}
|
||
_ = conn.WriteJSON(&wsCommon.WsPushMsg{Type: "ack", Message: "已终止对话"})
|
||
}
|
||
|
||
// getToolCancel 获取当前 agent 运行的取消函数
|
||
func getToolCancel(conn *wsCommon.WsConnection) context.CancelFunc {
|
||
cancel, _ := wsCommon.GetMetaT[context.CancelFunc](conn, "toolCancel")
|
||
return cancel
|
||
}
|
||
|
||
// errChatTerminated 前端终止对话的错误标记（写入 exec_chat.error_message）
|
||
var errChatTerminated = errors.New("用户已终止对话")
|
||
|
||
// recordChat 把一次普通对话写入 exec_chat：答案传 OSS 存 result_file_url，
|
||
// 友好提示写 error_message，原始错误写 error，token 与费用（模型网关返回的累计 cost）落库
|
||
func recordChat(ctx context.Context, id int64, answer string, msg string, runErr error, totalTokens int64, totalCost float64, duration int64) error {
|
||
var resultFileUrl string
|
||
if runErr == nil && answer != "" {
|
||
url, uploadErr := gateway.Upload(ctx, fmt.Sprintf("chat_%v_%d.txt", id, time.Now().UnixMilli()), []byte(answer))
|
||
if uploadErr != nil {
|
||
glog.Errorf(ctx, "普通对话答案上传OSS失败: %v", uploadErr)
|
||
} else {
|
||
resultFileUrl = url
|
||
}
|
||
}
|
||
var errorDetail string
|
||
if runErr != nil {
|
||
errorDetail = runErr.Error()
|
||
}
|
||
_, err := sessionDao.ExecChatDao.Update(ctx, &sessionDto.UpdateExecChatReq{
|
||
Id: id,
|
||
Duration: duration,
|
||
ResultFileUrl: resultFileUrl,
|
||
TotalTokens: int(totalTokens),
|
||
TotalFee: totalCost,
|
||
ErrorMessage: msg,
|
||
Error: errorDetail,
|
||
})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
// 执行成功：重新执行复用了同一条记录，OmitEmpty 的 Update 会跳过空 error_message/error，
|
||
// 需显式清空，避免上一次失败的报错残留
|
||
if runErr == nil {
|
||
if _, err := sessionDao.ExecChatDao.ClearError(ctx, id); err != nil {
|
||
glog.Errorf(ctx, "exec_chat 报错信息清空失败: %v", err)
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// pushAgentEvent 把 ReAct 过程事件转为 WS 推送消息
|
||
func pushAgentEvent(conn *wsCommon.WsConnection, ev runner.ReActEvent) {
|
||
if conn.IsClosed() {
|
||
return
|
||
}
|
||
switch ev.Type {
|
||
case runner.ReActEventRoundStart:
|
||
_ = conn.WriteJSON(&wsCommon.WsPushMsg{
|
||
Type: string(ev.Type),
|
||
Message: "运行开始",
|
||
Data: map[string]interface{}{"recordId": ev.Id},
|
||
})
|
||
case runner.ReActEventModelCall:
|
||
_ = conn.WriteJSON(&wsCommon.WsPushMsg{
|
||
Type: string(ev.Type),
|
||
Message: "模型思考中",
|
||
Data: map[string]interface{}{"step": ev.Step, "maxStep": ev.MaxStep},
|
||
})
|
||
case runner.ReActEventToolCall:
|
||
_ = conn.WriteJSON(&wsCommon.WsPushMsg{
|
||
Type: string(ev.Type),
|
||
Message: "调用工具",
|
||
Data: map[string]interface{}{"description": ev.Description},
|
||
})
|
||
case runner.ReActEventToolResult:
|
||
_ = conn.WriteJSON(&wsCommon.WsPushMsg{
|
||
Type: string(ev.Type),
|
||
Message: "工具返回",
|
||
Data: map[string]interface{}{"description": ev.Description},
|
||
})
|
||
case runner.ReActEventAnswerChunk:
|
||
_ = conn.WriteJSON(&wsCommon.WsPushMsg{
|
||
Type: string(ev.Type),
|
||
Message: "思考中",
|
||
Data: map[string]interface{}{"delta": ev.Delta},
|
||
})
|
||
case runner.ReActEventReasoningChunk:
|
||
_ = conn.WriteJSON(&wsCommon.WsPushMsg{
|
||
Type: string(ev.Type),
|
||
Message: "思考中",
|
||
Data: map[string]interface{}{"delta": ev.Delta},
|
||
})
|
||
case runner.ReActEventAnswer:
|
||
_ = conn.WriteJSON(&wsCommon.WsPushMsg{
|
||
Type: string(ev.Type),
|
||
Message: "作答完成",
|
||
Data: map[string]interface{}{"answer": ev.Answer},
|
||
})
|
||
case runner.ReActEventError:
|
||
_ = conn.WriteJSON(&wsCommon.WsPushMsg{Type: string(ev.Type), Message: ev.Message, Error: ev.Error})
|
||
}
|
||
}
|