68 lines
2.0 KiB
Go
68 lines
2.0 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
type codexQuotaOverdraftAccountTestCoordinator interface {
|
|
ObserveAccount(account *Account, preferredModel string)
|
|
ObserveBusinessSuccess(account *Account, preferredModel string)
|
|
HandleQuota429(ctx context.Context, account *Account, headers http.Header, body []byte, preferredModel string) bool
|
|
}
|
|
|
|
func (s *AccountTestService) SetCodexQuotaOverdraftCoordinator(coordinator *CodexQuotaOverdraftCoordinator) {
|
|
if s != nil {
|
|
s.codexQuotaOverdraft = coordinator
|
|
}
|
|
}
|
|
|
|
func (s *AccountTestService) prepareCodexQuotaOverdraftTestRequest(
|
|
ctx context.Context,
|
|
account *Account,
|
|
payload []byte,
|
|
) (context.Context, []byte, bool) {
|
|
if !s.codexQuotaOverdraftTestEnabled(account) || !codexQuotaOverdraftInjectionEligible(account, time.Now().UTC()) {
|
|
return ctx, payload, false
|
|
}
|
|
ctx = WithCodexQuotaOverdraftScheduling(ctx)
|
|
updated, changed, err := injectCodexQuotaOverdraft(payload)
|
|
if err != nil || !changed {
|
|
return ctx, payload, false
|
|
}
|
|
markCodexQuotaOverdraftInjected(ctx, account.ID)
|
|
return ctx, updated, true
|
|
}
|
|
|
|
func (s *AccountTestService) handleCodexQuotaOverdraftTest429(
|
|
ctx context.Context,
|
|
account *Account,
|
|
headers http.Header,
|
|
body []byte,
|
|
preferredModel string,
|
|
) bool {
|
|
return s.codexQuotaOverdraft != nil && s.codexQuotaOverdraftTestEnabled(account) &&
|
|
s.codexQuotaOverdraft.HandleQuota429(ctx, account, headers, body, preferredModel)
|
|
}
|
|
|
|
func (s *AccountTestService) observeCodexQuotaOverdraftTestResult(
|
|
account *Account,
|
|
preferredModel string,
|
|
injected bool,
|
|
) {
|
|
if s.codexQuotaOverdraft == nil || !s.codexQuotaOverdraftTestEnabled(account) {
|
|
return
|
|
}
|
|
if injected {
|
|
s.codexQuotaOverdraft.ObserveBusinessSuccess(account, preferredModel)
|
|
return
|
|
}
|
|
s.codexQuotaOverdraft.ObserveAccount(account, preferredModel)
|
|
}
|
|
|
|
func (s *AccountTestService) codexQuotaOverdraftTestEnabled(account *Account) bool {
|
|
return s != nil && s.cfg != nil && s.cfg.Gateway.CodexQuotaOverdraftEnabled &&
|
|
isCodexQuotaOverdraftAccount(account)
|
|
}
|