62 lines
2.7 KiB
Go
62 lines
2.7 KiB
Go
|
package promptkit_test
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"go.maximumdirect.net/promptkit"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestCoordinatePublicBoundary(t *testing.T) {
|
||
|
|
ctx := context.Background()
|
||
|
|
engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(&fakeLLMClient{response: &promptkit.GenerateResponse{Content: "answer"}}))
|
||
|
|
initial, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
for _, request := range []promptkit.CoordinateRequest{{}, {Start: &promptkit.Exchange{}, Next: func(context.Context, promptkit.ConversationTurn) (*promptkit.RunRequest, error) { return nil, nil }}, {Start: initial.Exchange, AdditionalTurns: -1, Next: func(context.Context, promptkit.ConversationTurn) (*promptkit.RunRequest, error) { return nil, nil }}} {
|
||
|
|
if _, err := engine.Coordinate(ctx, request); !errors.Is(err, promptkit.ErrInvalidRequest) {
|
||
|
|
t.Fatalf("invalid request error = %v", err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
if _, err := (*promptkit.Engine)(nil).Coordinate(ctx, promptkit.CoordinateRequest{}); !errors.Is(err, promptkit.ErrInvalidConfig) {
|
||
|
|
t.Fatalf("nil engine error = %v", err)
|
||
|
|
}
|
||
|
|
callbacks := 0
|
||
|
|
selected := &promptkit.RunRequest{PromptID: "prompt", ContinueFrom: &promptkit.Exchange{}}
|
||
|
|
result, err := engine.Coordinate(ctx, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 2, Next: func(_ context.Context, turn promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
|
||
|
|
callbacks++
|
||
|
|
if callbacks == 1 {
|
||
|
|
if turn.Result != nil || turn.RemainingTurns != 2 {
|
||
|
|
t.Fatalf("initial turn = %#v", turn)
|
||
|
|
}
|
||
|
|
return selected, nil
|
||
|
|
}
|
||
|
|
if turn.Result == nil || turn.Exchange.Details().ParentRunID == "" {
|
||
|
|
t.Fatalf("delivered turn = %#v", turn)
|
||
|
|
}
|
||
|
|
if callbacks == 2 {
|
||
|
|
turn.Result.Exchange = nil
|
||
|
|
*turn.Exchange = promptkit.Exchange{}
|
||
|
|
return selected, nil
|
||
|
|
}
|
||
|
|
if turn.Result.ParentRunID == initial.RunID {
|
||
|
|
t.Fatal("callback mutation redirected branch")
|
||
|
|
}
|
||
|
|
return nil, nil
|
||
|
|
}})
|
||
|
|
if err != nil || result.Outcome != promptkit.CoordinationStopped || callbacks != 3 || selected.ContinueFrom == nil {
|
||
|
|
t.Fatalf("result=%#v error=%v callbacks=%d", result, err, callbacks)
|
||
|
|
}
|
||
|
|
stop, err := engine.Coordinate(ctx, promptkit.CoordinateRequest{Start: initial.Exchange, Next: func(context.Context, promptkit.ConversationTurn) (*promptkit.RunRequest, error) { return nil, nil }})
|
||
|
|
if err != nil || stop.Outcome != promptkit.CoordinationStopped {
|
||
|
|
t.Fatalf("zero budget = %#v, %v", stop, err)
|
||
|
|
}
|
||
|
|
sentinel := errors.New("callback")
|
||
|
|
_, err = engine.Coordinate(ctx, promptkit.CoordinateRequest{Start: initial.Exchange, Next: func(context.Context, promptkit.ConversationTurn) (*promptkit.RunRequest, error) { return nil, sentinel }})
|
||
|
|
if !errors.Is(err, sentinel) {
|
||
|
|
t.Fatalf("callback error = %v", err)
|
||
|
|
}
|
||
|
|
}
|