pkg/aflow: handle per-minute quota violations

1 files changed, 13 insertions, 3 deletions

diff --git a/pkg/aflow/llm_agent.go b/pkg/aflow/llm_agent.go
index f60a19a83..05e875f28 100644
--- a/pkg/aflow/llm_agent.go
+++ b/pkg/aflow/llm_agent.go
@@ -9,6 +9,8 @@ import (
 	"maps"
 	"net/http"
 	"reflect"
+	"regexp"
+	"strconv"
 	"strings"
 	"time"
 
@@ -355,14 +357,22 @@ func (a *LLMAgent) generateContent(ctx *Context, cfg *genai.GenerateContentConfi
 			continue
 		}
 		if err != nil && errors.As(err, &apiErr) && apiErr.Code == http.StatusTooManyRequests &&
-			strings.Contains(apiErr.Message, "Quota exceeded for metric") &&
-			strings.Contains(apiErr.Message, "generate_requests_per_model_per_day") {
-			return resp, &modelQuotaError{ctx.modelName(a.Model)}
+			strings.Contains(apiErr.Message, "Quota exceeded for metric") {
+			if match := rePleaseRetry.FindStringSubmatch(apiErr.Message); match != nil {
+				sec, _ := strconv.Atoi(match[1])
+				time.Sleep(time.Duration(sec+1) * time.Second)
+				continue
+			}
+			if strings.Contains(apiErr.Message, "generate_requests_per_model_per_day") {
+				return resp, &modelQuotaError{ctx.modelName(a.Model)}
+			}
 		}
 		return resp, err
 	}
 }
 
+var rePleaseRetry = regexp.MustCompile("Please retry in ([0-9]+)[.s]")
+
 func (a *LLMAgent) generateContentCached(ctx *Context, cfg *genai.GenerateContentConfig,
 	req []*genai.Content, candidate int) (*genai.GenerateContentResponse, error) {
 	type Cached struct {