From f469568d85e95c8db6964e05013c341198693473 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Tue, 4 Mar 2025 12:03:19 +0800
Subject: [PATCH 1/2] Improve request context handling in chat streaming

- Use request context for client connection tracking
- Replace CloseNotify with context cancellation
- Enhance streaming method to handle request context termination
---
 neo/assistant/api.go | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/neo/assistant/api.go b/neo/assistant/api.go
index f8b0bef8..69bda80c 100644
--- a/neo/assistant/api.go
+++ b/neo/assistant/api.go
@@ -319,7 +319,7 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
 	var result interface{} = nil
 	var err error = nil
 
-	// Chat with AI in background
+	requestCtx := c.Request.Context()
 	go func() {
 		var res interface{} = nil
 		res, err = ast.streamChat(c, ctx, messages, options, clientBreak, contents, callback...)
@@ -338,7 +338,8 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
 			return nil, err
 		}
 		return result, nil
-	case <-c.Writer.CloseNotify():
+
+	case <-requestCtx.Done():
 		clientBreak <- true
 		return nil, nil
 	}

From dafd6b029b0c9327b92197866e7c2d7cc319c129 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Tue, 4 Mar 2025 15:27:40 +0800
Subject: [PATCH 2/2] Enhance retry mechanism with flexible action handling

- Modify HookRetry to support returning NextAction
- Update streamChat to handle different retry response types
- Add type switching for retry hook return values
- Implement execution of NextAction in retry flow
- Improve error handling and response processing during retry
---
 neo/assistant/api.go   | 22 ++++++++++++++++++----
 neo/assistant/hooks.go | 18 +++++++++++++-----
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/neo/assistant/api.go b/neo/assistant/api.go
index 69bda80c..64ffc10a 100644
--- a/neo/assistant/api.go
+++ b/neo/assistant/api.go
@@ -623,19 +623,32 @@ func (ast *Assistant) streamChat(
 		ctx.Retry = true                    // Set the retry mode
 
 		// Hook retry
-		prompt, retryErr := ast.HookRetry(c, ctx, messages, contents, exception.Trim(retry))
+		promptAny, retryErr := ast.HookRetry(c, ctx, messages, contents, exception.Trim(retry))
 		if retryErr != nil {
 			color.Red("%s, try to fix the error %d times, but failed with %s", exception.Trim(retry), ctx.RetryTimes, exception.Trim(retryErr))
 			chatMessage.New().Error(retry.Error()).Done().Callback(cb).Write(c.Writer)
 			return nil, retry
 		}
 
-		// if the prompt is empty, return the error
-		if prompt == "" {
+		if promptAny == nil {
 			chatMessage.New().Error(retry.Error()).Done().Callback(cb).Write(c.Writer)
 			return nil, retry
 		}
 
+		var prompt string = ""
+		switch v := promptAny.(type) {
+		case NextAction:
+			result, err := v.Execute(c, ctx, contents, cb)
+			if err != nil {
+				chatMessage.New().Error(err.Error()).Done().Callback(cb).Write(c.Writer)
+				return nil, retry
+			}
+			return result, nil
+
+		case string:
+			prompt = v
+		}
+
 		// Add the prompt to the messages
 		retryMessages, retryErr := ast.retryMessages(messages, prompt)
 		if retryErr != nil {
@@ -645,7 +658,8 @@ func (ast *Assistant) streamChat(
 		}
 
 		// Retry the chat
-		return ast.streamChat(c, ctx, retryMessages, options, clientBreak, contents, cb)
+		retryContents := chatMessage.NewContents()
+		return ast.execute(c, ctx, retryMessages, options, retryContents, cb)
 	}
 
 	// Handle error
diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go
index 32adc7ce..31aaec2c 100644
--- a/neo/assistant/hooks.go
+++ b/neo/assistant/hooks.go
@@ -135,7 +135,7 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
 }
 
 // HookRetry Handle retry of assistant response
-func (ast *Assistant) HookRetry(c *gin.Context, context chatctx.Context, input []message.Message, contents *chatMessage.Contents, errmsg string) (string, error) {
+func (ast *Assistant) HookRetry(c *gin.Context, context chatctx.Context, input []message.Message, contents *chatMessage.Contents, errmsg string) (interface{}, error) {
 	ctx := ast.createBackgroundContext()
 	output := []message.Data{}
 	if len(input) < 1 {
@@ -158,12 +158,20 @@ func (ast *Assistant) HookRetry(c *gin.Context, context chatctx.Context, input [
 		return "", err
 	}
 
-	res, ok := v.(string)
-	if !ok {
-		return "", fmt.Errorf("invalid return type: %T", v)
+	switch v := v.(type) {
+	case string:
+		return v, nil
+	case map[string]interface{}:
+		var next NextAction
+		raw, _ := jsoniter.MarshalToString(v)
+		err := jsoniter.UnmarshalFromString(raw, &next)
+		if err != nil {
+			return "", err
+		}
+		return next, nil
 	}
 
-	return res, nil
+	return "", nil
 }
 
 // HookDone Handle completion of assistant response