newkirk
91 lines
2.5 KiB
Go
91 lines
2.5 KiB
Go
package providers
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/sony/gobreaker"
|
|
"gophergate/internal/models"
|
|
)
|
|
|
|
type CircuitBreakerProvider struct {
|
|
provider Provider
|
|
cb *gobreaker.CircuitBreaker
|
|
}
|
|
|
|
func NewCircuitBreakerProvider(p Provider) Provider {
|
|
name := p.Name()
|
|
var maxRequests uint32 = 5
|
|
var interval = 60 * time.Second
|
|
var timeout = 5 * time.Minute
|
|
|
|
settings := gobreaker.Settings{
|
|
Name: name,
|
|
MaxRequests: maxRequests,
|
|
Interval: interval,
|
|
Timeout: timeout,
|
|
ReadyToTrip: func(counts gobreaker.Counts) bool {
|
|
// Trip after 3 consecutive failures
|
|
return counts.ConsecutiveFailures > 3
|
|
},
|
|
}
|
|
return &CircuitBreakerProvider{
|
|
provider: p,
|
|
cb: gobreaker.NewCircuitBreaker(settings),
|
|
}
|
|
}
|
|
|
|
func (cbp *CircuitBreakerProvider) Name() string {
|
|
return cbp.provider.Name()
|
|
}
|
|
|
|
func (cbp *CircuitBreakerProvider) ChatCompletion(ctx context.Context, req *models.UnifiedRequest) (*models.ChatCompletionResponse, error) {
|
|
result, err := cbp.cb.Execute(func() (interface{}, error) {
|
|
return cbp.provider.ChatCompletion(ctx, req)
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return result.(*models.ChatCompletionResponse), nil
|
|
}
|
|
|
|
func (cbp *CircuitBreakerProvider) ChatCompletionStream(ctx context.Context, req *models.UnifiedRequest) (<-chan *models.ChatCompletionStreamResponse, error) {
|
|
result, err := cbp.cb.Execute(func() (interface{}, error) {
|
|
return cbp.provider.ChatCompletionStream(ctx, req)
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return result.(<-chan *models.ChatCompletionStreamResponse), nil
|
|
}
|
|
|
|
func (cbp *CircuitBreakerProvider) ImageGeneration(ctx context.Context, req *models.ImageGenerationRequest) (*models.ImageGenerationResponse, error) {
|
|
result, err := cbp.cb.Execute(func() (interface{}, error) {
|
|
return cbp.provider.ImageGeneration(ctx, req)
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return result.(*models.ImageGenerationResponse), nil
|
|
}
|
|
|
|
func (cbp *CircuitBreakerProvider) Responses(ctx context.Context, req *models.ResponsesRequest) (*models.ResponsesResponse, error) {
|
|
result, err := cbp.cb.Execute(func() (interface{}, error) {
|
|
return cbp.provider.Responses(ctx, req)
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return result.(*models.ResponsesResponse), nil
|
|
}
|
|
|
|
func (cbp *CircuitBreakerProvider) ResponsesStream(ctx context.Context, req *models.ResponsesRequest) (<-chan *models.ResponsesStreamChunk, error) {
|
|
result, err := cbp.cb.Execute(func() (interface{}, error) {
|
|
return cbp.provider.ResponsesStream(ctx, req)
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return result.(<-chan *models.ResponsesStreamChunk), nil
|
|
}
|