better error handling

internal/gpt/common.go

@@ -0,0 +1,42 @@
+package gpt
+
+import (
+	"time"
+
+	"github.com/sashabaranov/go-openai"
+)
+
+func sleepBeforeRetry(i int) {
+	time.Sleep(time.Duration(i*3) * time.Second)
+}
+
+func isServiceUnavailableError(err error) bool {
+	e, ok := isOpenAIError(err)
+	if !ok {
+		return false
+	}
+
+	if e.HTTPStatusCode == 503 {
+		return true
+	}
+
+	return false
+}
+
+func isTokenExceededError(err error) bool {
+	e, ok := isOpenAIError(err)
+	if !ok {
+		return false
+	}
+
+	if s, ok := e.Code.(string); ok && s == "context_length_exceeded" {
+		return true
+	}
+
+	return false
+}
+
+func isOpenAIError(err error) (*openai.APIError, bool) {
+	e, ok := err.(*openai.APIError)
+	return e, ok
+}

internal/gpt/completions.go

@@ -2,8 +2,6 @@ package gpt
 
 import (
 	"context"
-	"errors"
-	"time"
 
 	"github.com/sashabaranov/go-openai"
 )
@@ -16,68 +14,48 @@ func (g *Gpt) CreateCompletion(ctx context.Context, history []openai.ChatComplet
 		Content: userMsg,
 	})
 
-	response, err := g.createCompletionWithTimeout(ctx, messageHistory)
+	res, err := g.complReqWithTimeout(ctx, messageHistory)
 	if err != nil {
 		return []openai.ChatCompletionMessage{}, err
 	}
 
 	messageHistory = append(messageHistory, openai.ChatCompletionMessage{
 		Role:    openai.ChatMessageRoleAssistant,
-		Content: response.Choices[0].Message.Content,
+		Content: res.Choices[0].Message.Content,
 	})
 
 	return messageHistory, err
 }
 
-// createCompletionWithTimeout makes a request to get a GPT completion with a specified timeout.
-func (g *Gpt) createCompletionWithTimeout(ctx context.Context, msg []openai.ChatCompletionMessage) (openai.ChatCompletionResponse, error) {
-	var response openai.ChatCompletionResponse
+// complReqWithTimeout makes a request to get a GPT completion with a specified timeout.
+func (g *Gpt) complReqWithTimeout(ctx context.Context, msg []openai.ChatCompletionMessage) (openai.ChatCompletionResponse, error) {
+	var res openai.ChatCompletionResponse
 	var err error
 
-	ctx, cancel := context.WithTimeout(ctx, g.gptTimeout)
-	defer cancel()
-
 	for i := 0; i < g.maxAttempts; i++ {
-		response, err = g.client.CreateChatCompletion(
+		ctx, cancel := context.WithTimeout(ctx, g.gptTimeout)
+		defer cancel()
+
+		res, err = g.client.CreateChatCompletion(
 			ctx,
 			openai.ChatCompletionRequest{
 				Model:    g.model,
 				Messages: msg,
 			},
 		)
-		if ctx.Err() == context.Canceled {
-			return openai.ChatCompletionResponse{}, ctx.Err()
-		}
-		if err == nil {
-			if len(response.Choices) < 1 || response.Choices[0].Message.Content == "" {
-				err = errors.New("empty response")
-			} else {
-				return response, nil
-			}
-		}
 
-		if isTokenExceededError(err) {
+		if ctx.Err() == context.Canceled {
+			return res, ctx.Err()
+		} else if isTokenExceededError(err) {
 			msg = trimFirstMsgFromHistory(msg)
-			continue
+		} else if !isServiceUnavailableError(err) {
+			break
 		}
 
-		time.Sleep(time.Duration(i*3) * time.Second)
-	}
-
-	return openai.ChatCompletionResponse{}, err
-}
-
-func isTokenExceededError(err error) bool {
-	e, ok := err.(*openai.APIError)
-	if !ok {
-		return false
-	}
-
-	if s, ok := e.Code.(string); ok && s == "context_length_exceeded" {
-		return true
+		sleepBeforeRetry(i)
 	}
 
-	return false
+	return res, err
 }
 
 func trimFirstMsgFromHistory(msg []openai.ChatCompletionMessage) []openai.ChatCompletionMessage {

internal/gpt/images.go

@@ -2,38 +2,35 @@ package gpt
 
 import (
 	"context"
-	"errors"
-	"time"
 
 	"github.com/sashabaranov/go-openai"
 )
 
 // CreateImage makes a request to get a DALL-E image URL.
 func (g *Gpt) CreateImage(ctx context.Context, prompt string) (openai.ImageResponse, error) {
-	var response openai.ImageResponse
+	var res openai.ImageResponse
 	var err error
 
-	ctx, cancel := context.WithTimeout(ctx, g.gptTimeout)
-	defer cancel()
-
 	for i := 0; i < g.maxAttempts; i++ {
-		response, err = g.client.CreateImage(
+		ctx, cancel := context.WithTimeout(ctx, g.gptTimeout)
+		defer cancel()
+
+		res, err = g.client.CreateImage(
 			ctx,
 			openai.ImageRequest{
 				Prompt:         prompt,
 				Size:           "1024x1024",
 				ResponseFormat: openai.CreateImageResponseFormatURL,
 			},
 		)
+
 		if ctx.Err() == context.Canceled {
-			return openai.ImageResponse{}, ctx.Err()
+			return res, ctx.Err()
+		} else if !isServiceUnavailableError(err) {
+			break
 		}
-		if err == nil && len(response.Data) < 1 {
-			err = errors.New("empty response")
-		} else if err == nil {
-			return response, nil
-		}
-		time.Sleep(time.Duration(i*2) * time.Second)
+
+		sleepBeforeRetry(i)
 	}
 
 	return openai.ImageResponse{}, err