Merge pull request #852 from trheyi/main

Refactor chat history and message processing

neo/assistant/api.go

@@ -9,8 +9,6 @@ import (
 	"github.com/gin-gonic/gin"
 	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/fs"
-	"github.com/yaoapp/gou/process"
-	"github.com/yaoapp/kun/utils"
 	chatctx "github.com/yaoapp/yao/neo/context"
 	chatMessage "github.com/yaoapp/yao/neo/message"
 )
@@ -115,35 +113,37 @@ func (ast *Assistant) execute(c *gin.Context, ctx chatctx.Context, input []chatM
 func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context, contents *chatMessage.Contents) error {
 	switch next.Action {
 
-	case "process":
+	// It's not used, because the process could be executed in the hook script
-		if next.Payload == nil {
+	// It may remove in the future
-			return fmt.Errorf("payload is required")
+	// case "process":
-		}
+	// 	if next.Payload == nil {
+	// 		return fmt.Errorf("payload is required")
+	// 	}
 
-		name, ok := next.Payload["name"].(string)
+	// 	name, ok := next.Payload["name"].(string)
-		if !ok {
+	// 	if !ok {
-			return fmt.Errorf("process name should be string")
+	// 		return fmt.Errorf("process name should be string")
-		}
+	// 	}
 
-		args := []interface{}{}
+	// 	args := []interface{}{}
-		if v, ok := next.Payload["args"].([]interface{}); ok {
+	// 	if v, ok := next.Payload["args"].([]interface{}); ok {
-			args = v
+	// 		args = v
-		}
+	// 	}
 
-		// Add context and writer to args
+	// 	// Add context and writer to args
-		args = append(args, ctx, c.Writer)
+	// 	args = append(args, ctx, c.Writer)
-		p, err := process.Of(name, args...)
+	// 	p, err := process.Of(name, args...)
-		if err != nil {
+	// 	if err != nil {
-			return fmt.Errorf("get process error: %s", err.Error())
+	// 		return fmt.Errorf("get process error: %s", err.Error())
-		}
+	// 	}
 
-		err = p.Execute()
+	// 	err = p.Execute()
-		if err != nil {
+	// 	if err != nil {
-			return fmt.Errorf("execute process error: %s", err.Error())
+	// 		return fmt.Errorf("execute process error: %s", err.Error())
-		}
+	// 	}
-		defer p.Release()
+	// 	defer p.Release()
 
-		return nil
+	// 	return nil
 
 	case "assistant":
 		if next.Payload == nil {
@@ -201,6 +201,7 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context, contents *c
 			options = v
 		}
 
+		input.Hidden = true // not show in the history
 		messages, err := assistant.withHistory(ctx, input)
 		if err != nil {
 			return fmt.Errorf("with history error: %s", err.Error())
@@ -263,9 +264,6 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
 		if err != nil {
 			chatMessage.New().Error(err).Done().Write(c.Writer)
 		}
-
-		ast.saveChatHistory(ctx, messages, contents)
-		fmt.Printf("saveChatHistory %v\n", ctx.ChatID)
 		done <- true
 	}()
 
@@ -418,30 +416,8 @@ func (ast *Assistant) streamChat(
 			// Complete the stream
 			if msg.IsDone {
 
-				// if value == "" {
+				// Send the last message to the client
-				// 	msg.Write(c.Writer)
+				if delta != "" {
-				// }
-
-				// Remove the last empty data
-				contents.RemoveLastEmpty()
-
-				res, hookErr := ast.HookDone(c, ctx, messages, contents)
-				if hookErr == nil && res != nil {
-					if res.Next != nil {
-
-						fmt.Println("---- Execute Next ---")
-						utils.Dump(res.Next)
-						fmt.Println("---- Execute Next end ---")
-						err := res.Next.Execute(c, ctx, contents)
-						if err != nil {
-							chatMessage.New().Error(err.Error()).Done().Write(c.Writer)
-						}
-
-						done <- true
-						return 0 // break
-					}
-
-				} else if delta != "" {
 					chatMessage.New().
 						Map(map[string]interface{}{
 							"assistant_id":     ast.ID,
@@ -455,22 +431,36 @@ func (ast *Assistant) streamChat(
 						Write(c.Writer)
 				}
 
-				// Hook execute error
+				// Remove the last empty data
+				contents.RemoveLastEmpty()
+				res, hookErr := ast.HookDone(c, ctx, messages, contents)
+
+				// Some error occurred in the hook, return the error
 				if hookErr != nil {
 					chatMessage.New().Error(hookErr.Error()).Done().Write(c.Writer)
 					done <- true
 					return 0 // break
 				}
 
-				msg := chatMessage.New().Done()
+				// Save the chat history
-				if res != nil && res.Output != nil {
+				ast.saveChatHistory(ctx, messages, contents)
-					msg = chatMessage.New().
+
-						Map(map[string]interface{}{
+				// If the hook is successful, execute the next action
-							"text": res.Input,
+				if res != nil && res.Next != nil {
-							"done": true,
+					err := res.Next.Execute(c, ctx, contents)
-						})
+					if err != nil {
+						chatMessage.New().Error(err.Error()).Done().Write(c.Writer)
 					}
-				msg.Write(c.Writer)
+					done <- true
+					return 0 // break
+				}
+
+				// The default output
+				output := chatMessage.New().Done()
+				if res != nil && res.Output != nil {
+					output = chatMessage.New().Map(map[string]interface{}{"text": res.Output, "done": true})
+				}
+				output.Write(c.Writer)
 				done <- true
 				return 0 // break
 			}
@@ -516,20 +506,9 @@ func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessag
 			},
 		}
 
-		// contents
+		// if the user message is hidden, just save the assistant message
-		fmt.Println("---contents ---")
+		if userMessage.Hidden {
-		if contents.Data != nil {
+			data = []map[string]interface{}{data[1]}
-			fmt.Println("---contents.Data ---")
-			for _, content := range contents.Data {
-				fmt.Println(content.Map())
-			}
-			fmt.Println("---contents.Data end ---")
-		}
-		fmt.Println("---contents end ---")
-
-		// Add mentions
-		if userMessage.Mentions != nil {
-			data[0]["mentions"] = userMessage.Mentions
 		}
 
 		storage.SaveHistory(ctx.Sid, data, ctx.ChatID, ctx.Map())

neo/message/message.go

@@ -33,6 +33,7 @@ type Message struct {
 	Mentions        []Mention              `json:"menions,omitempty"`          // Mentions for the message ( for user  role = user )
 	Data            map[string]interface{} `json:"-"`                          // data for the message
 	Pending         bool                   `json:"-"`                          // pending for the message
+	Hidden          bool                   `json:"hidden,omitempty"`           // hidden for the message (not show in the UI and history)
 }
 
 // Mention represents a mention