45 lines
1.2 KiB
Go
45 lines
1.2 KiB
Go
package providers
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/sony/gobreaker"
|
|
"gophergate/internal/models"
|
|
)
|
|
|
|
type CircuitBreakerProvider struct {
|
|
provider Provider
|
|
cb *gobreaker.CircuitBreaker
|
|
}
|
|
|
|
func NewCircuitBreakerProvider(p Provider) Provider {
|
|
settings := gobreaker.Settings{
|
|
Name: p.Name(),
|
|
}
|
|
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) {
|
|
// Circuit breaker for streaming is tricky. We'll just call the provider directly.
|
|
// Future: Implement a way to track stream failures in the circuit breaker.
|
|
return cbp.provider.ChatCompletionStream(ctx, req)
|
|
}
|