Merge pull request #889 from trheyi/main

Improved Request Context Handling and Enhanced Retry Mechanism

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
 	}
@@ -622,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 {
@@ -644,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

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