Merge pull request #868 from trheyi/main

Add callback support for assistant chat processing

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</think>\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</tool>\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

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 {
+		if len(args) > 3 {
-			optionsRaw, err := bridge.GoValue(args[2], info.Context())
+			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())
 		}

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)
 }
 

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())