pkg/aflow: retry LLM requests on transient errors

Update google#6573

Author: dvyukov

pkg/aflow/llm_agent.go

@@ -7,6 +7,8 @@ import (
 	"fmt"
 	"maps"
 	"reflect"
+	"strings"
+	"time"
 
 	"github.com/google/syzkaller/pkg/aflow/trajectory"
 	"google.golang.org/genai"
@@ -98,7 +100,7 @@ func (a *LLMAgent) chat(ctx *Context, cfg *genai.GenerateContentConfig, tools ma
 		if err := ctx.startSpan(reqSpan); err != nil {
 			return "", nil, err
 		}
-		resp, err := ctx.generateContent(cfg, req)
+		resp, err := a.generateContent(ctx, cfg, req)
 		if err != nil {
 			return "", nil, ctx.finishSpan(reqSpan, err)
 		}
@@ -216,6 +218,21 @@ func (a *LLMAgent) parseResponse(resp *genai.GenerateContentResponse) (
 	return
 }
 
+func (a *LLMAgent) generateContent(ctx *Context, cfg *genai.GenerateContentConfig,
+	req []*genai.Content) (*genai.GenerateContentResponse, error) {
+	backoff := time.Second
+	for try := 0; ; try++ {
+		resp, err := ctx.generateContent(cfg, req)
+		if err != nil && try < 100 &&
+			strings.Contains(err.Error(), "Error 503, Message: The model is overloaded. Please try again later.") {
+			time.Sleep(backoff)
+			backoff = min(backoff+time.Second, 10*time.Second)
+			continue
+		}
+		return resp, err
+	}
+}
+
 func (a *LLMAgent) verify(vctx *verifyContext) {
 	vctx.requireNotEmpty(a.Name, "Name", a.Name)
 	vctx.requireNotEmpty(a.Name, "Reply", a.Reply)