Enhance Neo API assistant response handling with new hook methods

- Introduced HookDone and HookFail methods to manage completion and failure scenarios in assistant responses, improving error handling and output customization.
- Updated streamChat method to utilize these hooks, allowing for more flexible response management based on the assistant's output and error states.
- Enhanced ResHookDone and ResHookFail structs to include next action handling, input messages, and error information, providing better control over 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

@@ -135,6 +135,13 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [
 			// Handle error
 			if msg.Type == "error" {
 				value := msg.String()
+				res, hookErr := ast.HookFail(c, ctx, messages, string(*content), fmt.Errorf("%s", value))
+				if hookErr == nil && res != nil && (res.Output != "" || res.Error != "") {
+					value = res.Output
+					if res.Error != "" {
+						value = res.Error
+					}
+				}
 				chatMessage.New().Error(value).Done().Write(c.Writer)
 				return 0 // break
 			}
@@ -143,28 +150,21 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [
 			*content = msg.Append(*content)
 			value := msg.String()
 			if value != "" {
-
 				// Handle stream
-				res, err := ast.HookStream(c, ctx, messages, value)
+				res, err := ast.HookStream(c, ctx, messages, string(*content))
-				if err != nil {
+				if err == nil && res != nil {
-					return 0 // break
-				}
-
-				// Custom output from hook
 					if res.Output != "" {
 						value = res.Output
 					}
-
+					if res.Next != nil && res.Next.Action == "exit" {
-				// Custom next action from hook
-				if res.Next != nil {
-					switch res.Next.Action {
-					case "exit":
 						done <- true
 						return 0 // break
 					}
+					if res.Silent {
+						return 1 // continue
+					}
 				}
 
-				if !res.Silent {
 				chatMessage.New().
 					Map(map[string]interface{}{
 						"text": value,
@@ -172,13 +172,37 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [
 					}).
 					Write(c.Writer)
 			}
-			}
 
 			// Complete the stream
 			if msg.IsDone {
 				if value == "" {
 					msg.Write(c.Writer)
 				}
+
+				// Call HookDone
+				res, hookErr := ast.HookDone(c, ctx, messages, string(*content))
+				if hookErr == nil && res != nil {
+					if res.Output != "" {
+						chatMessage.New().
+							Map(map[string]interface{}{
+								"text": res.Output,
+								"done": true,
+							}).
+							Write(c.Writer)
+					}
+					if res.Next != nil && res.Next.Action == "exit" {
+						done <- true
+						return 0 // break
+					}
+				} else if value != "" {
+					chatMessage.New().
+						Map(map[string]interface{}{
+							"text": value,
+							"done": true,
+						}).
+						Write(c.Writer)
+				}
+
 				done <- true
 				return 0 // break
 			}

neo/assistant/hooks.go

@@ -99,6 +99,90 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
 	return response, nil
 }
 
+// HookDone Handle completion of assistant response
+func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookDone, error) {
+	// Create timeout context
+	ctx, cancel := ast.createTimeoutContext(c)
+	defer cancel()
+
+	v, err := ast.call(ctx, "Done", context, input, output, c.Writer)
+	if err != nil {
+		if err.Error() == HookErrorMethodNotFound {
+			return nil, nil
+		}
+		return nil, err
+	}
+
+	response := &ResHookDone{
+		Input:  input,
+		Output: output,
+	}
+
+	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
+			}
+		}
+	case string:
+		response.Output = v
+	}
+
+	return response, nil
+}
+
+// HookFail Handle failure of assistant response
+func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input []message.Message, output string, err error) (*ResHookFail, error) {
+	// Create timeout context
+	ctx, cancel := ast.createTimeoutContext(c)
+	defer cancel()
+
+	v, callErr := ast.call(ctx, "Fail", context, input, output, err.Error(), c.Writer)
+	if callErr != nil {
+		if callErr.Error() == HookErrorMethodNotFound {
+			return nil, nil
+		}
+		return nil, callErr
+	}
+
+	response := &ResHookFail{
+		Input:  input,
+		Output: output,
+		Error:  err.Error(),
+	}
+
+	switch v := v.(type) {
+	case map[string]interface{}:
+		if res, ok := v["output"].(string); ok {
+			response.Output = res
+		}
+		if res, ok := v["error"].(string); ok {
+			response.Error = 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.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

@@ -44,6 +44,21 @@ type ResHookStream struct {
 	Output string      `json:"output,omitempty"` // The output
 }
 
+// ResHookDone the response of the done hook
+type ResHookDone struct {
+	Next   *NextAction       `json:"next,omitempty"`
+	Input  []message.Message `json:"input,omitempty"`
+	Output string            `json:"output,omitempty"`
+}
+
+// ResHookFail the response of the fail hook
+type ResHookFail struct {
+	Next   *NextAction       `json:"next,omitempty"`
+	Input  []message.Message `json:"input,omitempty"`
+	Output string            `json:"output,omitempty"`
+	Error  string            `json:"error,omitempty"`
+}
+
 // NextAction the next action
 type NextAction struct {
 	Action  string                 `json:"action"`