#43 Defined a LangModel struct and merged all health tracking code there

pkg/config/fields/error_budget.go

@@ -0,0 +1 @@
+package fields

pkg/providers/config.go

@@ -31,14 +31,14 @@ func DefaultLangModelConfig() *LangModelConfig {
 	}
 }
 
-func (c *LangModelConfig) ToModel(tel *telemetry.Telemetry) (LanguageModel, error) {
+func (c *LangModelConfig) ToModel(tel *telemetry.Telemetry) (*LangModel, error) {
 	if c.OpenAI != nil {
 		client, err := openai.NewClient(c.OpenAI, tel)
 		if err != nil {
 			return nil, fmt.Errorf("error initing openai client: %v", err)
 		}
 
-		return client, nil
+		return NewLangModel(c.ID, client), nil
 	}
 
 	return nil, ErrProviderNotFound

pkg/providers/provider.go

@@ -3,16 +3,55 @@ package providers
 import (
 	"context"
 
+	"glide/pkg/routers/health"
+
 	"glide/pkg/api/schemas"
 )
 
-// ModelProvider defines an interface all model providers should support
-type ModelProvider interface {
+// LangModelProvider defines an interface a provider should fulfill to be able to serve language chat requests
+type LangModelProvider interface {
 	Provider() string
+	Chat(ctx context.Context, request *schemas.UnifiedChatRequest) (*schemas.UnifiedChatResponse, error)
 }
 
-// LanguageModel defines the interface a provider should fulfill to be able to serve language chat requests
-type LanguageModel interface {
-	ID() string
-	Chat(ctx context.Context, request *schemas.UnifiedChatRequest) (*schemas.UnifiedChatResponse, error)
+// LangModel
+type LangModel struct {
+	modelID     string
+	client      LangModelProvider
+	rateLimit   *health.RateLimitTracker
+	errorBudget *health.TokenBucket // TODO: centralize provider API health tracking in the registry
+}
+
+func NewLangModel(modelID string, client LangModelProvider) *LangModel {
+	return &LangModel{
+		modelID:     modelID,
+		client:      client,
+		rateLimit:   health.NewRateLimitTracker(),
+		errorBudget: health.NewTokenBucket(1, 10), // TODO: set from configs
+	}
+}
+
+func (m *LangModel) ID() string {
+	return m.modelID
+}
+
+func (m *LangModel) Provider() string {
+	return m.client.Provider()
+}
+
+func (m *LangModel) Healthy() bool {
+	return !m.rateLimit.Limited() && m.errorBudget.HasTokens()
+}
+
+func (m *LangModel) Chat(ctx context.Context, request *schemas.UnifiedChatRequest) (*schemas.UnifiedChatResponse, error) {
+	resp, err := m.client.Chat(ctx, request)
+
+	if err == nil {
+		// successful response
+		return resp, err
+	}
+
+	// TODO: track all availability issues
+
+	return resp, err
 }

pkg/routers/config.go

@@ -49,10 +49,10 @@ type LangRouterConfig struct {
 }
 
 // BuildModels creates LanguageModel slice out of the given config
-func (c *LangRouterConfig) BuildModels(tel *telemetry.Telemetry) ([]providers.LanguageModel, error) {
+func (c *LangRouterConfig) BuildModels(tel *telemetry.Telemetry) ([]*providers.LangModel, error) {
 	var errs error
 
-	models := make([]providers.LanguageModel, 0, len(c.Models))
+	models := make([]*providers.LangModel, 0, len(c.Models))
 
 	for _, modelConfig := range c.Models {
 		if !modelConfig.Enabled {

pkg/routers/router.go

@@ -4,8 +4,9 @@ import (
 	"context"
 	"errors"
 
+	"glide/pkg/providers"
+
 	"glide/pkg/api/schemas"
-	"glide/pkg/routers/health"
 	"glide/pkg/routers/routing"
 	"glide/pkg/telemetry"
 )
@@ -15,7 +16,7 @@ var ErrNoModels = errors.New("no models configured for router")
 type LangRouter struct {
 	Config    *LangRouterConfig
 	routing   routing.LangModelRouting
-	models    *[]health.LangModelHealthTracker
+	models    []*providers.LangModel
 	telemetry *telemetry.Telemetry
 }
 
@@ -25,16 +26,10 @@ func NewLangRouter(cfg *LangRouterConfig, tel *telemetry.Telemetry) (*LangRouter
 		return nil, err
 	}
 
-	modelTrackers := make([]health.LangModelHealthTracker, 0, len(models))
-
-	for _, model := range models {
-		modelTrackers = append(modelTrackers, *health.NewLangModelHealthTracker(model))
-	}
-
 	router := &LangRouter{
 		Config:    cfg,
-		models:    &modelTrackers,
-		routing:   routing.NewPriorityRouting(&modelTrackers),
+		models:    models,
+		routing:   routing.NewPriorityRouting(models),
 		telemetry: tel,
 	}
 
@@ -46,7 +41,7 @@ func (r *LangRouter) ID() string {
 }
 
 func (r *LangRouter) Chat(ctx context.Context, request *schemas.UnifiedChatRequest) (*schemas.UnifiedChatResponse, error) {
-	if len(*r.models) == 0 {
+	if len(r.models) == 0 {
 		return nil, ErrNoModels
 	}
 

pkg/routers/routing/priority.go

@@ -3,7 +3,7 @@ package routing
 import (
 	"sync/atomic"
 
-	"glide/pkg/routers/health"
+	"glide/pkg/providers"
 )
 
 const (
@@ -15,10 +15,10 @@ const (
 //	Priority of models are defined as position of the model on the list
 //	(e.g. the first model definition has the highest priority, then the second model definition and so on)
 type PriorityRouting struct {
-	models *[]health.LangModelHealthTracker
+	models []*providers.LangModel
 }
 
-func NewPriorityRouting(models *[]health.LangModelHealthTracker) *PriorityRouting {
+func NewPriorityRouting(models []*providers.LangModel) *PriorityRouting {
 	return &PriorityRouting{
 		models: models,
 	}
@@ -35,11 +35,11 @@ func (r *PriorityRouting) Iterator() LangModelIterator {
 
 type PriorityIterator struct {
 	idx    *atomic.Uint64
-	models *[]health.LangModelHealthTracker
+	models []*providers.LangModel
 }
 
-func (r PriorityIterator) Next() (*health.LangModelHealthTracker, error) {
-	models := *r.models
+func (r PriorityIterator) Next() (*providers.LangModel, error) {
+	models := r.models
 	idx := r.idx.Load()
 
 	for int(idx) < len(models) {
@@ -48,7 +48,7 @@ func (r PriorityIterator) Next() (*health.LangModelHealthTracker, error) {
 		r.idx.Add(1)
 
 		if model.Healthy() {
-			return &model, nil
+			return model, nil
 		}
 
 		// otherwise, try to pick the next model on the list

pkg/routers/routing/strategies.go

@@ -3,7 +3,7 @@ package routing
 import (
 	"errors"
 
-	"glide/pkg/routers/health"
+	"glide/pkg/providers"
 )
 
 var ErrNoHealthyModels = errors.New("no healthy models found")
@@ -16,5 +16,5 @@ type LangModelRouting interface {
 }
 
 type LangModelIterator interface {
-	Next() (*health.LangModelHealthTracker, error)
+	Next() (*providers.LangModel, error)
 }