177 lines
7.5 KiB
Go
177 lines
7.5 KiB
Go
package promptkit_test
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit"
|
|
"gitea.maximumdirect.net/eric/promptkit/testfixture"
|
|
)
|
|
|
|
type observationClientFunc func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error)
|
|
|
|
func (f observationClientFunc) Generate(ctx context.Context, req promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
return f(ctx, req)
|
|
}
|
|
|
|
func TestDelegatedOperationsIsolateObservations(t *testing.T) {
|
|
for _, mode := range []string{"run", "observed run", "prepared", "observed prepared", "prepare", "prepare execution"} {
|
|
t.Run(mode, func(t *testing.T) {
|
|
type contextKey struct{}
|
|
ctx, cancel := context.WithTimeout(context.WithValue(context.Background(), contextKey{}, "caller-value"), time.Minute)
|
|
defer cancel()
|
|
deadline, _ := ctx.Deadline()
|
|
child, err := testfixture.NewEngine(testfixture.Config{},
|
|
testfixture.WithPromptFS(contractPromptFS("child", "profile", "child message"), "."),
|
|
testfixture.WithProfiles(testfixture.Profile{ID: "profile", Endpoint: "https://example.invalid/v1", Model: "child-model"}),
|
|
promptkit.WithLLMClient(observationClientFunc(func(callCtx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
gotDeadline, ok := callCtx.Deadline()
|
|
if !ok || gotDeadline != deadline || callCtx.Value(contextKey{}) != "caller-value" || callCtx.Done() != ctx.Done() {
|
|
t.Error("delegation lost caller context behavior")
|
|
}
|
|
return &promptkit.GenerateResponse{Content: "child", Usage: promptkit.TokenUsage{PromptTokens: 7}}, nil
|
|
})))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var childEvents []promptkit.ExecutionEvent
|
|
var childResult *promptkit.RunResult
|
|
outer := observationEngine(t, observationClientFunc(func(callCtx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
req := promptkit.RunRequest{PromptID: "child"}
|
|
if mode == "observed run" || mode == "observed prepared" {
|
|
req.Observer = func(event promptkit.ExecutionEvent) { childEvents = append(childEvents, event) }
|
|
}
|
|
var callErr error
|
|
switch mode {
|
|
case "run", "observed run":
|
|
childResult, callErr = child.Run(callCtx, req)
|
|
case "prepared", "observed prepared", "prepare execution":
|
|
var handle *promptkit.PreparedExecution
|
|
handle, callErr = child.PrepareExecution(callCtx, req)
|
|
if callErr == nil {
|
|
defer handle.Discard()
|
|
if mode != "prepare execution" {
|
|
childResult, callErr = child.RunPrepared(callCtx, handle)
|
|
}
|
|
}
|
|
case "prepare":
|
|
_, callErr = child.Prepare(callCtx, req)
|
|
}
|
|
return &promptkit.GenerateResponse{Content: "outer", Usage: promptkit.TokenUsage{PromptTokens: 1}}, callErr
|
|
}))
|
|
var outerEvents []promptkit.ExecutionEvent
|
|
result, err := outer.Run(ctx, promptkit.RunRequest{PromptID: "prompt", Observer: func(event promptkit.ExecutionEvent) { outerEvents = append(outerEvents, event) }})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(outerEvents) != 3 {
|
|
t.Fatalf("outer events=%+v", outerEvents)
|
|
}
|
|
for _, event := range outerEvents[1:] {
|
|
if event.RunID != result.RunID || event.PromptID != "prompt" || event.Model != "model" || event.Usage.PromptTokens.Value != 1 {
|
|
t.Fatalf("child contaminated outer event: %+v", event)
|
|
}
|
|
}
|
|
if outerEvents[2].Attempts != 1 {
|
|
t.Fatalf("outer attempts=%d", outerEvents[2].Attempts)
|
|
}
|
|
if childResult != nil && childResult.RunID == result.RunID {
|
|
t.Fatal("delegated run reused outer run ID")
|
|
}
|
|
if mode == "observed run" || mode == "observed prepared" {
|
|
if len(childEvents) != 3 || childEvents[0].Kind != promptkit.ExecutionStarted || childEvents[2].Kind != promptkit.ExecutionFinished || childEvents[2].Attempts != 1 || childEvents[2].Usage.PromptTokens.Value != 7 || childEvents[2].RunID != childResult.RunID {
|
|
t.Fatalf("child lifecycle=%+v", childEvents)
|
|
}
|
|
} else if len(childEvents) != 0 {
|
|
t.Fatalf("unobserved child events=%+v", childEvents)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestConcurrentDelegationDoesNotShareRecorder(t *testing.T) {
|
|
child := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
return &promptkit.GenerateResponse{Content: "child"}, nil
|
|
}))
|
|
ids := make(chan string, 2)
|
|
outer := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
errs := make(chan error, 2)
|
|
for range 2 {
|
|
go func() {
|
|
r, err := child.Run(ctx, promptkit.RunRequest{PromptID: "prompt"})
|
|
if r != nil {
|
|
ids <- r.RunID
|
|
}
|
|
errs <- err
|
|
}()
|
|
}
|
|
return &promptkit.GenerateResponse{Content: "outer"}, errors.Join(<-errs, <-errs)
|
|
}))
|
|
var events []promptkit.ExecutionEvent
|
|
r, err := outer.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", Observer: func(e promptkit.ExecutionEvent) { events = append(events, e) }})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(events) != 3 || events[2].Attempts != 1 {
|
|
t.Fatalf("outer events=%+v", events)
|
|
}
|
|
first, second := <-ids, <-ids
|
|
if first == second || first == r.RunID || second == r.RunID {
|
|
t.Fatal("delegated identities are not independent")
|
|
}
|
|
}
|
|
|
|
type classifiedObservationTestError struct{ cause error }
|
|
|
|
func (e *classifiedObservationTestError) Error() string { return "injected failure" }
|
|
func (e *classifiedObservationTestError) Unwrap() error { return e.cause }
|
|
|
|
func TestInjectedAttemptAndTerminalFailureCategoriesAgree(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
cause error
|
|
want promptkit.ExecutionFailureKind
|
|
}{
|
|
{"nil success", nil, promptkit.ExecutionFailureMalformedResponse},
|
|
{"malformed", promptkit.ErrMalformedResponse, promptkit.ExecutionFailureMalformedResponse},
|
|
{"transport", promptkit.ErrTransport, promptkit.ExecutionFailureTransport},
|
|
{"arbitrary", errors.New("transport failed"), promptkit.ExecutionFailureGeneration},
|
|
{"canceled", errors.Join(promptkit.ErrTransport, context.Canceled), promptkit.ExecutionFailureCanceled},
|
|
{"deadline", errors.Join(promptkit.ErrMalformedResponse, context.DeadlineExceeded), promptkit.ExecutionFailureDeadline},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
var injected error
|
|
var typed *classifiedObservationTestError
|
|
if tc.cause != nil {
|
|
typed = &classifiedObservationTestError{cause: tc.cause}
|
|
injected = fmt.Errorf("wrapped: %w", typed)
|
|
}
|
|
calls := 0
|
|
engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
|
calls++
|
|
return nil, injected
|
|
}))
|
|
var events []promptkit.ExecutionEvent
|
|
retries := 3
|
|
result, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", TransientRetries: &retries, Observer: func(e promptkit.ExecutionEvent) { events = append(events, e) }})
|
|
if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || calls != 1 || len(events) != 3 {
|
|
t.Fatalf("result=%+v err=%v calls=%d events=%+v", result, err, calls, events)
|
|
}
|
|
if events[1].Failure != tc.want || events[2].Failure != tc.want {
|
|
t.Fatalf("attempt=%s terminal=%s want=%s", events[1].Failure, events[2].Failure, tc.want)
|
|
}
|
|
if injected != nil {
|
|
var got *classifiedObservationTestError
|
|
if !errors.Is(err, injected) || !errors.Is(err, tc.cause) || !errors.As(err, &got) || got != typed {
|
|
t.Fatalf("lost injected identity: %v", err)
|
|
}
|
|
} else if !errors.Is(err, promptkit.ErrMalformedResponse) {
|
|
t.Fatalf("nil success error=%v", err)
|
|
}
|
|
})
|
|
}
|
|
}
|