From 0a7bbd79ddce993add61e99ffe0e9983dd56257d Mon Sep 17 00:00:00 2001 From: Dmitry Vyukov Date: Mon, 26 Jan 2026 15:07:37 +0100 Subject: pkg/aflow: handle per-minute quota violations --- pkg/aflow/llm_agent.go | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) (limited to 'pkg') 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 { -- cgit mrf-deployment