62 lines
2.0 KiB
Go
62 lines
2.0 KiB
Go
|
package promptkit_test
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"go.maximumdirect.net/promptkit"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestPreparedNotificationSelectionAndOrdering(t *testing.T) {
|
||
|
|
for _, split := range []bool{false, true} {
|
||
|
|
t.Run(map[bool]string{false: "run", true: "prepared"}[split], func(t *testing.T) {
|
||
|
|
calls := 0
|
||
|
|
engine := observationEngine(t, observationClientFunc(func(_ context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
||
|
|
if calls != 1 {
|
||
|
|
t.Fatalf("generation before notification: %d", calls)
|
||
|
|
}
|
||
|
|
return &promptkit.GenerateResponse{Content: "ok"}, nil
|
||
|
|
}))
|
||
|
|
req := promptkit.RunRequest{PromptID: "prompt", OnPrepared: func(_ context.Context, details promptkit.PreparedRun) {
|
||
|
|
calls++
|
||
|
|
if details.PromptID != "prompt" || len(details.Messages) == 0 {
|
||
|
|
t.Fatalf("details=%+v", details)
|
||
|
|
}
|
||
|
|
}}
|
||
|
|
if split {
|
||
|
|
handle, err := engine.PrepareExecution(context.Background(), req)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if calls != 0 {
|
||
|
|
t.Fatal("preparation notified")
|
||
|
|
}
|
||
|
|
req.OnPrepared = func(context.Context, promptkit.PreparedRun) { t.Fatal("selection changed") }
|
||
|
|
_, err = engine.RunPrepared(context.Background(), handle)
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
} else if _, err := engine.Run(context.Background(), req); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if calls != 1 {
|
||
|
|
t.Fatalf("notifications=%d", calls)
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestCancellationInPreparedNotificationPreventsGeneration(t *testing.T) {
|
||
|
|
ctx, cancel := context.WithCancel(context.Background())
|
||
|
|
defer cancel()
|
||
|
|
engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
||
|
|
t.Fatal("generation after cancellation")
|
||
|
|
return nil, nil
|
||
|
|
}))
|
||
|
|
_, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", OnPrepared: func(context.Context, promptkit.PreparedRun) { cancel() }})
|
||
|
|
if !errors.Is(err, context.Canceled) || !errors.Is(err, promptkit.ErrLLMGenerate) {
|
||
|
|
t.Fatalf("error=%v", err)
|
||
|
|
}
|
||
|
|
}
|