Merge pull request #886 from trheyi/main

Add retry mechanism for chat streaming with error handling

neo/assistant/api.go

@@ -7,9 +7,11 @@ import (
 	"os"
 	"strings"
 
+	"github.com/fatih/color"
 	"github.com/gin-gonic/gin"
 	jsoniter "github.com/json-iterator/go"
 	"github.com/yaoapp/gou/fs"
+	"github.com/yaoapp/kun/exception"
 	"github.com/yaoapp/kun/log"
 	chatctx "github.com/yaoapp/yao/neo/context"
 	chatMessage "github.com/yaoapp/yao/neo/message"
@@ -365,9 +367,11 @@ func (ast *Assistant) streamChat(
 
 	toolsCount := 0
 	currentMessageID := ""
+	var retry error = nil
 	var result interface{} = nil // To save the result
 	var content string = ""      // To save the content
 	err := ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
+
 		select {
 		case <-clientBreak:
 			return 0 // break
@@ -568,7 +572,7 @@ func (ast *Assistant) streamChat(
 
 				// Some error occurred in the hook, return the error
 				if hookErr != nil {
-					chatMessage.New().Error(hookErr.Error()).Done().Callback(cb).Write(c.Writer)
+					retry = hookErr
 					return 0 // break
 				}
 
@@ -610,6 +614,38 @@ func (ast *Assistant) streamChat(
 		}
 	})
 
+	// retry
+	if retry != nil {
+
+		// Update the retry times
+		ctx.RetryTimes = ctx.RetryTimes + 1 // Increment the retry times
+		ctx.Retry = true                    // Set the retry mode
+
+		// Hook retry
+		prompt, 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 == "" {
+			return nil, retry
+		}
+
+		// Add the prompt to the messages
+		retryMessages, retryErr := ast.retryMessages(messages, prompt)
+		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
+		}
+
+		// Retry the chat
+		return ast.streamChat(c, ctx, retryMessages, options, clientBreak, contents, cb)
+	}
+
 	// Handle error
 	if err != nil {
 		return nil, err
@@ -635,6 +671,27 @@ func (ast *Assistant) streamChat(
 	return strings.TrimSpace(content), nil
 }
 
+func (ast *Assistant) retryMessages(messages []chatMessage.Message, prompt string) ([]chatMessage.Message, error) {
+
+	// Get the last user message
+	var lastIndex int
+	for i := len(messages) - 1; i >= 0; i-- {
+		if messages[i].Role == "user" {
+			messages[i].Text = prompt
+			lastIndex = i
+			break
+		}
+	}
+
+	if lastIndex == 0 {
+		return nil, fmt.Errorf("no user message found")
+	}
+
+	// Remove the messages after the last user message
+	messages = messages[:lastIndex+1]
+	return messages, nil
+}
+
 // saveChatHistory saves the chat history if storage is available
 func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessage.Message, contents *chatMessage.Contents) {
 	if len(contents.Data) > 0 && ctx.Sid != "" && len(messages) > 0 {

neo/assistant/hooks.go

@@ -134,6 +134,38 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
 	return response, nil
 }
 
+// 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) {
+	ctx := ast.createBackgroundContext()
+	output := []message.Data{}
+	if len(input) < 1 {
+		return "", fmt.Errorf("no input")
+	}
+
+	var lastInput message.Message = input[len(input)-1]
+	for _, data := range contents.Data {
+		if data.Type == "think" {
+			continue
+		}
+		output = append(output, data)
+	}
+
+	v, err := ast.call(ctx, "Retry", c, contents, context, lastInput.String(), output, errmsg)
+	if err != nil {
+		if err.Error() == HookErrorMethodNotFound {
+			return "", nil
+		}
+		return "", err
+	}
+
+	res, ok := v.(string)
+	if !ok {
+		return "", fmt.Errorf("invalid return type: %T", v)
+	}
+
+	return res, nil
+}
+
 // HookDone Handle completion of assistant response
 func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, contents *chatMessage.Contents) (*ResHookDone, error) {
 	// Create timeout context

neo/assistant/object.go

@@ -9,6 +9,7 @@ import (
 	chatctx "github.com/yaoapp/yao/neo/context"
 	"github.com/yaoapp/yao/neo/message"
 	chatMessage "github.com/yaoapp/yao/neo/message"
+	sui "github.com/yaoapp/yao/sui/core"
 	"rogchap.com/v8go"
 )
 
@@ -69,6 +70,9 @@ func (ast *Assistant) InitObject(v8ctx *v8.Context, c *gin.Context, context chat
 	v8ctx.WithFunction("Get", jsGet)
 	v8ctx.WithFunction("Del", jsDel)
 	v8ctx.WithFunction("Clear", jsClear)
+
+	// Template methods
+	v8ctx.WithFunction("Replace", jsReplace)
 }
 
 // jsSet function, set a value to the shared space
@@ -301,6 +305,40 @@ func jsSend(info *v8go.FunctionCallbackInfo) *v8go.Value {
 	}
 }
 
+func jsReplace(info *v8go.FunctionCallbackInfo) *v8go.Value {
+	args := info.Args()
+	if len(args) < 2 {
+		return bridge.JsException(info.Context(), "Replace requires at least two arguments")
+	}
+
+	if !args[0].IsString() {
+		return bridge.JsException(info.Context(), "the first argument must be a string")
+	}
+	tmpl := args[0].String()
+
+	raw, err := bridge.GoValue(args[1], info.Context())
+	if err != nil {
+		return bridge.JsException(info.Context(), err.Error())
+	}
+
+	data, ok := raw.(map[string]interface{})
+	if !ok {
+		return bridge.JsException(info.Context(), "the second argument must be a map")
+	}
+
+	replaced, _ := sui.Data(data).Replace(tmpl)
+	if err != nil {
+		return bridge.JsException(info.Context(), err.Error())
+	}
+
+	jsReplaced, err := bridge.JsValue(info.Context(), replaced)
+	if err != nil {
+		return bridge.JsException(info.Context(), err.Error())
+	}
+
+	return jsReplaced
+}
+
 // global get the global variables
 func global(info *v8go.FunctionCallbackInfo) (global *GlobalVariables, err error) {
 	return getGlobal(info.Context(), info.This())

neo/context/context.go

@@ -24,6 +24,7 @@ type Context struct {
 	Signal      interface{}            `json:"signal,omitempty"`
 	Silent      bool                   `json:"silent,omitempty"`      // Silent mode
 	Retry       bool                   `json:"retry,omitempty"`       // Retry mode
+	RetryTimes  uint8                  `json:"retry_times,omitempty"` // Retry times
 	Upload      *FileUpload            `json:"upload,omitempty"`
 	Version     bool                   `json:"version,omitempty"` // Version support
 	RAG         bool                   `json:"rag,omitempty"`     // RAG support
@@ -119,6 +120,14 @@ func (ctx *Context) Map() map[string]interface{} {
 		data["silent"] = ctx.Silent
 	}
 
+	// Retry mode
+	if ctx.Retry {
+		data["retry"] = ctx.Retry
+	}
+
+	// Retry times
+	data["retry_times"] = ctx.RetryTimes
+
 	if ctx.Path != "" {
 		data["pathname"] = ctx.Path
 	}