From 1ecddf59dfc509a65af1cfe99c833dfbe3ca551a Mon Sep 17 00:00:00 2001
From: shenlan <manbuzhe2009@qq.com>
Date: Sun, 3 Aug 2025 18:50:22 +0800
Subject: [PATCH] fix: attempt callChutes initial request plus retries

---
 server/api/askai.go | 120 ++++++++++++++++++++++++++++----------------
 1 file changed, 76 insertions(+), 44 deletions(-)

diff --git a/server/api/askai.go b/server/api/askai.go
index b77a3a4..fa59e7e 100644
--- a/server/api/askai.go
+++ b/server/api/askai.go
@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"os"
 	"path/filepath"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	"gopkg.in/yaml.v3"
@@ -41,40 +42,53 @@ const chutesURL = "https://llm.chutes.ai/v1/chat/completions"
 type serverConfig struct {
 	Env   map[string]string `yaml:"env"`
 	Model []string          `yaml:"model"`
+	AskAI struct {
+		Timeout int `yaml:"timeout"` // seconds
+		Retries int `yaml:"retries"`
+	} `yaml:"askai"`
 }
 
-// loadTokenAndModel attempts to read CHUTES_API_TOKEN and model from
+// loadConfig attempts to read CHUTES_API_TOKEN, model, timeout and retries from
 // environment variables, falling back to config/server.yaml.
-func loadTokenAndModel() (string, string) {
+func loadConfig() (string, string, time.Duration, int) {
 	token := os.Getenv("CHUTES_API_TOKEN")
 	model := os.Getenv("CHUTES_API_MODEL")
-	if token != "" && model != "" {
-		return token, model
-	}
+	timeout := 30 * time.Second
+	retries := 3
 	path := filepath.Join("server", "config", "server.yaml")
 	data, err := os.ReadFile(path)
-	if err != nil {
-		return token, model
+	if err == nil {
+		var cfg serverConfig
+		if err := yaml.Unmarshal(data, &cfg); err == nil {
+			if token == "" {
+				token = cfg.Env["CHUTES_API_TOKEN"]
+			}
+			if model == "" && len(cfg.Model) > 0 {
+				model = cfg.Model[0]
+			}
+			if cfg.AskAI.Timeout > 0 {
+				timeout = time.Duration(cfg.AskAI.Timeout) * time.Second
+			}
+			if cfg.AskAI.Retries > 0 {
+				retries = cfg.AskAI.Retries
+			}
+		}
 	}
-	var cfg serverConfig
-	if err := yaml.Unmarshal(data, &cfg); err != nil {
-		return token, model
+	if timeout > 30*time.Second {
+		timeout = 30 * time.Second
 	}
-	if token == "" {
-		token = cfg.Env["CHUTES_API_TOKEN"]
-	}
-	if model == "" && len(cfg.Model) > 0 {
-		model = cfg.Model[0]
+	if retries > 3 {
+		retries = 3
 	}
 	if model == "" {
 		model = "deepseek-ai/DeepSeek-R1"
 	}
-	return token, model
+	return token, model, timeout, retries
 }
 
 // callChutes sends the question to the hosted LLM service and returns the reply.
 func callChutes(question string) (string, error) {
-	token, model := loadTokenAndModel()
+	token, model, timeout, retries := loadConfig()
 	if token == "" {
 		return "", errors.New("CHUTES_API_TOKEN not set")
 	}
@@ -94,35 +108,53 @@ func callChutes(question string) (string, error) {
 	if err != nil {
 		return "", err
 	}
-	req, err := http.NewRequest("POST", url, bytes.NewBuffer(data))
-	if err != nil {
-		return "", err
-	}
-	req.Header.Set("Authorization", "Bearer "+token)
-	req.Header.Set("Content-Type", "application/json")
 
-	resp, err := http.DefaultClient.Do(req)
-	if err != nil {
-		return "", err
-	}
-	defer resp.Body.Close()
-	if resp.StatusCode != http.StatusOK {
-		b, _ := io.ReadAll(resp.Body)
-		return "", fmt.Errorf("chutes API error: %s", string(b))
-	}
+	client := &http.Client{Timeout: timeout}
+	var lastErr error
+	for i := 0; i <= retries; i++ {
+		req, err := http.NewRequest("POST", url, bytes.NewBuffer(data))
+		if err != nil {
+			return "", err
+		}
+		req.Header.Set("Authorization", "Bearer "+token)
+		req.Header.Set("Content-Type", "application/json")
 
-	var res struct {
-		Choices []struct {
-			Message struct {
-				Content string `json:"content"`
-			} `json:"message"`
-		} `json:"choices"`
+		resp, err := client.Do(req)
+		if err != nil {
+			lastErr = err
+			continue
+		}
+
+		b, err := io.ReadAll(resp.Body)
+		resp.Body.Close()
+		if err != nil {
+			lastErr = err
+			continue
+		}
+		if resp.StatusCode != http.StatusOK {
+			lastErr = fmt.Errorf("chutes API error: %s", string(b))
+			continue
+		}
+
+		var res struct {
+			Choices []struct {
+				Message struct {
+					Content string `json:"content"`
+				} `json:"message"`
+			} `json:"choices"`
+		}
+		if err := json.Unmarshal(b, &res); err != nil {
+			lastErr = err
+			continue
+		}
+		if len(res.Choices) == 0 {
+			lastErr = errors.New("no choices returned")
+			continue
+		}
+		return res.Choices[0].Message.Content, nil
 	}
-	if err := json.NewDecoder(resp.Body).Decode(&res); err != nil {
-		return "", err
+	if lastErr == nil {
+		lastErr = errors.New("request failed")
 	}
-	if len(res.Choices) == 0 {
-		return "", errors.New("no choices returned")
-	}
-	return res.Choices[0].Message.Content, nil
+	return "", lastErr
 }