Agent-Native/runtime/llm/fallback.go

package llm

import "context"

// WithFallback 在尚未转发任何增量时，遇到 5xx 就改用链上下一个模型。
// fallbacks 为空时返回原来的 StreamFunc，不改变行为。
func WithFallback(base StreamFunc, fallbacks []string, onSwitch func(from, to string)) StreamFunc {
	chain := nonEmpty(fallbacks)
	if base == nil || len(chain) == 0 {
		return base
	}
	return func(ctx context.Context, model string, messages []Message, tools []ToolDecl, opts StreamOptions) <-chan Event {
		out := make(chan Event)
		go func() {
			defer close(out)
			models := append([]string{model}, chain...)
			for i, current := range models {
				forwarded := false
				switchNext := false
				for ev := range base(ctx, current, messages, tools, opts) {
					if isDelta(ev) {
						forwarded = true
					}
					if ev.Kind == KindError && !forwarded && serverError(ev) && i+1 < len(models) {
						switchNext = true
						if onSwitch != nil {
							onSwitch(current, models[i+1])
						}
						break
					}
					select {
					case out <- ev:
					case <-ctx.Done():
						return
					}
					if ev.Kind == KindMessageEnd || ev.Kind == KindError {
						return
					}
				}
				if !switchNext {
					return
				}
			}
		}()
		return out
	}
}

func nonEmpty(in []string) []string {
	var out []string
	seen := map[string]bool{}
	for _, s := range in {
		s = trimModel(s)
		if s == "" || seen[s] {
			continue
		}
		seen[s] = true
		out = append(out, s)
	}
	return out
}

func trimModel(s string) string {
	for len(s) > 0 && (s[0] == ' ' || s[0] == '\t') {
		s = s[1:]
	}
	for len(s) > 0 && (s[len(s)-1] == ' ' || s[len(s)-1] == '\t') {
		s = s[:len(s)-1]
	}
	return s
}

func isDelta(ev Event) bool {
	return ev.Kind == KindTextDelta || ev.Kind == KindReasoningDelta || ev.Kind == KindToolCallDelta
}

func serverError(ev Event) bool {
	return ev.HTTPStatus >= 500 && ev.HTTPStatus <= 599
}