79 lines
3.0 KiB
Go
79 lines
3.0 KiB
Go
|
package promptkit_test
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"reflect"
|
||
|
|
"sync"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"go.maximumdirect.net/promptkit"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestComposeObserversCopiesAndOrders(t *testing.T) {
|
||
|
|
if promptkit.ComposeObservers() != nil || promptkit.ComposeObservers(nil, nil) != nil {
|
||
|
|
t.Fatal("empty composition is not nil")
|
||
|
|
}
|
||
|
|
var calls []int
|
||
|
|
first := func(promptkit.ExecutionEvent) { calls = append(calls, 1) }
|
||
|
|
second := func(promptkit.ExecutionEvent) { calls = append(calls, 2) }
|
||
|
|
values := []promptkit.ExecutionObserver{first, nil, second, first}
|
||
|
|
observer := promptkit.ComposeObservers(values...)
|
||
|
|
values[0] = second
|
||
|
|
observer(promptkit.ExecutionEvent{})
|
||
|
|
if !reflect.DeepEqual(calls, []int{1, 2, 1}) {
|
||
|
|
t.Fatalf("calls=%v", calls)
|
||
|
|
}
|
||
|
|
defer func() {
|
||
|
|
if recover() != "panic" || len(calls) != 3 {
|
||
|
|
t.Errorf("panic or calls=%v", calls)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
promptkit.ComposeObservers(func(promptkit.ExecutionEvent) { panic("panic") }, second)(promptkit.ExecutionEvent{})
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestCoordinateObserverCompositionAndNestedSelection(t *testing.T) {
|
||
|
|
engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(&fakeLLMClient{response: &promptkit.GenerateResponse{Content: "answer"}}))
|
||
|
|
ctx := context.Background()
|
||
|
|
initial, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
|
||
|
|
if err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
var sequence []string
|
||
|
|
var nested int
|
||
|
|
request := &promptkit.RunRequest{PromptID: "prompt", Observer: func(promptkit.ExecutionEvent) { sequence = append(sequence, "request") }}
|
||
|
|
result, err := engine.Coordinate(ctx, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 1, Observer: func(promptkit.ExecutionEvent) { sequence = append(sequence, "coordinator") }, Next: func(ctx context.Context, turn promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
|
||
|
|
if turn.Result == nil {
|
||
|
|
return request, nil
|
||
|
|
}
|
||
|
|
_, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", Observer: func(promptkit.ExecutionEvent) { nested++ }})
|
||
|
|
return nil, err
|
||
|
|
}})
|
||
|
|
if err != nil || result.Outcome != promptkit.CoordinationStopped || nested != 3 || !reflect.DeepEqual(sequence, []string{"request", "coordinator", "request", "coordinator", "request", "coordinator"}) {
|
||
|
|
t.Fatalf("result=%#v error=%v events=%v nested=%d", result, err, sequence, nested)
|
||
|
|
}
|
||
|
|
if request.ContinueFrom != nil {
|
||
|
|
t.Fatal("coordinator changed caller request")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestComposedObserversSupportConcurrentRuns(t *testing.T) {
|
||
|
|
var mu sync.Mutex
|
||
|
|
count := 0
|
||
|
|
observer := promptkit.ComposeObservers(func(promptkit.ExecutionEvent) { mu.Lock(); count++; mu.Unlock() })
|
||
|
|
engine := newSimpleEngine(t, "question", promptkit.WithLLMClient(&preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "answer"}}))
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
for i := 0; i < 8; i++ {
|
||
|
|
wg.Add(1)
|
||
|
|
go func() {
|
||
|
|
defer wg.Done()
|
||
|
|
if _, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", Observer: observer}); err != nil {
|
||
|
|
t.Error(err)
|
||
|
|
}
|
||
|
|
}()
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
if count != 24 {
|
||
|
|
t.Fatalf("event count=%d", count)
|
||
|
|
}
|
||
|
|
}
|