fix: Don't retry OpenAI calls on 400 errors; Better Error Messages (#213)

* better error messages

* don't retry on 400 errors

* BadRequest errors should be HTTP400 responses

Author: danielchalef

pkg/extractors/summarizer.go

@@ -316,7 +316,10 @@ func incrementalSummarizer(
 	return summary, tokensUsed, nil
 }
 
-func generateProgressiveSummarizerPrompt(appState *models.AppState, promptData SummaryPromptTemplateData) (string, error) {
+func generateProgressiveSummarizerPrompt(
+	appState *models.AppState,
+	promptData SummaryPromptTemplateData,
+) (string, error) {
 	customSummaryPromptTemplateAnthropic := appState.Config.CustomPrompts.SummarizerPrompts.Anthropic
 	customSummaryPromptTemplateOpenAI := appState.Config.CustomPrompts.SummarizerPrompts.OpenAI
 

pkg/llms/llm_base.go

@@ -3,6 +3,7 @@ package llms
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"time"
 
 	"github.com/getzep/zep/pkg/models"
@@ -131,13 +132,32 @@ func Float64ToFloat32Matrix(in [][]float64) [][]float32 {
 }
 
 func NewRetryableHTTPClient(retryMax int, timeout time.Duration) *retryablehttp.Client {
-	retryableHttpClient := retryablehttp.NewClient()
-	retryableHttpClient.RetryMax = retryMax
-	retryableHttpClient.HTTPClient.Timeout = timeout
-	retryableHttpClient.Logger = log
-	retryableHttpClient.Backoff = retryablehttp.DefaultBackoff
+	retryableHTTPClient := retryablehttp.NewClient()
+	retryableHTTPClient.RetryMax = retryMax
+	retryableHTTPClient.HTTPClient.Timeout = timeout
+	retryableHTTPClient.Logger = log
+	retryableHTTPClient.Backoff = retryablehttp.DefaultBackoff
+	retryableHTTPClient.CheckRetry = retryPolicy
+
+	return retryableHTTPClient
+}
+
+// retryPolicy is a retryablehttp.CheckRetry function. It is used to determine
+// whether a request should be retried or not.
+func retryPolicy(ctx context.Context, resp *http.Response, err error) (bool, error) {
+	// do not retry on context.Canceled or context.DeadlineExceeded
+	if ctx.Err() != nil {
+		return false, ctx.Err()
+	}
+
+	// Do not retry 400 errors as they're used by OpenAI to indicate maximum
+	// context length exceeded
+	if resp != nil && resp.StatusCode == 400 {
+		return false, err
+	}
 
-	return retryableHttpClient
+	shouldRetry, _ := retryablehttp.DefaultRetryPolicy(ctx, resp, err)
+	return shouldRetry, nil
 }
 
 // useOpenAIEmbeddings is true if OpenAI embeddings are enabled

pkg/server/handlertools/tools.go

@@ -9,6 +9,7 @@ import (
 	"strings"
 
 	"github.com/getzep/zep/internal"
+	"github.com/getzep/zep/pkg/models"
 
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
@@ -71,7 +72,7 @@ func RenderError(w http.ResponseWriter, err error, status int) {
 		// Don't log not found errors
 		log.Error(err)
 	}
-	if strings.Contains(err.Error(), "is deleted") {
+	if strings.Contains(err.Error(), "is deleted") || errors.Is(err, models.ErrBadRequest) {
 		status = http.StatusBadRequest
 	}
 	http.Error(w, err.Error(), status)

pkg/store/postgres/documents.go

@@ -11,6 +11,7 @@ import (
 	"github.com/getzep/zep/pkg/models"
 	"github.com/google/uuid"
 	"github.com/uptrace/bun"
+	"github.com/uptrace/bun/driver/pgdriver"
 )
 
 func NewDocumentCollectionDAO(
@@ -65,8 +66,8 @@ func (dc *DocumentCollectionDAO) Create(
 		Returning("*").
 		Exec(ctx)
 	if err != nil {
-		if strings.Contains(err.Error(), "duplicate key value violates unique constraint") {
-			return fmt.Errorf("collection with name %s already exists", dc.Name)
+		if err, ok := err.(pgdriver.Error); ok && err.IntegrityViolation() {
+			return models.NewBadRequestError("collection already exists: " + dc.Name)
 		}
 		return fmt.Errorf("failed to insert collection: %w", err)
 	}
@@ -270,6 +271,9 @@ func (dc *DocumentCollectionDAO) CreateDocuments(
 		Returning("uuid").
 		Exec(ctx)
 	if err != nil {
+		if err, ok := err.(pgdriver.Error); ok && err.IntegrityViolation() {
+			return nil, models.NewBadRequestError("document_id already exists")
+		}
 		if strings.Contains(err.Error(), "different vector dimensions") {
 			return nil, store.NewEmbeddingMismatchError(err)
 		}

pkg/store/postgres/session.go

@@ -5,11 +5,13 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"strings"
 	"sync"
 	"time"
 
 	"github.com/getzep/zep/pkg/models"
 	"github.com/uptrace/bun"
+	"github.com/uptrace/bun/driver/pgdriver"
 )
 
 var _ models.SessionManager = &SessionDAO{}
@@ -47,6 +49,16 @@ func (dao *SessionDAO) Create(
 		Returning("*").
 		Exec(ctx)
 	if err != nil {
+		if err, ok := err.(pgdriver.Error); ok && err.IntegrityViolation() {
+			if strings.Contains(err.Error(), "user") {
+				return nil, models.NewBadRequestError(
+					"user does not exist with user_id: " + *session.UserID,
+				)
+			}
+			return nil, models.NewBadRequestError(
+				"session already exists with session_id: " + session.SessionID,
+			)
+		}
 		return nil, fmt.Errorf("failed to create session: %w", err)
 	}
 

pkg/store/postgres/userstore.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/getzep/zep/pkg/models"
 	"github.com/uptrace/bun"
+	"github.com/uptrace/bun/driver/pgdriver"
 )
 
 var _ models.UserStore = &UserStoreDAO{}
@@ -37,6 +38,11 @@ func (dao *UserStoreDAO) Create(
 	}
 	_, err := dao.db.NewInsert().Model(userDB).Returning("*").Exec(ctx)
 	if err != nil {
+		if err, ok := err.(pgdriver.Error); ok && err.IntegrityViolation() {
+			return nil, models.NewBadRequestError(
+				"user already exists with user_id: " + user.UserID,
+			)
+		}
 		return nil, err
 	}
 

pkg/store/postgres/userstore_test.go

@@ -3,6 +3,7 @@ package postgres
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/getzep/zep/pkg/testutils"
 
@@ -66,6 +67,9 @@ func TestUserStoreDAO(t *testing.T) {
 		createdUser, err := userStore.Create(ctx, user)
 		assert.NoError(t, err)
 
+		// Wait a second
+		<-time.After(1 * time.Second)
+
 		// Update the user with zero values
 		userUpdate := &models.UpdateUserRequest{
 			UserID:    user.UserID,