From e440f1ff81c0c47d124c1099fffd2b204adde89b Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 13 Jan 2025 17:32:51 +0800 Subject: [PATCH] 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 | 70 +++++++++++++++++++++++------------ neo/assistant/hooks.go | 84 ++++++++++++++++++++++++++++++++++++++++++ neo/assistant/types.go | 15 ++++++++ 3 files changed, 146 insertions(+), 23 deletions(-) diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 5b3052c8..ae09c751 100644 --- a/neo/assistant/api.go +++ b/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,35 +150,27 @@ 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) - if err != nil { - return 0 // break - } - - // 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": + res, err := ast.HookStream(c, ctx, messages, string(*content)) + if err == nil && res != nil { + if res.Output != "" { + value = res.Output + } + if res.Next != nil && res.Next.Action == "exit" { done <- true return 0 // break } + if res.Silent { + return 1 // continue + } } - if !res.Silent { - chatMessage.New(). - Map(map[string]interface{}{ - "text": value, - "done": msg.IsDone, - }). - Write(c.Writer) - } + chatMessage.New(). + Map(map[string]interface{}{ + "text": value, + "done": msg.IsDone, + }). + Write(c.Writer) } // Complete the stream @@ -179,6 +178,31 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [ 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 } diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go index c1425b58..c6dd65b2 100644 --- a/neo/assistant/hooks.go +++ b/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) diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 058e926b..80862a01 100644 --- a/neo/assistant/types.go +++ b/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"`