196 lines
9.2 KiB
Go
196 lines
9.2 KiB
Go
package promptkit_test
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/ConfigKit"
|
|
"gitea.maximumdirect.net/eric/promptkit"
|
|
)
|
|
|
|
func TestPreparedObservationStartsOnlyAfterClaimAndKeepsObserver(t *testing.T) {
|
|
engine := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "hello")
|
|
var events []promptkit.ExecutionEvent
|
|
request := promptkit.RunRequest{PromptID: "prepared", Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }}
|
|
prepared, err := engine.PrepareExecution(context.Background(), request)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
request.Observer = func(promptkit.ExecutionEvent) { t.Error("later observer was invoked") }
|
|
if len(events) != 0 {
|
|
t.Fatalf("preparation emitted events: %+v", events)
|
|
}
|
|
result, err := engine.RunPrepared(context.Background(), prepared)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(events) != 3 || events[0].Kind != promptkit.ExecutionStarted || events[1].Kind != promptkit.ExecutionAttempt || events[2].Kind != promptkit.ExecutionFinished {
|
|
t.Fatalf("events=%+v", events)
|
|
}
|
|
if events[0].PromptID != "prepared" || events[0].ProfileID != "profile" || events[0].Model != "model" || events[0].RunID != result.RunID || events[2].RunID != result.RunID || events[2].Attempts != 1 {
|
|
t.Fatalf("result=%+v events=%+v", result, events)
|
|
}
|
|
if again, err := engine.RunPrepared(context.Background(), prepared); again != nil || !errors.Is(err, promptkit.ErrInvalidRequest) || len(events) != 3 {
|
|
t.Fatalf("reused result=%+v error=%v events=%+v", again, err, events)
|
|
}
|
|
}
|
|
|
|
func TestPreparedObservationInvalidForeignAndDiscardedHandlesAreSilent(t *testing.T) {
|
|
engine := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "hello")
|
|
foreign := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "hello")
|
|
var events []promptkit.ExecutionEvent
|
|
request := promptkit.RunRequest{PromptID: "prepared", Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }}
|
|
prepared, err := engine.PrepareExecution(context.Background(), request)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if result, err := foreign.RunPrepared(context.Background(), prepared); result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) || len(events) != 0 {
|
|
t.Fatalf("foreign result=%+v error=%v events=%+v", result, err, events)
|
|
}
|
|
prepared.Discard()
|
|
if result, err := engine.RunPrepared(context.Background(), prepared); result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) || len(events) != 0 {
|
|
t.Fatalf("discarded result=%+v error=%v events=%+v", result, err, events)
|
|
}
|
|
if result, err := engine.RunPrepared(context.Background(), nil); result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) || len(events) != 0 {
|
|
t.Fatalf("nil result=%+v error=%v events=%+v", result, err, events)
|
|
}
|
|
}
|
|
|
|
func TestPreparedObservationCountsCompletedUsageBeforeRepairFailure(t *testing.T) {
|
|
client := &preparedRecordingClient{responses: []*promptkit.GenerateResponse{{Content: "invalid json", Usage: promptkit.TokenUsage{PromptTokens: 4}}, nil}}
|
|
engine := newPreparedContractEngine(t, client, "hello")
|
|
var events []promptkit.ExecutionEvent
|
|
prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared", Validation: &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}, Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
result, err := engine.RunPrepared(context.Background(), prepared)
|
|
if result != nil || !errors.Is(err, promptkit.ErrLLMGenerate) || len(events) != 4 {
|
|
t.Fatalf("result=%+v error=%v events=%+v", result, err, events)
|
|
}
|
|
terminal := events[3]
|
|
if terminal.Outcome != promptkit.ExecutionOutcomeError || terminal.Failure != promptkit.ExecutionFailureMalformedResponse || terminal.Attempts != 2 || terminal.RepairCalls != 1 || terminal.Usage.PromptTokens.Value != 4 || terminal.Usage.PromptTokens.Complete {
|
|
t.Fatalf("terminal=%+v", terminal)
|
|
}
|
|
}
|
|
|
|
func TestPreparedObservationClaimRaceAndPanicCleanup(t *testing.T) {
|
|
engine := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "hello")
|
|
var mu sync.Mutex
|
|
var events []promptkit.ExecutionEvent
|
|
prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared", Observer: func(event promptkit.ExecutionEvent) { mu.Lock(); events = append(events, event); mu.Unlock() }})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var wg sync.WaitGroup
|
|
var successes, invalid int
|
|
var countMu sync.Mutex
|
|
for range 8 {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
_, err := engine.RunPrepared(context.Background(), prepared)
|
|
countMu.Lock()
|
|
defer countMu.Unlock()
|
|
if err == nil {
|
|
successes++
|
|
} else if errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
invalid++
|
|
} else {
|
|
t.Errorf("unexpected error: %v", err)
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
if successes != 1 || invalid != 7 || len(events) != 3 {
|
|
t.Fatalf("successes=%d invalid=%d events=%+v", successes, invalid, events)
|
|
}
|
|
|
|
panicking, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared", Observer: func(event promptkit.ExecutionEvent) {
|
|
if event.Kind == promptkit.ExecutionStarted {
|
|
panic("callback")
|
|
}
|
|
}})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
func() {
|
|
defer func() {
|
|
if recover() == nil {
|
|
t.Error("expected panic")
|
|
}
|
|
}()
|
|
_, _ = engine.RunPrepared(context.Background(), panicking)
|
|
}()
|
|
if result, err := engine.RunPrepared(context.Background(), panicking); result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) {
|
|
t.Fatalf("reused after panic: result=%+v error=%v", result, err)
|
|
}
|
|
}
|
|
|
|
func TestPreparedObservationFinishesAfterCredentialAndCapacityFailures(t *testing.T) {
|
|
const envName = "PROMPTKIT_PREPARED_OBSERVATION_KEY"
|
|
t.Setenv(envName, "available")
|
|
engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &configkit.Input{
|
|
Source: configkit.Mapping("credential preparation", map[string]any{
|
|
"profiles": map[string]any{"profile": map[string]any{
|
|
"endpoint": "https://example.invalid/v1", "model": "model", "api_key_required": true,
|
|
}},
|
|
"prompts": map[string]any{"prepared": map[string]any{"1": map[string]any{
|
|
"default_profile": "profile",
|
|
"messages": []any{map[string]any{"role": "user", "content": "hello"}},
|
|
"output": map[string]any{"format": "text", "validation_mode": "none"},
|
|
}}},
|
|
}),
|
|
}}}, promptkit.WithLLMClient(&fakeLLMClient{response: &promptkit.GenerateResponse{Content: "ok"}}))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var events []promptkit.ExecutionEvent
|
|
prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared", Execution: &promptkit.ExecutionTargetOverride{APIKeyEnv: envName}, Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Setenv(envName, "")
|
|
result, err := engine.RunPrepared(context.Background(), prepared)
|
|
if result != nil || !errors.Is(err, promptkit.ErrInvalidRequest) || len(events) != 2 || events[1].Failure != promptkit.ExecutionFailureInvalidRequest || events[1].Attempts != 0 {
|
|
t.Fatalf("credential result=%+v error=%v events=%+v", result, err, events)
|
|
}
|
|
|
|
release := make(chan struct{})
|
|
client := newCapacityGateClient(release, 2)
|
|
limited := newBackendCapacityEngine(t, client, 1, capacityInt(0), nil)
|
|
first := make(chan capacityRunResult, 1)
|
|
go runCapacityRequest(limited, context.Background(), promptkit.RunRequest{PromptID: "prompt"}, first)
|
|
awaitCapacityRequest(t, client.started)
|
|
events = nil
|
|
prepared, err = limited.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prompt", Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
result, err = limited.RunPrepared(context.Background(), prepared)
|
|
close(release)
|
|
if firstResult := awaitCapacityRun(t, first); firstResult.err != nil {
|
|
t.Fatalf("first: %v", firstResult.err)
|
|
}
|
|
if result != nil || !errors.Is(err, promptkit.ErrCapacityExceeded) || len(events) != 2 || events[1].Failure != promptkit.ExecutionFailureCapacity || events[1].Attempts != 0 {
|
|
t.Fatalf("capacity result=%+v error=%v events=%+v", result, err, events)
|
|
}
|
|
}
|
|
|
|
func TestPreparedObservationCanceledClaimStillFinishes(t *testing.T) {
|
|
engine := newPreparedContractEngine(t, &preparedRecordingClient{response: &promptkit.GenerateResponse{Content: "ok"}}, "hello")
|
|
var events []promptkit.ExecutionEvent
|
|
prepared, err := engine.PrepareExecution(context.Background(), promptkit.RunRequest{PromptID: "prepared", Observer: func(event promptkit.ExecutionEvent) { events = append(events, event) }})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
result, err := engine.RunPrepared(ctx, prepared)
|
|
if result != nil || !errors.Is(err, context.Canceled) || len(events) != 3 || events[1].Dispatched || events[2].Failure != promptkit.ExecutionFailureCanceled || events[2].Attempts != 0 {
|
|
t.Fatalf("result=%+v error=%v events=%+v", result, err, events)
|
|
}
|
|
}
|