From bc13b0d83fd51b1bf4c034057878729faaaa8323 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Tue, 29 Apr 2025 20:35:27 +0800
Subject: [PATCH] feat: Improve error handling and result processing in
 Assistant API

- Update streamChat method to set a default prompt for error messages, enhancing user guidance.
- Refactor case handling in streamChat to better manage boolean and map results, improving response clarity.
- Modify HookRetry method to streamline result extraction and handle action presence more effectively, ensuring robust error handling.
---
 neo/assistant/api.go   | 15 ++++++++++++---
 neo/assistant/hooks.go | 23 ++++++++++++++++-------
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/neo/assistant/api.go b/neo/assistant/api.go
index 86d5106b..d52849ac 100644
--- a/neo/assistant/api.go
+++ b/neo/assistant/api.go
@@ -687,9 +687,18 @@ func (ast *Assistant) streamChat(
 			return nil, retry
 		}
 
-		var prompt string = ""
+		// Default prompt
+		var prompt string = fmt.Sprintf("Try to fix the error following the error message. error:\n %s", exception.Trim(retry))
 		switch v := promptAny.(type) {
-		case NextAction:
+		case bool: // Ignore the error, and return the specific result
+			if v == false {
+				return nil, retry
+			}
+
+		case map[string]interface{}: // Ignore the error, and return the specific result
+			return v, nil
+
+		case NextAction: // Execute the next action
 			result, err := v.Execute(c, ctx, contents, cb)
 			if err != nil {
 				// chatMessage.New().Error(err.Error()).Done().Callback(cb).Write(c.Writer)
@@ -697,7 +706,7 @@ func (ast *Assistant) streamChat(
 			}
 			return result, nil
 
-		case string:
+		case string: // Add the prompt to the messages
 			prompt = v
 		}
 
diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go
index 76ffad8f..acf39281 100644
--- a/neo/assistant/hooks.go
+++ b/neo/assistant/hooks.go
@@ -164,16 +164,25 @@ func (ast *Assistant) HookRetry(c *gin.Context, context chatctx.Context, input [
 	}
 
 	switch v := v.(type) {
-	case string:
+	case string, bool:
 		return v, nil
+
 	case map[string]interface{}:
-		var next NextAction
-		raw, _ := jsoniter.MarshalToString(v)
-		err := jsoniter.UnmarshalFromString(raw, &next)
-		if err != nil {
-			return nil, err
+
+		// Has Action
+		if _, has := v["action"]; has {
+			var next NextAction
+			raw, _ := jsoniter.MarshalToString(v)
+			err := jsoniter.UnmarshalFromString(raw, &next)
+			if err != nil {
+				return nil, err
+			}
+			return &next, nil
 		}
-		return &next, nil
+
+		// Ignore the error, and return the specific result
+		return v, nil
+
 	}
 
 	return nil, nil