Enhance Neo API assistant interaction with improved message handling and streaming support

- Added message handling in the Execute method to support input messages from the response.
- Updated the handleChatStream method to pass the context to the streamChat method, improving context management during chat streaming.
- Introduced the HookStream method to handle streaming responses, allowing for custom output and next action handling based on the assistant's response.
- Enhanced the ResHookStream struct to include silent output control and next action management, providing better flexibility in assistant interactions.

These changes improve the robustness and maintainability of the Neo API, paving the way for enhanced assistant functionalities and message management.

neo/assistant/api.go

@@ -81,6 +81,11 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
 		options = res.Options
 	}
 
+	// messages
+	if res.Input != nil {
+		messages = res.Input
+	}
+
 	// Only proceed with chat stream if no specific next action was handled
 	return ast.handleChatStream(c, ctx, messages, options)
 }
@@ -93,7 +98,7 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
 
 	// Chat with AI in background
 	go func() {
-		err := ast.streamChat(c, messages, options, clientBreak, done, &content)
+		err := ast.streamChat(c, ctx, messages, options, clientBreak, done, &content)
 		if err != nil {
 			chatMessage.New().Error(err).Done().Write(c.Writer)
 		}
@@ -113,7 +118,7 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
 }
 
 // streamChat handles the streaming chat interaction
-func (ast *Assistant) streamChat(c *gin.Context, messages []message.Message, options map[string]interface{},
+func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages []message.Message, options map[string]interface{},
 	clientBreak chan bool, done chan bool, content *[]byte) error {
 
 	return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
@@ -138,12 +143,35 @@ func (ast *Assistant) streamChat(c *gin.Context, messages []message.Message, opt
 			*content = msg.Append(*content)
 			value := msg.String()
 			if value != "" {
-				chatMessage.New().
+
-					Map(map[string]interface{}{
+				// Handle stream
-						"text": value,
+				res, err := ast.HookStream(c, ctx, messages, value)
-						"done": msg.IsDone,
+				if err != nil {
-					}).
+					return 0 // break
-					Write(c.Writer)
+				}
+
+				// Custom output from hook
+				if res.Output != "" {
+					value = res.Output
+				}
+
+				// Custom next action from hook
+				if res.Next != nil {
+					switch res.Next.Action {
+					case "exit":
+						done <- true
+						return 0 // break
+					}
+				}
+
+				if !res.Silent {
+					chatMessage.New().
+						Map(map[string]interface{}{
+							"text": value,
+							"done": msg.IsDone,
+						}).
+						Write(c.Writer)
+				}
 			}
 
 			// Complete the stream

neo/assistant/hooks.go

@@ -10,26 +10,6 @@ import (
 	"github.com/yaoapp/yao/neo/message"
 )
 
-const (
-	// HookErrorMethodNotFound is the error message for method not found
-	HookErrorMethodNotFound = "method not found"
-)
-
-// ResHookInit the response of the init hook
-type ResHookInit struct {
-	AssistantID string                 `json:"assistant_id,omitempty"`
-	ChatID      string                 `json:"chat_id,omitempty"`
-	Next        *NextAction            `json:"next,omitempty"`
-	Input       []message.Message      `json:"input,omitempty"`
-	Options     map[string]interface{} `json:"options,omitempty"`
-}
-
-// NextAction the next action
-type NextAction struct {
-	Action  string                 `json:"action"`
-	Payload map[string]interface{} `json:"payload,omitempty"`
-}
-
 // HookInit initialize the assistant
 func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error) {
 	// Create timeout context
@@ -54,6 +34,16 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []
 			response.ChatID = res
 		}
 
+		if res, ok := v["next"].(map[string]interface{}); ok {
+			response.Next = &NextAction{}
+			if name, ok := res["action"].(string); ok {
+				response.Next.Action = name
+			}
+			if payload, ok := res["payload"].(map[string]interface{}); ok {
+				response.Next.Payload = payload
+			}
+		}
+
 	case string:
 		response.AssistantID = v
 		response.ChatID = context.ChatID
@@ -66,6 +56,49 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []
 	return response, nil
 }
 
+// HookStream Handle streaming response from LLM
+func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookStream, error) {
+
+	// Create timeout context
+	ctx, cancel := ast.createTimeoutContext(c)
+	defer cancel()
+
+	v, err := ast.call(ctx, "Stream", context, input, output, c.Writer)
+	if err != nil {
+		if err.Error() == HookErrorMethodNotFound {
+			return nil, nil
+		}
+		return nil, err
+	}
+
+	response := &ResHookStream{}
+	switch v := v.(type) {
+	case map[string]interface{}:
+		if res, ok := v["output"].(string); ok {
+			response.Output = res
+		}
+		if res, ok := v["next"].(map[string]interface{}); ok {
+			response.Next = &NextAction{}
+			if name, ok := res["action"].(string); ok {
+				response.Next.Action = name
+			}
+			if payload, ok := res["payload"].(map[string]interface{}); ok {
+				response.Next.Payload = payload
+			}
+		}
+
+		// Custom silent from hook
+		if res, ok := v["silent"].(bool); ok {
+			response.Silent = res
+		}
+
+	case string:
+		response.Output = v
+	}
+
+	return response, nil
+}
+
 // createTimeoutContext creates a timeout context with 5 seconds timeout
 func (ast *Assistant) createTimeoutContext(c *gin.Context) (context.Context, context.CancelFunc) {
 	ctx, cancel := context.WithTimeout(c.Request.Context(), 5*time.Second)

neo/assistant/types.go

@@ -13,6 +13,11 @@ import (
 	api "github.com/yaoapp/yao/openai"
 )
 
+const (
+	// HookErrorMethodNotFound is the error message for method not found
+	HookErrorMethodNotFound = "method not found"
+)
+
 // API the assistant API interface
 type API interface {
 	Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error
@@ -23,6 +28,28 @@ type API interface {
 	HookInit(c *gin.Context, ctx chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error)
 }
 
+// ResHookInit the response of the init hook
+type ResHookInit struct {
+	AssistantID string                 `json:"assistant_id,omitempty"`
+	ChatID      string                 `json:"chat_id,omitempty"`
+	Next        *NextAction            `json:"next,omitempty"`
+	Input       []message.Message      `json:"input,omitempty"`
+	Options     map[string]interface{} `json:"options,omitempty"`
+}
+
+// ResHookStream the response of the stream hook
+type ResHookStream struct {
+	Silent bool        `json:"silent,omitempty"` // Whether to suppress the output
+	Next   *NextAction `json:"next,omitempty"`   // The next action
+	Output string      `json:"output,omitempty"` // The output
+}
+
+// NextAction the next action
+type NextAction struct {
+	Action  string                 `json:"action"`
+	Payload map[string]interface{} `json:"payload,omitempty"`
+}
+
 // RAG the RAG interface
 type RAG struct {
 	Engine     driver.Engine