Merge pull request #851 from trheyi/main

Refactor streaming chat and message processing

neo/assistant/api.go

@@ -206,11 +206,19 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context, contents *c
 			return fmt.Errorf("with history error: %s", err.Error())
 		}
 
-		fmt.Println("---messages ---")
+		// Create a new Text
-		utils.Dump(messages)
+		// Send loading message and mark as new
-		fmt.Println(`chatID: `, ctx.ChatID)
+		msg := chatMessage.New().Map(map[string]interface{}{
+			"new":   true,
+			"role":  "assistant",
+			"type":  "loading",
+			"props": map[string]interface{}{"placeholder": "Calling " + assistant.Name},
+		})
+		msg.Assistant(assistant.ID, assistant.Name, assistant.Avatar)
+		msg.Write(c.Writer)
+		newContents := chatMessage.NewContents()
 
-		return assistant.execute(c, ctx, messages, options, contents)
+		return assistant.execute(c, ctx, messages, options, newContents)
 
 	case "exit":
 		return nil
@@ -360,33 +368,39 @@ func (ast *Assistant) streamChat(
 					}
 
 					// New message with the tails
-					newMsg, err := chatMessage.NewString(tails, id)
+					if tails != "" {
-					if err != nil {
+						newMsg, err := chatMessage.NewString(tails, id)
-						return
+						if err != nil {
+							return
+						}
+						messages = append(messages, *newMsg)
 					}
-					messages = append(messages, *newMsg)
 				})
 
 				// Handle stream
-				res, err := ast.HookStream(c, ctx, messages, msg, contents)
+				// The stream hook is not used, because there's no need to handle the stream output
-				if err == nil && res != nil {
+				// if some thing need to be handled in future, we can use the stream hook again
+				// ------------------------------------------------------------------------------
+				// res, err := ast.HookStream(c, ctx, messages, msg, contents)
+				// if err == nil && res != nil {
 
-					if res.Next != nil {
+				// 	if res.Next != nil {
-						err = res.Next.Execute(c, ctx, contents)
+				// 		err = res.Next.Execute(c, ctx, contents)
-						if err != nil {
+				// 		if err != nil {
-							chatMessage.New().Error(err.Error()).Done().Write(c.Writer)
+				// 			chatMessage.New().Error(err.Error()).Done().Write(c.Writer)
-						}
+				// 		}
 
-						done <- true
+				// 		done <- true
-						return 0 // break
+				// 		return 0 // break
-					}
+				// 	}
 
-					if res.Silent {
+				// 	if res.Silent {
-						return 1 // continue
+				// 		return 1 // continue
-					}
+				// 	}
-				}
+				// }
+				// ------------------------------------------------------------------------------
 
-				// Write the message to the client
+				// Write the message to the stream
 				output := chatMessage.New().Map(map[string]interface{}{
 					"text":  delta,
 					"type":  msg.Type,
@@ -414,6 +428,10 @@ func (ast *Assistant) streamChat(
 				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)

neo/assistant/hooks.go

@@ -3,6 +3,7 @@ package assistant
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/gin-gonic/gin"
@@ -138,7 +139,40 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
 	// Create timeout context
 	ctx := ast.createBackgroundContext()
 
-	v, err := ast.call(ctx, "Done", c, contents, context, input, contents.Data)
+	// format the output
+	// 1. Remove thinking message
+	// 2. Parse the tool call message content
+	output := []message.Data{}
+	if contents != nil && contents.Data != nil {
+		for _, data := range contents.Data {
+			if data.Type == "think" {
+				continue
+			}
+
+			// parse the tool call message content
+			if data.Type == "tool" && data.Props != nil {
+				props := map[string]interface{}{}
+				if text, ok := data.Props["text"].(string); ok {
+
+					// Remove <tool> and </tool> tags
+					text = strings.ReplaceAll(text, "<tool>", "")
+					text = strings.ReplaceAll(text, "</tool>", "")
+
+					// Parse the text into props
+					err := jsoniter.UnmarshalFromString(text, &props)
+					if err != nil {
+						props["error"] = err.Error()
+					}
+				}
+
+				output = append(output, message.Data{Type: "tool", Props: props})
+				continue
+			}
+			output = append(output, data)
+		}
+	}
+
+	v, err := ast.call(ctx, "Done", c, contents, context, input, output)
 	if err != nil {
 		if err.Error() == HookErrorMethodNotFound {
 			return nil, nil