diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 4e275572..4bd611bc 100644 --- a/neo/assistant/api.go +++ b/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 { diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go index e099ee10..32adc7ce 100644 --- a/neo/assistant/hooks.go +++ b/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 diff --git a/neo/assistant/object.go b/neo/assistant/object.go index c97041c3..5d672d2f 100644 --- a/neo/assistant/object.go +++ b/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()) diff --git a/neo/context/context.go b/neo/context/context.go index 8d0d9442..77d1c43f 100644 --- a/neo/context/context.go +++ b/neo/context/context.go @@ -22,8 +22,9 @@ type Context struct { Namespace string `json:"namespace,omitempty"` Config map[string]interface{} `json:"config,omitempty"` Signal interface{} `json:"signal,omitempty"` - Silent bool `json:"silent,omitempty"` // Silent mode - Retry bool `json:"retry,omitempty"` // Retry mode + 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 }