Repository: weaviate/weaviate
Pull request #7562: Handle missing or not parseable numbers from openai
Author: login: dirkkul
Created: 2025-03-18T08:36:30Z | Merged: 2025-03-18T10:23:44Z

### What's being changed:

The batching-vectorization algorithm uses the remaining number of tokens/requests to decide if requests to the vectorizer can be made. With openAI these numbers are part of the header of every response we get from them. However, we noticed that this number sometimes jumps to 0, which could happen if the header does not include the given values or they are not parseable. This PR ignores the update if we cannot get the values from the header.

### Review checklist

- [ ] Documentation has been updated, if necessary. Link to changed documentation:
- [ ] Chaos pipeline run or not necessary. Link to pipeline:
- [ ] All new code is covered by tests where it is reasonable.
- [ ] Performance tests have been run or not necessary.

<!-- Uncomment the following section if this PR requires changes in related projects (e.g., documentation, client libraries).

GitHub actions will automatically create an issue in the corresponding repository for each checked box below. (See `.github/workflows/create-cross-functional-issues.yml`)

### Cross-functional impact

- [ ] This change requires public documentation (weaviate-io) to be updated. Check the box to automatically create a corresponding issue.
- Does it require a change in the client libraries? If yes, please check the boxes for the affected client libraries.
    - [ ] Python (weaviate-python-client)
    - [ ] JavaScript/TypeScript (typescript-client)
    - [ ] Go (weaviate-go-client)
    - [ ] Java (java-client)

-->



Diff:
diff --git a/modules/text2vec-openai/clients/openai.go b/modules/text2vec-openai/clients/openai.go
index 2c257eeabc1..b318af3385c 100644
--- a/modules/text2vec-openai/clients/openai.go
+++ b/modules/text2vec-openai/clients/openai.go
@@ -24,6 +24,7 @@ import (
 	"time"
 
 	"github.com/weaviate/weaviate/entities/moduletools"
+	"github.com/weaviate/weaviate/usecases/logrusext"
 
 	"github.com/weaviate/weaviate/usecases/modulecomponents"
 
@@ -108,6 +109,7 @@ type client struct {
 	httpClient         *http.Client
 	buildUrlFn         func(baseURL, resourceName, deploymentID, apiVersion string, isAzure bool) (string, error)
 	logger             logrus.FieldLogger
+	sampledLogger      *logrusext.Sampler
 }
 
 func New(openAIApiKey, openAIOrganization, azureApiKey string, timeout time.Duration, logger logrus.FieldLogger) *client {
@@ -118,8 +120,9 @@ func New(openAIApiKey, openAIOrganization, azureApiKey string, timeout time.Dura
 		httpClient: &http.Client{
 			Timeout: timeout,
 		},
-		buildUrlFn: buildUrl,
-		logger:     logger,
+		buildUrlFn:    buildUrl,
+		logger:        logger,
+		sampledLogger: logrusext.NewSampler(logger, 5, time.Minute),
 	}
 }
 
@@ -184,7 +187,7 @@ func (v *client) vectorize(ctx context.Context, input []string, model string, co
 	if res.StatusCode != 200 || resBody.Error != nil {
 		return nil, nil, 0, v.getError(res.StatusCode, requestID, resBody.Error, config.IsAzure)
 	}
-	rateLimit := ent.GetRateLimitsFromHeader(res.Header)
+	rateLimit := ent.GetRateLimitsFromHeader(v.sampledLogger, res.Header, config.IsAzure)
 
 	texts := make([]string, len(resBody.Data))
 	embeddings := make([][]float32, len(resBody.Data))
diff --git a/modules/text2vec-openai/ent/vectorization_result.go b/modules/text2vec-openai/ent/vectorization_result.go
index 945119f445b..0ceedf4cb9b 100644
--- a/modules/text2vec-openai/ent/vectorization_result.go
+++ b/modules/text2vec-openai/ent/vectorization_result.go
@@ -16,12 +16,14 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/sirupsen/logrus"
+	"github.com/weaviate/weaviate/usecases/logrusext"
 	"github.com/weaviate/weaviate/usecases/modulecomponents"
 )
 
 const dummyLimit = 10000000
 
-func GetRateLimitsFromHeader(header http.Header) *modulecomponents.RateLimits {
+func GetRateLimitsFromHeader(l *logrusext.Sampler, header http.Header, isAzure bool) *modulecomponents.RateLimits {
 	requestsReset, err := time.ParseDuration(header.Get("x-ratelimit-reset-requests"))
 	if err != nil {
 		requestsReset = 0
@@ -37,30 +39,41 @@ func GetRateLimitsFromHeader(header http.Header) *modulecomponents.RateLimits {
 	remainingTokens := getHeaderInt(header, "x-ratelimit-remaining-tokens")
 
 	// azure returns 0 as limit, make sure this does not block anything by setting a high value
-	if limitTokens == 0 {
+	if isAzure {
 		limitTokens = dummyLimit
-	}
-	if limitRequests == 0 {
 		limitRequests = dummyLimit
 	}
+
+	updateWithMissingValues := false
+	if limitRequests < 0 || limitTokens < 0 || remainingRequests < 0 || remainingTokens < 0 {
+		updateWithMissingValues = true
+
+		// logging all headers as there should not be anything sensitive according to the documentation:
+		// https://platform.openai.com/docs/api-reference/debugging-requests
+		l.WithSampling(func(l logrus.FieldLogger) {
+			l.WithField("headers", header).
+				Warn("rate limit headers are missing or invalid, going to keep using the old values")
+		})
+	}
 	return &modulecomponents.RateLimits{
-		LimitRequests:     limitRequests,
-		LimitTokens:       limitTokens,
-		RemainingRequests: remainingRequests,
-		RemainingTokens:   remainingTokens,
-		ResetRequests:     time.Now().Add(requestsReset),
-		ResetTokens:       time.Now().Add(tokensReset),
+		LimitRequests:           limitRequests,
+		LimitTokens:             limitTokens,
+		RemainingRequests:       remainingRequests,
+		RemainingTokens:         remainingTokens,
+		ResetRequests:           time.Now().Add(requestsReset),
+		ResetTokens:             time.Now().Add(tokensReset),
+		UpdateWithMissingValues: updateWithMissingValues,
 	}
 }
 
 func getHeaderInt(header http.Header, key string) int {
 	value := header.Get(key)
 	if value == "" {
-		return 0
+		return -1
 	}
 	i, err := strconv.Atoi(value)
 	if err != nil {
-		return 0
+		return -1
 	}
 	return i
 }
diff --git a/usecases/logrusext/sampling.go b/usecases/logrusext/sampling.go
new file mode 100644
index 00000000000..2486c2e7d96
--- /dev/null
+++ b/usecases/logrusext/sampling.go
@@ -0,0 +1,57 @@
+//                           _       _
+// __      _____  __ ___   ___  __ _| |_ ___
+// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
+//  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
+//   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
+//
+//  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
+//
+//  CONTACT: hello@weaviate.io
+//
+
+package logrusext
+
+import (
+	"sync"
+	"time"
+
+	"github.com/sirupsen/logrus"
+)
+
+// Sampler is a very simple log sampler which prints up to N log messages
+// every tick duration.
+type Sampler struct {
+	l logrus.FieldLogger
+
+	mu        sync.Mutex
+	counter   int
+	limit     int
+	tick      time.Duration
+	lastReset time.Time
+}
+
+func NewSampler(l logrus.FieldLogger, n int, tick time.Duration) *Sampler {
+	return &Sampler{
+		l:     l,
+		limit: n,
+		tick:  tick,
+	}
+}
+
+func (s *Sampler) WithSampling(fn func(l logrus.FieldLogger)) {
+	now := time.Now()
+
+	s.mu.Lock()
+	counter := s.counter
+	if now.Sub(s.lastReset) > s.tick {
+		counter = 0
+		s.lastReset = now
+	}
+	counter++
+	s.counter = counter
+	s.mu.Unlock()
+
+	if counter <= s.limit {
+		fn(s.l)
+	}
+}
diff --git a/usecases/logrusext/sampling_test.go b/usecases/logrusext/sampling_test.go
new file mode 100644
index 00000000000..64be549db78
--- /dev/null
+++ b/usecases/logrusext/sampling_test.go
@@ -0,0 +1,44 @@
+//                           _       _
+// __      _____  __ ___   ___  __ _| |_ ___
+// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
+//  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
+//   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
+//
+//  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
+//
+//  CONTACT: hello@weaviate.io
+//
+
+package logrusext
+
+import (
+	"testing"
+	"time"
+
+	"github.com/sirupsen/logrus"
+	"github.com/sirupsen/logrus/hooks/test"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSampler_Simple(t *testing.T) {
+	logger, hook := test.NewNullLogger()
+	sampler := NewSampler(logger, 3, 100*time.Millisecond)
+
+	for range 10 {
+		sampler.WithSampling(func(l logrus.FieldLogger) {
+			l.Infof("hello")
+		})
+	}
+
+	require.Len(t, hook.AllEntries(), 3)
+	hook.Reset()
+
+	time.Sleep(200 * time.Millisecond)
+
+	for range 5 {
+		sampler.WithSampling(func(l logrus.FieldLogger) {
+			l.Infof("hello")
+		})
+	}
+	require.Len(t, hook.AllEntries(), 3)
+}
diff --git a/usecases/modulecomponents/client_results.go b/usecases/modulecomponents/client_results.go
index 7862784caa9..d31897bb711 100644
--- a/usecases/modulecomponents/client_results.go
+++ b/usecases/modulecomponents/client_results.go
@@ -18,17 +18,18 @@ import (
 )
 
 type RateLimits struct {
-	LastOverwrite        time.Time
-	AfterRequestFunction func(limits *RateLimits, tokensUsed int, deductRequest bool)
-	LimitRequests        int
-	LimitTokens          int
-	RemainingRequests    int
-	RemainingTokens      int
-	ReservedRequests     int
-	ReservedTokens       int
-	ResetRequests        time.Time
-	ResetTokens          time.Time
-	Label                string
+	LastOverwrite           time.Time
+	AfterRequestFunction    func(limits *RateLimits, tokensUsed int, deductRequest bool)
+	LimitRequests           int
+	LimitTokens             int
+	RemainingRequests       int
+	RemainingTokens         int
+	ReservedRequests        int
+	ReservedTokens          int
+	ResetRequests           time.Time
+	ResetTokens             time.Time
+	Label                   string
+	UpdateWithMissingValues bool
 }
 
 func (rl *RateLimits) ResetAfterRequestFunction(tokensUsed int) {
@@ -81,6 +82,9 @@ func (rl *RateLimits) CanSendFullBatch(numRequests int, batchTokens int, addMetr
 }
 
 func (rl *RateLimits) UpdateWithRateLimit(other *RateLimits) {
+	if other.UpdateWithMissingValues {
+		return
+	}
 	rl.LimitRequests = other.LimitRequests
 	rl.LimitTokens = other.LimitTokens
 	rl.ResetRequests = other.ResetRequests


Reviews:
user: login: orca-security-eu[bot]
body: ### Orca Security Scan Summary
| Status  | Check | Issues by priority |   |
| ------- | ----- | ------------------ | - |
| <img width="16" alt="Passed" src="https://raw.githubusercontent.com/orcasecurity/orca-cli/main/resources/images/prcomment/status/passed.png" title="Passed"> Passed | Secrets | <img width="12" alt="high" src="https://raw.githubusercontent.com/orcasecurity/orca-cli/main/resources/images/prcomment/priority/high.png" title="High"> 0 &emsp; <img width="12" alt="medium" src="https://raw.githubusercontent.com/orcasecurity/orca-cli/main/resources/images/prcomment/priority/medium.png" title="Medium"> 0 &emsp; <img width="12" alt="low" src="https://raw.githubusercontent.com/orcasecurity/orca-cli/main/resources/images/prcomment/priority/low.png" title="Low"> 0 &emsp; <img width="12" alt="info" src="https://raw.githubusercontent.com/orcasecurity/orca-cli/main/resources/images/prcomment/priority/info.png" title="Info"> 0 | <a href="https://eu.app.orcasecurity.io/shift-left/file_system_secret_detection/scan-log/97ba313d-c95f-4d94-ba40-0a746f626bd2" target="_blank">View in Orca</a> |
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-18T08:37:13Z
user: login: faustuzas
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T08:42:06Z
user: login: dirkkul
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T10:00:58Z
user: login: faustuzas
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T10:05:50Z

Review comments:
diff_hunk: @@ -37,30 +37,34 @@ func GetRateLimitsFromHeader(header http.Header) *modulecomponents.RateLimits {
 	remainingTokens := getHeaderInt(header, "x-ratelimit-remaining-tokens")
 
 	// azure returns 0 as limit, make sure this does not block anything by setting a high value
-	if limitTokens == 0 {
+	if isAzure {
 		limitTokens = dummyLimit
-	}
-	if limitRequests == 0 {
 		limitRequests = dummyLimit
 	}
+
+	updateWithMissingValues := false
+	if limitRequests < 0 || limitTokens < 0 || remainingRequests < 0 || remainingTokens < 0 {
+		updateWithMissingValues = true
path: modules/text2vec-openai/ent/vectorization_result.go
user: login: faustuzas
body: it would be very interesting to log all headers here (omitting the sensitive ones of course). We might discover that we are parsing stuff incorrectly or there is some additional "error" metric or something.
created_at: 2025-03-18T08:42:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -45,6 +47,13 @@ func GetRateLimitsFromHeader(header http.Header, isAzure bool) *modulecomponents
 	updateWithMissingValues := false
 	if limitRequests < 0 || limitTokens < 0 || remainingRequests < 0 || remainingTokens < 0 {
 		updateWithMissingValues = true
+
+		// logging all headers as there should not be anything sensitive according to the documentation:
+		// https://platform.openai.com/docs/api-reference/debugging-requests
+		l.WithSampling(func(l logrus.FieldLogger) {
+			l.WithField("headers", header).
+				Warn("rate limit headers are missing or invalid, using dummy")
path: modules/text2vec-openai/ent/vectorization_result.go
user: login: dirkkul
body: we are not using a dummy value, but will skip updating the the value in the batching algo
created_at: 2025-03-18T10:00:58Z
author_association: CONTRIBUTOR

PR comments:
user: login: sonarqubecloud[bot]
created_at: 2025-03-18T10:04:47Z
body: ## [![Quality Gate Passed](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/checks/QualityGateBadge/qg-passed-20px.png 'Quality Gate Passed')](https://sonarcloud.io/dashboard?id=weaviate_weaviate&pullRequest=7562) **Quality Gate passed**  
Issues  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [1 New issue](https://sonarcloud.io/project/issues?id=weaviate_weaviate&pullRequest=7562&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/accepted-16px.png '') [0 Accepted issues](https://sonarcloud.io/project/issues?id=weaviate_weaviate&pullRequest=7562&issueStatuses=ACCEPTED)

Measures  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0 Security Hotspots](https://sonarcloud.io/project/security_hotspots?id=weaviate_weaviate&pullRequest=7562&issueStatuses=OPEN,CONFIRMED&sinceLeakPeriod=true)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Coverage on New Code](https://sonarcloud.io/component_measures?id=weaviate_weaviate&pullRequest=7562&metric=new_coverage&view=list)  
![](https://sonarsource.github.io/sonarcloud-github-static-resources/v2/common/passed-16px.png '') [0.0% Duplication on New Code](https://sonarcloud.io/component_measures?id=weaviate_weaviate&pullRequest=7562&metric=new_duplicated_lines_density&view=list)  
  
[See analysis details on SonarQube Cloud](https://sonarcloud.io/dashboard?id=weaviate_weaviate&pullRequest=7562)


author_association: NONE