AIRouter/internal/router/overload.go

package router

import (
	"fmt"
	"strings"

	"github.com/edgeai/gateway/internal/config"
)

// OverloadStrategy defines how to handle requests when the primary model is overloaded.
type OverloadStrategy int

const (
	StrategySameModelOtherInstance OverloadStrategy = iota // try same model on another instance
	StrategySmallerLocalModel                              // fall back to a smaller local model
	StrategyBackupEdgeNode                                 // route to a backup edge node
	StrategyReject                                         // reject the request
)

// OverloadResolver resolves overload situations using configured strategies.
type OverloadResolver struct {
	strategies []OverloadStrategy
	modelMap   *LogicalModelMapping
}

// NewOverloadResolver creates an overload resolver from routing config.
func NewOverloadResolver(cfg *config.RoutingConfig, modelMap *LogicalModelMapping) *OverloadResolver {
	r := &OverloadResolver{
		modelMap: modelMap,
	}
	for _, s := range cfg.OverloadStrategy {
		switch strings.ToLower(s) {
		case "same_model_other_instance":
			r.strategies = append(r.strategies, StrategySameModelOtherInstance)
		case "smaller_local_model":
			r.strategies = append(r.strategies, StrategySmallerLocalModel)
		case "backup_edge_node":
			r.strategies = append(r.strategies, StrategyBackupEdgeNode)
		case "reject":
			r.strategies = append(r.strategies, StrategyReject)
		}
	}
	// Default: reject if no strategies configured
	if len(r.strategies) == 0 {
		r.strategies = []OverloadStrategy{StrategyReject}
	}
	return r
}

// OverloadResult contains the result of an overload resolution attempt.
type OverloadResult struct {
	Strategy  OverloadStrategy
	Target    *ModelTarget // resolved fallback target, nil if rejected
	Rejected  bool
	Reason    string
}

// Resolve attempts to find a fallback target when the primary model is overloaded.
// currentModel is the logical model that is overloaded.
// excludeEndpoints is a set of endpoints already tried (to avoid loops).
func (r *OverloadResolver) Resolve(currentModel string, excludeEndpoints map[string]bool) OverloadResult {
	for _, strategy := range r.strategies {
		switch strategy {
		case StrategySameModelOtherInstance:
			// Find same actual model on a different endpoint
			target := r.findSameModelDifferentEndpoint(currentModel, excludeEndpoints)
			if target != nil {
				return OverloadResult{
					Strategy: StrategySameModelOtherInstance,
					Target:   target,
				}
			}
		case StrategySmallerLocalModel:
			// Find a model with smaller context window (proxy for "smaller")
			target := r.findSmallerModel(currentModel, excludeEndpoints)
			if target != nil {
				return OverloadResult{
					Strategy: StrategySmallerLocalModel,
					Target:   target,
				}
			}
		case StrategyBackupEdgeNode:
			// Find any available model not yet tried
			target := r.findAnyAvailable(excludeEndpoints)
			if target != nil {
				return OverloadResult{
					Strategy: StrategyBackupEdgeNode,
					Target:   target,
				}
			}
		case StrategyReject:
			return OverloadResult{
				Strategy: StrategyReject,
				Rejected:  true,
				Reason:    "all overload strategies exhausted, request rejected",
			}
		}
	}
	return OverloadResult{
		Strategy: StrategyReject,
		Rejected:  true,
		Reason:    fmt.Sprintf("no fallback available for model %s", currentModel),
	}
}

func (r *OverloadResolver) findSameModelDifferentEndpoint(currentModel string, exclude map[string]bool) *ModelTarget {
	r.modelMap.mu.RLock()
	defer r.modelMap.mu.RUnlock()

	current, ok := r.modelMap.mapping[currentModel]
	if !ok {
		return nil
	}

	for logical, target := range r.modelMap.mapping {
		if logical == currentModel {
			continue
		}
		if target.ActualModel == current.ActualModel && !exclude[target.Endpoint] {
			return target
		}
	}
	return nil
}

func (r *OverloadResolver) findSmallerModel(currentModel string, exclude map[string]bool) *ModelTarget {
	r.modelMap.mu.RLock()
	defer r.modelMap.mu.RUnlock()

	current, ok := r.modelMap.mapping[currentModel]
	if !ok {
		return nil
	}

	var best *ModelTarget
	for logical, target := range r.modelMap.mapping {
		if logical == currentModel {
			continue
		}
		if exclude[target.Endpoint] {
			continue
		}
		// Pick model with smaller context window as "smaller" proxy
		if target.ContextWindow < current.ContextWindow {
			if best == nil || target.ContextWindow > best.ContextWindow {
				best = target
			}
		}
	}
	return best
}

func (r *OverloadResolver) findAnyAvailable(exclude map[string]bool) *ModelTarget {
	r.modelMap.mu.RLock()
	defer r.modelMap.mu.RUnlock()

	for _, target := range r.modelMap.mapping {
		if !exclude[target.Endpoint] {
			return target
		}
	}
	return nil
}

// StrategyName returns a human-readable name for a strategy.
func StrategyName(s OverloadStrategy) string {
	switch s {
	case StrategySameModelOtherInstance:
		return "same_model_other_instance"
	case StrategySmallerLocalModel:
		return "smaller_local_model"
	case StrategyBackupEdgeNode:
		return "backup_edge_node"
	case StrategyReject:
		return "reject"
	}
	return "unknown"
}