From b36fef0e0c48c2a70a3f13efb8ad0923b80f9aac Mon Sep 17 00:00:00 2001
From: BeaconCat <BeaconCat@users.noreply.github.com>
Date: Mon, 13 Apr 2026 11:42:43 +0800
Subject: [PATCH] feat: add --retries flag with exponential backoff for
 transient LLM errors

Retry on timeout, 5xx, and 429 (rate limit) with 1s/2s/4s backoff.
Default 3 retries, configurable via --retries. Context cancellation
is respected between retries.
---
 cmd/membench/llm_client.go | 64 +++++++++++++++++++++++++++++++-------
 cmd/membench/main.go       |  4 +++
 2 files changed, 57 insertions(+), 11 deletions(-)

diff --git a/cmd/membench/llm_client.go b/cmd/membench/llm_client.go
index d32de057d..a6fd58c09 100644
--- a/cmd/membench/llm_client.go
+++ b/cmd/membench/llm_client.go
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"strings"
 	"time"
@@ -17,6 +18,7 @@ type LLMClient struct {
 	Model      string
 	APIKey     string
 	NoThinking bool // send chat_template_kwargs to disable thinking (llama.cpp specific)
+	MaxRetries int  // max retry attempts for transient errors (0 = no retry)
 	Client     *http.Client
 }
 
@@ -27,6 +29,7 @@ type LLMClientOptions struct {
 	APIKey     string
 	Timeout    time.Duration
 	NoThinking bool
+	MaxRetries int // max retry attempts (default 3)
 }
 
 // NewLLMClient creates a client for an OpenAI-compatible chat completion API.
@@ -34,11 +37,16 @@ func NewLLMClient(opts LLMClientOptions) *LLMClient {
 	if opts.Timeout == 0 {
 		opts.Timeout = 120 * time.Second
 	}
+	maxRetries := opts.MaxRetries
+	if maxRetries <= 0 {
+		maxRetries = 3
+	}
 	return &LLMClient{
 		BaseURL:    strings.TrimRight(opts.BaseURL, "/"),
 		Model:      opts.Model,
 		APIKey:     opts.APIKey,
 		NoThinking: opts.NoThinking,
+		MaxRetries: maxRetries,
 		Client: &http.Client{
 			Timeout: opts.Timeout,
 		},
@@ -101,19 +109,53 @@ func (c *LLMClient) Complete(ctx context.Context, systemPrompt, userPrompt strin
 		req.Header.Set("Authorization", "Bearer "+c.APIKey)
 	}
 
-	resp, err := c.Client.Do(req)
-	if err != nil {
-		return "", fmt.Errorf("http request: %w", err)
-	}
-	defer resp.Body.Close()
+	var respBody []byte
+	var lastErr error
+	for attempt := 0; attempt <= c.MaxRetries; attempt++ {
+		if attempt > 0 {
+			backoff := time.Duration(1<<(attempt-1)) * time.Second // 1s, 2s, 4s, ...
+			log.Printf("LLM retry %d/%d after %v: %v", attempt, c.MaxRetries, backoff, lastErr)
+			select {
+			case <-ctx.Done():
+				return "", ctx.Err()
+			case <-time.After(backoff):
+			}
+			// Rebuild request (body reader is consumed)
+			req, err = http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(jsonBody))
+			if err != nil {
+				return "", fmt.Errorf("create request: %w", err)
+			}
+			req.Header.Set("Content-Type", "application/json")
+			if c.APIKey != "" {
+				req.Header.Set("Authorization", "Bearer "+c.APIKey)
+			}
+		}
 
-	respBody, err := io.ReadAll(resp.Body)
-	if err != nil {
-		return "", fmt.Errorf("read response: %w", err)
-	}
+		var resp *http.Response
+		resp, lastErr = c.Client.Do(req)
+		if lastErr != nil {
+			continue // network/timeout error → retry
+		}
 
-	if resp.StatusCode != 200 {
-		return "", fmt.Errorf("API error %d: %s", resp.StatusCode, string(respBody))
+		respBody, lastErr = io.ReadAll(resp.Body)
+		resp.Body.Close()
+		if lastErr != nil {
+			continue
+		}
+
+		if resp.StatusCode == 429 || resp.StatusCode >= 500 {
+			lastErr = fmt.Errorf("API error %d: %s", resp.StatusCode, string(respBody))
+			continue // rate limit or server error → retry
+		}
+		if resp.StatusCode != 200 {
+			return "", fmt.Errorf("API error %d: %s", resp.StatusCode, string(respBody))
+		}
+
+		lastErr = nil
+		break
+	}
+	if lastErr != nil {
+		return "", fmt.Errorf("after %d retries: %w", c.MaxRetries, lastErr)
 	}
 
 	var chatResp chatResponse
diff --git a/cmd/membench/main.go b/cmd/membench/main.go
index 58930c675..ed61feeff 100644
--- a/cmd/membench/main.go
+++ b/cmd/membench/main.go
@@ -27,6 +27,7 @@ var (
 	flagNoThinking bool
 	flagLimit      int
 	flagTimeout    int
+	flagRetries    int
 )
 
 func main() {
@@ -66,6 +67,7 @@ func main() {
 		BoolVar(&flagNoThinking, "no-thinking", false, "disable thinking mode via chat_template_kwargs (llama.cpp + Qwen)")
 	evalCmd.Flags().IntVar(&flagLimit, "limit", 0, "max QA questions per sample (0 = all)")
 	evalCmd.Flags().IntVar(&flagTimeout, "timeout", 120, "HTTP timeout in seconds for LLM requests")
+	evalCmd.Flags().IntVar(&flagRetries, "retries", 3, "max retry attempts for transient LLM errors (timeout/5xx/429)")
 
 	reportCmd := &cobra.Command{
 		Use:   "report",
@@ -93,6 +95,7 @@ func main() {
 		BoolVar(&flagNoThinking, "no-thinking", false, "disable thinking mode via chat_template_kwargs (llama.cpp + Qwen)")
 	runCmd.Flags().IntVar(&flagLimit, "limit", 0, "max QA questions per sample (0 = all)")
 	runCmd.Flags().IntVar(&flagTimeout, "timeout", 120, "HTTP timeout in seconds for LLM requests")
+	runCmd.Flags().IntVar(&flagRetries, "retries", 3, "max retry attempts for transient LLM errors (timeout/5xx/429)")
 
 	rootCmd.AddCommand(ingestCmd, evalCmd, reportCmd, runCmd)
 
@@ -317,5 +320,6 @@ func buildLLMOptions() (LLMClientOptions, error) {
 		APIKey:     apiKey,
 		NoThinking: flagNoThinking,
 		Timeout:    time.Duration(flagTimeout) * time.Second,
+		MaxRetries: flagRetries,
 	}, nil
 }