diff --git a/pkg/providers/lang.go b/pkg/providers/lang.go index d3643432..3db4524d 100644 --- a/pkg/providers/lang.go +++ b/pkg/providers/lang.go @@ -2,9 +2,10 @@ package providers import ( "context" - "glide/pkg/routers/health" "time" + "glide/pkg/routers/health" + "glide/pkg/api/schemas" "glide/pkg/providers/clients" "glide/pkg/routers/latency" @@ -33,7 +34,7 @@ type LanguageModel struct { modelID string weight int client LangProvider - healthTracker *health.HealthTracker + healthTracker *health.Tracker chatLatency *latency.MovingAverage chatStreamLatency *latency.MovingAverage latencyUpdateInterval *time.Duration @@ -43,7 +44,7 @@ func NewLangModel(modelID string, client LangProvider, budget health.ErrorBudget return &LanguageModel{ modelID: modelID, client: client, - healthTracker: health.NewHealthTracker(budget), + healthTracker: health.NewTracker(budget), chatLatency: latency.NewMovingAverage(latencyConfig.Decay, latencyConfig.WarmupSamples), chatStreamLatency: latency.NewMovingAverage(latencyConfig.Decay, latencyConfig.WarmupSamples), latencyUpdateInterval: latencyConfig.UpdateInterval, diff --git a/pkg/providers/testing.go b/pkg/providers/testing.go index 8762aee8..2893108e 100644 --- a/pkg/providers/testing.go +++ b/pkg/providers/testing.go @@ -104,14 +104,14 @@ type LangModelMock struct { weight int } -func NewLangModelMock(ID string, healthy bool, avgLatency float64, weight int) *LangModelMock { +func NewLangModelMock(ID string, healthy bool, avgLatency float64, weight int) LangModelMock { chatLatency := latency.NewMovingAverage(0.06, 3) if avgLatency > 0.0 { chatLatency.Set(avgLatency) } - return &LangModelMock{ + return LangModelMock{ modelID: ID, healthy: healthy, chatLatency: chatLatency, diff --git a/pkg/routers/health/tracker.go b/pkg/routers/health/tracker.go index 92da9153..1be6832f 100644 --- a/pkg/routers/health/tracker.go +++ b/pkg/routers/health/tracker.go @@ -2,26 +2,28 @@ package health import ( "errors" + "glide/pkg/providers/clients" ) -type HealthTracker struct { +// Tracker tracks errors and general health of model provider +type Tracker struct { errBudget *TokenBucket rateLimit *RateLimitTracker } -func NewHealthTracker(budget ErrorBudget) *HealthTracker { - return &HealthTracker{ +func NewTracker(budget ErrorBudget) *Tracker { + return &Tracker{ rateLimit: NewRateLimitTracker(), errBudget: NewTokenBucket(budget.TimePerTokenMicro(), budget.Budget()), } } -func (t *HealthTracker) Healthy() bool { +func (t *Tracker) Healthy() bool { return !t.rateLimit.Limited() && t.errBudget.HasTokens() } -func (t *HealthTracker) TrackErr(err error) { +func (t *Tracker) TrackErr(err error) { var rateLimitErr *clients.RateLimitError if errors.As(err, &rateLimitErr) { diff --git a/pkg/routers/routing/least_latency_test.go b/pkg/routers/routing/least_latency_test.go index b5ddb570..70bd5d4e 100644 --- a/pkg/routers/routing/least_latency_test.go +++ b/pkg/routers/routing/least_latency_test.go @@ -115,7 +115,8 @@ func TestLeastLatencyRouting_Routing(t *testing.T) { } routing := LeastLatencyRouting{ - schedules: schedules, + latencyGetter: providers.ChatMockLatency, + schedules: schedules, } iterator := routing.Iterator()