AutoZero/agent/model_routing_test.go

package agent

import (
	"context"
	"errors"
	"io"
	"strings"
	"testing"

	"git.gopher.ink/mirrors2/AutoZero/llm"
	"git.gopher.ink/mirrors2/AutoZero/session"
	"git.gopher.ink/mirrors2/AutoZero/tools"
)

// routingStub 可切换模型的最小 Provider：记录每次 Stream 看到的模型，
// 按当前模型决定成功或失败。
type routingStub struct {
	model    string
	failOn   map[string]bool // 当前模型命中即失败（返回错误，不流出文本）
	failOnce map[string]int  // 命中且剩余次数 >0 才失败（供"失败一次后成功"）
	seen     []string        // 每次 Stream 的当前模型
}

func (p *routingStub) Stream(ctx context.Context, system string, msgs []llm.Message,
	tools []llm.ToolDef, opts llm.StreamOptions, onDelta func(llm.StreamDelta)) (llm.Message, error) {
	p.seen = append(p.seen, p.model)
	if p.failOnce != nil && p.failOnce[p.model] > 0 {
		p.failOnce[p.model]--
		return llm.Message{}, errors.New("upstream 500")
	}
	if p.failOn != nil && p.failOn[p.model] {
		return llm.Message{}, errors.New("upstream 400")
	}
	return llm.Message{Role: llm.RoleAssistant, Content: "ok from " + p.model}, nil
}
func (p *routingStub) SetModel(model string) { p.model = model }
func (p *routingStub) CurrentModel() string  { return p.model }

func newRoutingAgent(t *testing.T, p llm.Provider, fallback, multimodal string) *Agent {
	t.Helper()
	a := &Agent{Provider: p, Tools: tools.NewRegistry(), Session: &session.Session{Path: t.TempDir() + "/s.jsonl"},
		FallbackModel: fallback, MultimodalModel: multimodal,
		Opt: Options{Stdout: io.Discard, Stderr: io.Discard}}
	return a
}

// TestFallbackOnStepError：首选调用失败 → 粘性回落备选 → 重试成功。
func TestFallbackOnStepError(t *testing.T) {
	p := &routingStub{model: "primary", failOn: map[string]bool{"primary": true}}
	a := newRoutingAgent(t, p, "backup", "")
	if err := a.Run(context.Background(), "修复一个 bug"); err != nil {
		t.Fatalf("Run = %v, want nil（应回落成功）", err)
	}
	if p.CurrentModel() != "backup" {
		t.Errorf("final model = %s, want backup（粘性）", p.CurrentModel())
	}
	if len(p.seen) != 2 || p.seen[0] != "primary" || p.seen[1] != "backup" {
		t.Fatalf("seen = %v, want primary→backup", p.seen)
	}
	// 回复来自备选模型。
	last := a.Session.Messages[len(a.Session.Messages)-1]
	if !strings.Contains(last.Content, "ok from backup") {
		t.Errorf("reply = %q", last.Content)
	}
}

// TestFallbackStickyAcrossTurns：回落是粘性的——下一回合直接在备选上跑，
// 不再撞一次首选的失败。
func TestFallbackStickyAcrossTurns(t *testing.T) {
	p := &routingStub{model: "primary", failOnce: map[string]int{"primary": 1}}
	a := newRoutingAgent(t, p, "backup", "")
	if err := a.Run(context.Background(), "第一回合"); err != nil {
		t.Fatalf("turn1 = %v", err)
	}
	if err := a.Run(context.Background(), "第二回合"); err != nil {
		t.Fatalf("turn2 = %v", err)
	}
	for i, m := range p.seen {
		if m == "primary" && i > 0 {
			t.Fatalf("seen = %v，粘性回落不应再回首选", p.seen)
		}
	}
}

// TestMultimodalRouting：含图回合临时切多模态，回合结束恢复原模型；
// 不带图的回合不路由。
func TestMultimodalRouting(t *testing.T) {
	p := &routingStub{model: "primary"}
	a := newRoutingAgent(t, p, "", "vision")
	if err := a.Session.Append(llm.Message{Role: llm.RoleUser, Content: "看这张图",
		Images: []llm.Image{{Filename: "shot.png"}}}); err != nil {
		t.Fatal(err)
	}
	if err := a.Run(context.Background(), ""); err != nil {
		t.Fatalf("Run = %v", err)
	}
	if len(p.seen) == 0 || p.seen[0] != "vision" {
		t.Fatalf("seen = %v, want 含图回合走 vision", p.seen)
	}
	if p.CurrentModel() != "primary" {
		t.Errorf("回合后应恢复原模型，got %s", p.CurrentModel())
	}

	// 不带图：不路由。
	if err := a.Run(context.Background(), "纯文本回合"); err != nil {
		t.Fatal(err)
	}
	if p.seen[len(p.seen)-1] != "primary" {
		t.Errorf("纯文本回合应留在原模型，seen = %v", p.seen)
	}
}

// TestNoFallbackWhenStreamed：已有增量文本流出后出错，不重试（会叠加重复
// 内容），错误原样返回。
func TestNoFallbackWhenStreamed(t *testing.T) {
	p := &streamFailStub{model: "primary", failModel: "primary"}
	a := newRoutingAgent(t, p, "backup", "")
	if err := a.Run(context.Background(), "任务"); err == nil {
		t.Fatal("流中失败且已流出文本：错误应原样返回")
	}
	if len(p.seen) != 1 {
		t.Fatalf("seen = %d 次，不应重试", len(p.seen))
	}
}

// streamFailStub 先流出一段文本再报错（模拟流中断）。
type streamFailStub struct {
	model     string
	failModel string
	seen      []string
}

func (p *streamFailStub) Stream(ctx context.Context, system string, msgs []llm.Message,
	tools []llm.ToolDef, opts llm.StreamOptions, onDelta func(llm.StreamDelta)) (llm.Message, error) {
	p.seen = append(p.seen, p.model)
	if p.model == p.failModel {
		onDelta(llm.StreamDelta{Text: "部分输出"})
		return llm.Message{}, errors.New("stream broken")
	}
	return llm.Message{Role: llm.RoleAssistant, Content: "ok"}, nil
}
func (p *streamFailStub) SetModel(model string) { p.model = model }
func (p *streamFailStub) CurrentModel() string  { return p.model }