From 979acb27a87dcdb7a28e07efb8164134ac7fd298 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 18 Feb 2025 18:13:42 +0800 Subject: [PATCH] Add callback support for assistant chat processing - Introduce optional callback mechanism for chat messages - Modify Execute and streamChat methods to support callback functions - Add support for process and script method callbacks - Update message handling to trigger callbacks during chat streaming - Enhance flexibility of assistant chat interactions with dynamic callbacks --- neo/assistant/api.go | 42 ++++++++++++++--------- neo/assistant/hooks.go | 75 ++++++++++++++++++++++++++++++++++++++---- neo/assistant/types.go | 2 +- neo/message/message.go | 25 ++++++++++++++ 4 files changed, 122 insertions(+), 22 deletions(-) diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 94175a8c..39d0d106 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -46,17 +46,17 @@ func GetByConnector(connector string, name string) (*Assistant, error) { } // Execute implements the execute functionality -func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error { +func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}, callback ...interface{}) error { contents := chatMessage.NewContents() messages, err := ast.withHistory(ctx, input) if err != nil { return err } - return ast.execute(c, ctx, messages, options, contents) + return ast.execute(c, ctx, messages, options, contents, callback...) } // Execute implements the execute functionality -func (ast *Assistant) execute(c *gin.Context, ctx chatctx.Context, input []chatMessage.Message, userOptions map[string]interface{}, contents *chatMessage.Contents) error { +func (ast *Assistant) execute(c *gin.Context, ctx chatctx.Context, input []chatMessage.Message, userOptions map[string]interface{}, contents *chatMessage.Contents, callback ...interface{}) error { if contents == nil { contents = chatMessage.NewContents() @@ -123,11 +123,11 @@ func (ast *Assistant) execute(c *gin.Context, ctx chatctx.Context, input []chatM // Update assistant id ctx.AssistantID = res.AssistantID - return newAst.handleChatStream(c, ctx, input, options, contents) + return newAst.handleChatStream(c, ctx, input, options, contents, callback...) } // Only proceed with chat stream if no specific next action was handled - return ast.handleChatStream(c, ctx, input, options, contents) + return ast.handleChatStream(c, ctx, input, options, contents, callback...) } // Execute the next action @@ -289,13 +289,13 @@ func (ast *Assistant) Call(c *gin.Context, payload APIPayload) (interface{}, err } // handleChatStream manages the streaming chat interaction with the AI -func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []chatMessage.Message, options map[string]interface{}, contents *chatMessage.Contents) error { +func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []chatMessage.Message, options map[string]interface{}, contents *chatMessage.Contents, callback ...interface{}) error { clientBreak := make(chan bool, 1) done := make(chan bool, 1) // Chat with AI in background go func() { - err := ast.streamChat(c, ctx, messages, options, clientBreak, done, contents) + err := ast.streamChat(c, ctx, messages, options, clientBreak, done, contents, callback...) if err != nil { chatMessage.New().Error(err).Done().Write(c.Writer) } @@ -320,7 +320,14 @@ func (ast *Assistant) streamChat( options map[string]interface{}, clientBreak chan bool, done chan bool, - contents *chatMessage.Contents) error { + contents *chatMessage.Contents, + callback ...interface{}, +) error { + + var cb interface{} + if len(callback) > 0 { + cb = callback[0] + } errorRaw := "" isFirst := true @@ -363,7 +370,7 @@ func (ast *Assistant) streamChat( newMsg := chatMessage.New().Error(value).Done() newMsg.Retry = ctx.Retry newMsg.Silent = ctx.Silent - newMsg.Write(c.Writer) + newMsg.Callback(cb).Write(c.Writer) return 0 // break } @@ -380,8 +387,11 @@ func (ast *Assistant) streamChat( if isThinking && msg.Type != "think" { // add the think close tag end := chatMessage.New().Map(map[string]interface{}{"text": "\n\n", "type": "think", "delta": true}) - end.Write(c.Writer) end.ID = currentMessageID + end.Retry = ctx.Retry + end.Silent = ctx.Silent + + end.Callback(cb).Write(c.Writer) end.AppendTo(contents) contents.UpdateType("think", map[string]interface{}{"text": contents.Text()}, currentMessageID) isThinking = false @@ -405,8 +415,10 @@ func (ast *Assistant) streamChat( if msg.IsDone { end := chatMessage.New().Map(map[string]interface{}{"text": "}\n\n", "type": "tool", "delta": true}) - end.Write(c.Writer) end.ID = currentMessageID + end.Retry = ctx.Retry + end.Silent = ctx.Silent + end.Callback(cb).Write(c.Writer) end.AppendTo(contents) contents.UpdateType("tool", map[string]interface{}{"text": contents.Text()}, currentMessageID) isTool = false @@ -485,12 +497,11 @@ func (ast *Assistant) streamChat( output.Retry = ctx.Retry // Retry mode output.Silent = ctx.Silent // Silent mode - if isFirst { output.Assistant(ast.ID, ast.Name, ast.Avatar) isFirst = false } - output.Write(c.Writer) + output.Callback(cb).Write(c.Writer) } // Complete the stream @@ -510,6 +521,7 @@ func (ast *Assistant) streamChat( "retry": ctx.Retry, "silent": ctx.Silent, }). + Callback(cb). Write(c.Writer) } @@ -544,7 +556,7 @@ func (ast *Assistant) streamChat( output.Retry = ctx.Retry output.Silent = ctx.Silent } - output.Write(c.Writer) + output.Callback(cb).Write(c.Writer) done <- true return 0 // break } @@ -566,7 +578,7 @@ func (ast *Assistant) streamChat( } msg.Retry = ctx.Retry msg.Silent = ctx.Silent - msg.Done().Write(c.Writer) + msg.Done().Callback(cb).Write(c.Writer) } return nil diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go index 754df070..f4ed1d14 100644 --- a/neo/assistant/hooks.go +++ b/neo/assistant/hooks.go @@ -7,9 +7,11 @@ import ( "strings" "time" + "github.com/fatih/color" "github.com/gin-gonic/gin" "github.com/google/uuid" jsoniter "github.com/json-iterator/go" + "github.com/yaoapp/gou/process" "github.com/yaoapp/gou/runtime/v8/bridge" "github.com/yaoapp/kun/log" chatctx "github.com/yaoapp/yao/neo/context" @@ -19,7 +21,7 @@ import ( ) // HookInit initialize the assistant -func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []message.Message, options map[string]interface{}, contents *message.Contents) (*ResHookInit, error) { +func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []chatMessage.Message, options map[string]interface{}, contents *chatMessage.Contents) (*ResHookInit, error) { // Create timeout context ctx := ast.createBackgroundContext() v, err := ast.call(ctx, "Init", c, contents, context, input, options) @@ -310,7 +312,7 @@ func (ast *Assistant) call(ctx context.Context, method string, c *gin.Context, c // Add sendMessage function to the script context scriptCtx.WithFunction("SendMessage", sendMessage(c, contents)) - scriptCtx.WithFunction("Run", run(c, context)) + scriptCtx.WithFunction("Run", ast.run(c, context)) // Check if the method exists if !scriptCtx.Global().Has(method) { @@ -326,7 +328,7 @@ func (ast *Assistant) call(ctx context.Context, method string, c *gin.Context, c } // Execute the assistant -func run(c *gin.Context, context chatctx.Context) func(info *v8go.FunctionCallbackInfo) *v8go.Value { +func (ast *Assistant) run(c *gin.Context, context chatctx.Context) func(info *v8go.FunctionCallbackInfo) *v8go.Value { return func(info *v8go.FunctionCallbackInfo) *v8go.Value { // Get the args @@ -345,11 +347,72 @@ func run(c *gin.Context, context chatctx.Context) func(info *v8go.FunctionCallba } // input []chatMessage.Message + var cb func(msg *chatMessage.Message) input := args[1].String() + if len(args) > 2 { + + goValue, err := bridge.GoValue(args[2], info.Context()) + if err != nil { + return bridge.JsException(info.Context(), err.Error()) + } + + name := "" + userArgs := []interface{}{} + switch v := goValue.(type) { + case string: + name = v + case map[string]interface{}: + if fname, ok := v["name"].(string); ok { + name = fname + } + if args, ok := v["args"].([]interface{}); ok { + userArgs = args + } + } + + if strings.Contains(name, ".") { + cb = func(msg *chatMessage.Message) { + cbArgs := []interface{}{} + cbArgs = append(cbArgs, msg) + cbArgs = append(cbArgs, userArgs...) + p, err := process.Of(name, cbArgs...) + if err != nil { + log.Error("Failed to get the process: %s", err.Error()) + color.Red("Failed to get the process: %s", err.Error()) + return + } + err = p.Execute() + if err != nil { + log.Error("Failed to execute the process: %s", err.Error()) + color.Red("Failed to execute the process: %s", err.Error()) + return + } + defer p.Release() + } + } + + // Call self method + cb = func(msg *chatMessage.Message) { + cbArgs := []interface{}{} + cbArgs = append(cbArgs, msg) + cbArgs = append(cbArgs, userArgs...) + ctx, err := ast.Script.NewContext(context.Sid, nil) + if err != nil { + return + } + defer ctx.Close() + _, err = ctx.CallWith(context, name, cbArgs...) + if err != nil { + log.Error("Failed to call the method: %s", err.Error()) + color.Red("Failed to call the method: %s", err.Error()) + return + } + } + } options := map[string]interface{}{} - if len(args) > 2 { - optionsRaw, err := bridge.GoValue(args[2], info.Context()) + if len(args) > 3 { + optionsRaw, err := bridge.GoValue(args[3], info.Context()) if err != nil { return bridge.JsException(info.Context(), err.Error()) } @@ -374,7 +437,7 @@ func run(c *gin.Context, context chatctx.Context) func(info *v8go.FunctionCallba context.AssistantID = assistantID context.ChatID = fmt.Sprintf("chat_%s", uuid.New().String()) // New chat id context.Silent = true // Silent mode - err = assistant.Execute(c, context, input, options) // Execute the assistant + err = assistant.Execute(c, context, input, options, cb) // Execute the assistant if err != nil { return bridge.JsException(info.Context(), err.Error()) } diff --git a/neo/assistant/types.go b/neo/assistant/types.go index cad8fd16..394b21f2 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -26,7 +26,7 @@ type API interface { ReadBase64(ctx context.Context, fileID string) (string, error) GetPlaceholder() *Placeholder - Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error + Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}, callback ...interface{}) error Call(c *gin.Context, payload APIPayload) (interface{}, error) } diff --git a/neo/message/message.go b/neo/message/message.go index 125d6367..0236b547 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -693,6 +693,26 @@ func (m *Message) Bind(data map[string]interface{}) *Message { return m } +// Callback callback the message +func (m *Message) Callback(fn interface{}) *Message { + if fn != nil { + switch v := fn.(type) { + case func(msg *Message): + v(m) + break + + case func(): + v() + break + + default: + fmt.Println("no match callback") + break + } + } + return m +} + // Write writes the message to response writer func (m *Message) Write(w gin.ResponseWriter) bool { defer func() { @@ -702,6 +722,11 @@ func (m *Message) Write(w gin.ResponseWriter) bool { } }() + // Ignore silent messages + if m.Silent { + return true + } + data, err := jsoniter.Marshal(m) if err != nil { log.Error("%s", err.Error())