ModelRouter/internal/scheduler/scheduler.go

// Package scheduler implements request scheduling across providers: per-source
// concurrency caps (acquire with wait = queuing), AUTO model fallback chains,
// and exponential backoff via provider health.
package scheduler

import (
	"context"
	"fmt"

	"llmsproxy/internal/provider"
	"llmsproxy/internal/types"
)

// Scheduler drives one chat tool call across the candidate provider chain.
type Scheduler struct {
	// MaxRetries how many fallback providers to try before failing.
	MaxRetries int
}

func New(maxRetries int) *Scheduler {
	if maxRetries < 0 {
		maxRetries = 0
	}
	return &Scheduler{MaxRetries: maxRetries}
}

// Provider is the minimal interface the scheduler needs to schedule over.
type Provider interface {
	Name() string
	Available() bool
	ModelFor(reqModel string) string
	Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error)
	ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error)
	Image(ctx context.Context, req *types.ImageGenRequest) (*types.UnifiedResponse, error)
}

// FromRegistry converts *provider.Provider slices to the scheduler interface.
func FromRegistry(ps []*provider.Provider) []Provider {
	out := make([]Provider, len(ps))
	for i, p := range ps {
		out[i] = p
	}
	return out
}

// Chat runs a chat request across cands, falling back on failure. Each
// candidate receives a request pinned to its own model (ModelFor), so an AUTO
// chain fallback switches the model id per provider instead of reusing the
// first candidate's model name.
//
// On success it returns the response together with the name of the provider
// and the exact model id that actually served the request (used for stats).
func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatRequest) (*types.UnifiedResponse, string, string, error) {
	attempts := s.MaxRetries + 1
	var lastErr error
	for i := 0; i < attempts && i < len(cands); i++ {
		p := cands[i]
		r := *req
		r.Model = p.ModelFor(req.Model)
		resp, err := p.Chat(ctx, &r)
		if ctx.Err() != nil {
			return nil, "", "", ctx.Err()
		}
		if err == nil {
			return resp, p.Name(), r.Model, nil
		}
		lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
	}
	if lastErr == nil {
		if len(cands) == 0 {
			lastErr = fmt.Errorf("no provider available")
		}
	}
	return nil, "", "", lastErr
}

// ChatStream runs a streaming chat across cands, falling back early on connect
// errors. The request model is pinned per candidate like Chat. On success it
// returns the chunk channel plus the serving provider name and model id.
func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types.ChatRequest) (<-chan types.UnifiedChunk, string, string, error) {
	attempts := s.MaxRetries + 1
	var lastErr error
	for i := 0; i < attempts && i < len(cands); i++ {
		p := cands[i]
		r := *req
		r.Model = p.ModelFor(req.Model)
		resp, err := p.ChatStream(ctx, &r)
		if err == nil {
			return resp, p.Name(), r.Model, nil
		}
		lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
	}
	if lastErr == nil && len(cands) == 0 {
		lastErr = fmt.Errorf("no provider available")
	}
	return nil, "", "", lastErr
}

// Image runs an image-generation request across cands; returns the used
// provider name on success.
func (s *Scheduler) Image(ctx context.Context, cands []Provider, req *types.ImageGenRequest) (*types.UnifiedResponse, string, error) {
	attempts := s.MaxRetries + 1
	var lastErr error
	for i := 0; i < attempts && i < len(cands); i++ {
		p := cands[i]
		resp, err := p.Image(ctx, req)
		if err == nil {
			return resp, p.Name(), nil
		}
		lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
	}
	if lastErr == nil && len(cands) == 0 {
		lastErr = fmt.Errorf("no provider available")
	}
	return nil, "", lastErr
}