refactor AI response generation and enhance IMAP connection handling

internal/ai/ai.go

@@ -48,52 +48,63 @@ func (a *AI) GenerateReply(emailContent string, contextContent map[string]string
 	}
 
 	// Prepare the system message and user message
-	systemMsg := "You are a helpful assistant. Use the provided context to help answer the email professionally and accurately."
-	userMsg := fmt.Sprintf("Using the following context:\n%s\n\nPlease generate a reply for this email:\n%s", context, emailContent)
+	systemMsg := "You are a helpful assistant responding to emails. Analyze the language of the incoming email and translate the response to the same language. Format your response in HTML. Do not include any explanations or extra text - just write the email response directly in HTML format. Use appropriate HTML tags for formatting."
+	userMsg := fmt.Sprintf("Using the following context:\n%s\n\nPlease analyze the language of and generate a reply in the same language for this email:\n%s", context, emailContent)
 
 	messages := []Message{
 		{Role: "system", Content: systemMsg},
 		{Role: "user", Content: userMsg},
 	}
 
-	reqBody := OpenRouterRequest{
-		Model:       a.config.Model,
-		Messages:    messages,
-		Temperature: a.config.Temperature,
-		MaxTokens:   a.config.MaxTokens,
+	const maxRetries = 3
+	var lastErr error
+
+	for attempt := 0; attempt < maxRetries; attempt++ {
+		reqBody := OpenRouterRequest{
+			Model:       a.config.Model,
+			Messages:    messages,
+			Temperature: a.config.Temperature,
+			MaxTokens:   a.config.MaxTokens,
+		}
+
+		jsonData, err := json.Marshal(reqBody)
+		if err != nil {
+			return "", err
+		}
+
+		req, err := http.NewRequest("POST", "https://openrouter.ai/api/v1/chat/completions", bytes.NewBuffer(jsonData))
+		if err != nil {
+			return "", err
+		}
+
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", a.config.OpenRouterAPIKey))
+
+		resp, err := a.client.Do(req)
+		if err != nil {
+			lastErr = err
+			continue
+		}
+		defer resp.Body.Close()
+
+		if resp.StatusCode != http.StatusOK {
+			lastErr = fmt.Errorf("OpenRouter API returned status code: %d", resp.StatusCode)
+			continue
+		}
+
+		var result OpenRouterResponse
+		if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+			lastErr = err
+			continue
+		}
+
+		if len(result.Choices) == 0 || result.Choices[0].Message.Content == "" {
+			lastErr = fmt.Errorf("empty response received")
+			continue
+		}
+
+		return result.Choices[0].Message.Content, nil
 	}
 
-	jsonData, err := json.Marshal(reqBody)
-	if err != nil {
-		return "", err
-	}
-
-	req, err := http.NewRequest("POST", "https://openrouter.ai/api/v1/chat/completions", bytes.NewBuffer(jsonData))
-	if err != nil {
-		return "", err
-	}
-
-	req.Header.Set("Content-Type", "application/json")
-	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", a.config.OpenRouterAPIKey))
-
-	resp, err := a.client.Do(req)
-	if err != nil {
-		return "", err
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		return "", fmt.Errorf("OpenRouter API returned status code: %d", resp.StatusCode)
-	}
-
-	var result OpenRouterResponse
-	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
-		return "", err
-	}
-
-	if len(result.Choices) == 0 {
-		return "", fmt.Errorf("no response generated")
-	}
-
-	return result.Choices[0].Message.Content, nil
+	return "", fmt.Errorf("failed after %d attempts, last error: %v", maxRetries, lastErr)
 }