wehub-resource-sync
743 lines
27 KiB
Go
743 lines
27 KiB
Go
/*
|
|
* Copyright 2026 CloudWeGo Authors
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
package adk
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/cloudwego/eino/components/model"
|
|
"github.com/cloudwego/eino/schema"
|
|
)
|
|
|
|
type fakeChatModel struct {
|
|
generate func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error)
|
|
stream func(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error)
|
|
callbacksEnabled bool
|
|
}
|
|
|
|
func (m *fakeChatModel) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) {
|
|
return m.generate(ctx, input, opts...)
|
|
}
|
|
|
|
func (m *fakeChatModel) Stream(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return m.stream(ctx, input, opts...)
|
|
}
|
|
|
|
func (m *fakeChatModel) IsCallbacksEnabled() bool {
|
|
return m.callbacksEnabled
|
|
}
|
|
|
|
func drainMessageStream(sr *schema.StreamReader[*schema.Message]) ([]*schema.Message, error) {
|
|
defer sr.Close()
|
|
var out []*schema.Message
|
|
for {
|
|
chunk, err := sr.Recv()
|
|
if err == io.EOF {
|
|
return out, nil
|
|
}
|
|
if err != nil {
|
|
return out, err
|
|
}
|
|
out = append(out, chunk)
|
|
}
|
|
}
|
|
|
|
func streamWithMidError(chunks []*schema.Message, err error) *schema.StreamReader[*schema.Message] {
|
|
sr, sw := schema.Pipe[*schema.Message](2)
|
|
go func() {
|
|
defer sw.Close()
|
|
for _, c := range chunks {
|
|
sw.Send(c, nil)
|
|
}
|
|
sw.Send(nil, err)
|
|
}()
|
|
return sr
|
|
}
|
|
|
|
func streamWithMidErrorControlled(chunks []*schema.Message, err error, firstSent chan struct{}, release chan struct{}) *schema.StreamReader[*schema.Message] {
|
|
sr, sw := schema.Pipe[*schema.Message](2)
|
|
go func() {
|
|
defer sw.Close()
|
|
for i, c := range chunks {
|
|
sw.Send(c, nil)
|
|
if i == 0 && firstSent != nil {
|
|
close(firstSent)
|
|
if release != nil {
|
|
<-release
|
|
}
|
|
}
|
|
}
|
|
sw.Send(nil, err)
|
|
}()
|
|
return sr
|
|
}
|
|
|
|
func TestFailoverCurrentModelContext(t *testing.T) {
|
|
t.Run("set and get", func(t *testing.T) {
|
|
ctx := context.Background()
|
|
m := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return schema.AssistantMessage("ok", nil), nil
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil
|
|
},
|
|
}
|
|
ctx = typedSetFailoverCurrentModel[*schema.Message](ctx, m)
|
|
got, ok := typedGetFailoverCurrentModel[*schema.Message](ctx)
|
|
require.True(t, ok)
|
|
require.Same(t, m, got)
|
|
})
|
|
|
|
t.Run("wrong type", func(t *testing.T) {
|
|
ctx := context.WithValue(context.Background(), failoverCurrentModelKey{}, "bad")
|
|
_, ok := typedGetFailoverCurrentModel[*schema.Message](ctx)
|
|
require.False(t, ok)
|
|
})
|
|
|
|
t.Run("missing", func(t *testing.T) {
|
|
_, ok := typedGetFailoverCurrentModel[*schema.Message](context.Background())
|
|
require.False(t, ok)
|
|
})
|
|
}
|
|
|
|
func TestFailoverProxyModel(t *testing.T) {
|
|
t.Run("generate missing context", func(t *testing.T) {
|
|
p := &failoverProxyModel{}
|
|
_, err := p.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")})
|
|
require.Error(t, err)
|
|
})
|
|
|
|
t.Run("stream missing context", func(t *testing.T) {
|
|
p := &failoverProxyModel{}
|
|
_, err := p.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")})
|
|
require.Error(t, err)
|
|
})
|
|
|
|
t.Run("generate routes to current model", func(t *testing.T) {
|
|
var called int32
|
|
target := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
atomic.AddInt32(&called, 1)
|
|
return schema.AssistantMessage("routed", nil), nil
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("routed", nil)}), nil
|
|
},
|
|
}
|
|
ctx := typedSetFailoverCurrentModel[*schema.Message](context.Background(), target)
|
|
p := &failoverProxyModel{}
|
|
msg, err := p.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "routed", msg.Content)
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&called))
|
|
})
|
|
}
|
|
|
|
func TestFailoverModelWrapper_Generate(t *testing.T) {
|
|
t.Run("delegates when GetFailoverModel nil", func(t *testing.T) {
|
|
var called int32
|
|
inner := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
atomic.AddInt32(&called, 1)
|
|
return schema.AssistantMessage("inner", nil), nil
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("inner", nil)}), nil
|
|
},
|
|
}
|
|
w := newFailoverModelWrapper[*schema.Message](inner, &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 2,
|
|
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
|
|
GetFailoverModel: nil,
|
|
})
|
|
msg, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "inner", msg.Content)
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&called))
|
|
})
|
|
|
|
t.Run("failover to second model", func(t *testing.T) {
|
|
wantErr := errors.New("first failed")
|
|
var shouldCalls int32
|
|
var m1Calls int32
|
|
var m2Calls int32
|
|
|
|
m1 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
atomic.AddInt32(&m1Calls, 1)
|
|
return nil, wantErr
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
}
|
|
m2 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
atomic.AddInt32(&m2Calls, 1)
|
|
return schema.AssistantMessage("ok", nil), nil
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 1,
|
|
ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool {
|
|
atomic.AddInt32(&shouldCalls, 1)
|
|
return errors.Is(err, wantErr)
|
|
},
|
|
GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
require.Equal(t, uint(1), failoverCtx.FailoverAttempt)
|
|
return m2, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
|
|
failoverLastSuccessModel: m1,
|
|
})
|
|
msg, err := w.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
|
|
require.NoError(t, err)
|
|
require.Equal(t, "ok", msg.Content)
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls))
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls))
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
|
|
})
|
|
|
|
t.Run("canceled error delegates to ShouldFailover", func(t *testing.T) {
|
|
var shouldCalls int32
|
|
m1 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, context.Canceled
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 5,
|
|
ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool {
|
|
atomic.AddInt32(&shouldCalls, 1)
|
|
// User decides to stop on canceled error
|
|
return !errors.Is(err, context.Canceled)
|
|
},
|
|
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
return m1, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
|
|
failoverLastSuccessModel: m1,
|
|
})
|
|
_, err := w.Generate(ctx, []*schema.Message{schema.UserMessage("hi")})
|
|
require.ErrorIs(t, err, context.Canceled)
|
|
// ShouldFailover is called once and returns false, stopping failover
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
|
|
})
|
|
|
|
t.Run("stops when GetFailoverModel returns error", func(t *testing.T) {
|
|
wantErr := errors.New("get model failed")
|
|
var called int32
|
|
inner := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
atomic.AddInt32(&called, 1)
|
|
return schema.AssistantMessage("unused", nil), nil
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 3,
|
|
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
|
|
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
return nil, nil, wantErr
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](inner, cfg)
|
|
_, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")})
|
|
require.ErrorIs(t, err, wantErr)
|
|
require.Equal(t, int32(0), atomic.LoadInt32(&called))
|
|
})
|
|
|
|
t.Run("stops when GetFailoverModel returns nil model", func(t *testing.T) {
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 1,
|
|
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
|
|
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
return nil, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
msg, err := w.Generate(context.Background(), []*schema.Message{schema.UserMessage("hi")})
|
|
require.Nil(t, msg)
|
|
require.Error(t, err)
|
|
require.ErrorContains(t, err, "GetFailoverModel returned nil model")
|
|
})
|
|
}
|
|
|
|
func TestFailoverModelWrapper_Stream(t *testing.T) {
|
|
t.Run("returns stream when first attempt succeeds", func(t *testing.T) {
|
|
var shouldCalls int32
|
|
in := schema.UserMessage("hi")
|
|
|
|
m1 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, input []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
require.Len(t, input, 1)
|
|
require.Same(t, in, input[0])
|
|
return schema.StreamReaderFromArray([]*schema.Message{
|
|
schema.AssistantMessage("a", nil),
|
|
schema.AssistantMessage("b", nil),
|
|
}), nil
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 0,
|
|
ShouldFailover: func(context.Context, *schema.Message, error) bool {
|
|
atomic.AddInt32(&shouldCalls, 1)
|
|
return false
|
|
},
|
|
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
return m1, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
|
|
failoverLastSuccessModel: m1,
|
|
})
|
|
sr, err := w.Stream(ctx, []*schema.Message{in})
|
|
require.NoError(t, err)
|
|
msgs, err := drainMessageStream(sr)
|
|
require.NoError(t, err)
|
|
require.Len(t, msgs, 2)
|
|
require.Equal(t, "a", msgs[0].Content)
|
|
require.Equal(t, "b", msgs[1].Content)
|
|
require.Equal(t, int32(0), atomic.LoadInt32(&shouldCalls))
|
|
})
|
|
|
|
t.Run("failover when Stream returns error immediately", func(t *testing.T) {
|
|
wantErr := errors.New("stream init failed")
|
|
var shouldCalls int32
|
|
var m1Calls int32
|
|
var m2Calls int32
|
|
|
|
m1 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
atomic.AddInt32(&m1Calls, 1)
|
|
return nil, wantErr
|
|
},
|
|
}
|
|
m2 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
atomic.AddInt32(&m2Calls, 1)
|
|
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("ok", nil)}), nil
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 1,
|
|
ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool {
|
|
atomic.AddInt32(&shouldCalls, 1)
|
|
return errors.Is(err, wantErr)
|
|
},
|
|
GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
require.Equal(t, uint(1), failoverCtx.FailoverAttempt)
|
|
return m2, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
|
|
failoverLastSuccessModel: m1,
|
|
})
|
|
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
|
|
require.NoError(t, err)
|
|
msgs, err := drainMessageStream(sr)
|
|
require.NoError(t, err)
|
|
require.Len(t, msgs, 1)
|
|
require.Equal(t, "ok", msgs[0].Content)
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls))
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls))
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
|
|
})
|
|
|
|
t.Run("failover when stream errors mid-way", func(t *testing.T) {
|
|
streamErr := errors.New("mid error")
|
|
var shouldCalls int32
|
|
var seenOutput atomic.Value
|
|
var m1Calls int32
|
|
var m2Calls int32
|
|
|
|
m1 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
atomic.AddInt32(&m1Calls, 1)
|
|
return streamWithMidError([]*schema.Message{
|
|
schema.AssistantMessage("p1", nil),
|
|
schema.AssistantMessage("p2", nil),
|
|
}, streamErr), nil
|
|
},
|
|
}
|
|
m2 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
atomic.AddInt32(&m2Calls, 1)
|
|
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("final", nil)}), nil
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 1,
|
|
ShouldFailover: func(_ context.Context, out *schema.Message, err error) bool {
|
|
atomic.AddInt32(&shouldCalls, 1)
|
|
if errors.Is(err, streamErr) && out != nil {
|
|
seenOutput.Store(out.Content)
|
|
}
|
|
return errors.Is(err, streamErr)
|
|
},
|
|
GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
require.Equal(t, uint(1), failoverCtx.FailoverAttempt)
|
|
return m2, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
|
|
failoverLastSuccessModel: m1,
|
|
})
|
|
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
|
|
require.NoError(t, err)
|
|
msgs, err := drainMessageStream(sr)
|
|
require.NoError(t, err)
|
|
require.Len(t, msgs, 1)
|
|
require.Equal(t, "final", msgs[0].Content)
|
|
require.Equal(t, "p1p2", seenOutput.Load())
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls))
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&m2Calls))
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
|
|
})
|
|
|
|
t.Run("stop when ShouldFailover returns false for mid-way error", func(t *testing.T) {
|
|
streamErr := errors.New("mid error")
|
|
m1 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return streamWithMidError([]*schema.Message{schema.AssistantMessage("p", nil)}, streamErr), nil
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 3,
|
|
ShouldFailover: func(context.Context, *schema.Message, error) bool {
|
|
return false
|
|
},
|
|
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
return m1, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
|
|
failoverLastSuccessModel: m1,
|
|
})
|
|
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
|
|
require.Nil(t, sr)
|
|
require.ErrorIs(t, err, streamErr)
|
|
})
|
|
|
|
t.Run("canceled mid-way error delegates to ShouldFailover", func(t *testing.T) {
|
|
var shouldCalls int32
|
|
m1 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
return streamWithMidError([]*schema.Message{schema.AssistantMessage("p", nil)}, context.Canceled), nil
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 3,
|
|
ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool {
|
|
atomic.AddInt32(&shouldCalls, 1)
|
|
// User decides to stop on canceled error
|
|
return !errors.Is(err, context.Canceled)
|
|
},
|
|
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
return m1, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
|
|
failoverLastSuccessModel: m1,
|
|
})
|
|
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
|
|
require.Nil(t, sr)
|
|
require.ErrorIs(t, err, context.Canceled)
|
|
// ShouldFailover is called once and returns false, stopping failover
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
|
|
})
|
|
|
|
t.Run("stop when Stream returns error immediately and ShouldFailover returns false", func(t *testing.T) {
|
|
wantErr := errors.New("stream init failed")
|
|
var shouldCalls int32
|
|
var m1Calls int32
|
|
|
|
m1 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
atomic.AddInt32(&m1Calls, 1)
|
|
return nil, wantErr
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 3,
|
|
ShouldFailover: func(_ context.Context, _ *schema.Message, err error) bool {
|
|
atomic.AddInt32(&shouldCalls, 1)
|
|
require.ErrorIs(t, err, wantErr)
|
|
return false
|
|
},
|
|
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
return m1, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
ctx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
|
|
failoverLastSuccessModel: m1,
|
|
})
|
|
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
|
|
require.Nil(t, sr)
|
|
require.ErrorIs(t, err, wantErr)
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls))
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&shouldCalls))
|
|
})
|
|
|
|
t.Run("stops when GetFailoverModel returns nil model", func(t *testing.T) {
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 1,
|
|
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
|
|
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
return nil, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
sr, err := w.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")})
|
|
require.Nil(t, sr)
|
|
require.Error(t, err)
|
|
require.ErrorContains(t, err, "GetFailoverModel returned nil model")
|
|
})
|
|
|
|
t.Run("stops when GetFailoverModel returns error", func(t *testing.T) {
|
|
wantErr := errors.New("get model failed")
|
|
var called int32
|
|
inner := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
atomic.AddInt32(&called, 1)
|
|
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 3,
|
|
ShouldFailover: func(context.Context, *schema.Message, error) bool { return true },
|
|
GetFailoverModel: func(_ context.Context, _ *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
return nil, nil, wantErr
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](inner, cfg)
|
|
sr, err := w.Stream(context.Background(), []*schema.Message{schema.UserMessage("hi")})
|
|
require.Nil(t, sr)
|
|
require.ErrorIs(t, err, wantErr)
|
|
require.Equal(t, int32(0), atomic.LoadInt32(&called))
|
|
})
|
|
|
|
t.Run("stops when ctx canceled during mid-way error handling", func(t *testing.T) {
|
|
midErr := errors.New("mid error")
|
|
var shouldCalls int32
|
|
var m1Calls int32
|
|
var m2Calls int32
|
|
firstSent := make(chan struct{})
|
|
release := make(chan struct{})
|
|
|
|
m1 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
atomic.AddInt32(&m1Calls, 1)
|
|
return streamWithMidErrorControlled(
|
|
[]*schema.Message{schema.AssistantMessage("p", nil)},
|
|
midErr,
|
|
firstSent,
|
|
release,
|
|
), nil
|
|
},
|
|
}
|
|
m2 := &fakeChatModel{
|
|
callbacksEnabled: true,
|
|
generate: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.Message, error) {
|
|
return nil, errors.New("unused")
|
|
},
|
|
stream: func(_ context.Context, _ []*schema.Message, _ ...model.Option) (*schema.StreamReader[*schema.Message], error) {
|
|
atomic.AddInt32(&m2Calls, 1)
|
|
return schema.StreamReaderFromArray([]*schema.Message{schema.AssistantMessage("unused", nil)}), nil
|
|
},
|
|
}
|
|
|
|
cfg := &ModelFailoverConfig[*schema.Message]{
|
|
MaxRetries: 1,
|
|
ShouldFailover: func(context.Context, *schema.Message, error) bool {
|
|
atomic.AddInt32(&shouldCalls, 1)
|
|
return true
|
|
},
|
|
GetFailoverModel: func(_ context.Context, failoverCtx *FailoverContext[*schema.Message]) (model.BaseChatModel, []*schema.Message, error) {
|
|
require.Equal(t, uint(1), failoverCtx.FailoverAttempt)
|
|
return m2, nil, nil
|
|
},
|
|
}
|
|
|
|
w := newFailoverModelWrapper[*schema.Message](&failoverProxyModel{}, cfg)
|
|
baseCtx := withTypedChatModelAgentExecCtx(context.Background(), &chatModelAgentExecCtx{
|
|
failoverLastSuccessModel: m1,
|
|
})
|
|
ctx, cancel := context.WithCancel(baseCtx)
|
|
type result struct {
|
|
sr *schema.StreamReader[*schema.Message]
|
|
err error
|
|
}
|
|
ch := make(chan result, 1)
|
|
go func() {
|
|
sr, err := w.Stream(ctx, []*schema.Message{schema.UserMessage("hi")})
|
|
ch <- result{sr: sr, err: err}
|
|
}()
|
|
|
|
<-firstSent
|
|
cancel()
|
|
close(release)
|
|
|
|
res := <-ch
|
|
if res.sr != nil {
|
|
res.sr.Close()
|
|
}
|
|
require.Nil(t, res.sr)
|
|
require.ErrorIs(t, res.err, midErr)
|
|
require.Equal(t, int32(1), atomic.LoadInt32(&m1Calls))
|
|
require.Equal(t, int32(0), atomic.LoadInt32(&m2Calls))
|
|
require.Equal(t, int32(0), atomic.LoadInt32(&shouldCalls))
|
|
})
|
|
}
|
|
|
|
func TestTypedConsumeStream_EmptyAgenticStream(t *testing.T) {
|
|
sr, sw := schema.Pipe[*schema.AgenticMessage](1)
|
|
sw.Close()
|
|
|
|
msg, err := typedConsumeStream(sr)
|
|
assert.Nil(t, err, "empty stream should not return error")
|
|
assert.NotNil(t, msg, "empty stream should return non-nil message from ConcatAgenticMessages")
|
|
}
|
|
|
|
func TestTypedConsumeStream_AgenticMidStreamError(t *testing.T) {
|
|
midErr := errors.New("mid-stream failure")
|
|
sr := streamWithMidErrorAgentic(
|
|
[]*schema.AgenticMessage{agenticChunk("chunk1"), agenticChunk("chunk2")},
|
|
midErr,
|
|
)
|
|
|
|
msg, err := typedConsumeStream(sr)
|
|
assert.ErrorIs(t, err, midErr, "should return the mid-stream error")
|
|
assert.NotNil(t, msg, "should return concatenated partial message from received chunks")
|
|
}
|
|
|
|
func streamWithMidErrorAgentic(chunks []*schema.AgenticMessage, err error) *schema.StreamReader[*schema.AgenticMessage] {
|
|
sr, sw := schema.Pipe[*schema.AgenticMessage](len(chunks) + 1)
|
|
go func() {
|
|
defer sw.Close()
|
|
for _, c := range chunks {
|
|
sw.Send(c, nil)
|
|
}
|
|
sw.Send(nil, err)
|
|
}()
|
|
return sr
|
|
}
|
|
|
|
func agenticChunk(text string) *schema.AgenticMessage {
|
|
return &schema.AgenticMessage{
|
|
Role: schema.AgenticRoleTypeAssistant,
|
|
ContentBlocks: []*schema.ContentBlock{
|
|
schema.NewContentBlock(&schema.AssistantGenText{Text: text}),
|
|
},
|
|
}
|
|
}
|