promptkit/execution_group_contract_test.go

package promptkit_test

import (
	"context"
	"net/http"
	"testing"

	"go.maximumdirect.net/promptkit"
)

func TestExecutionGroupExplicitMembershipAndTerminalTotals(t *testing.T) {
	responses := []*http.Response{
		observedHTTPResponse(200, `{"choices":[{"message":{"content":"initial"}}],"usage":{"prompt_tokens":2}}`),
		observedHTTPResponse(503, `{}`),
		observedHTTPResponse(200, `{"choices":[{"message":{"content":"invalid"}}],"usage":{"prompt_tokens":3}}`),
		observedHTTPResponse(200, `{"choices":[{"message":{"content":"{\"ok\":true}"}}],"usage":{"prompt_tokens":4}}`),
		observedHTTPResponse(200, `{"choices":[{"message":{"content":"invalid"}}],"usage":{"prompt_tokens":5}}`),
		observedHTTPResponse(422, `{"error":{}}`),
		observedHTTPResponse(200, `{"choices":[{"message":{"content":"validator"}}],"usage":{"prompt_tokens":6}}`),
	}
	calls := 0
	one := 1
	config := simpleConfig("message")
	config.TransientRetries = &one
	config.HTTPClient = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
		response := responses[calls]
		calls++
		return response, nil
	})}
	engine, err := promptkit.NewEngine(config)
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()
	var all, selected promptkit.ExecutionGroup
	initial, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true, Observer: all.Observe})
	if err != nil {
		t.Fatal(err)
	}
	validation := &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}
	callbackCount := 0
	var branchID, validatorID string
	result, err := engine.Coordinate(ctx, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 1, Observer: promptkit.ComposeObservers(all.Observe, selected.Observe), Next: func(ctx context.Context, turn promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
		callbackCount++
		if turn.Result == nil {
			return &promptkit.RunRequest{PromptID: "prompt", Validation: validation}, nil
		}
		branchID = turn.Result.RunID
		return nil, nil
	}})
	if err != nil || result.Outcome != promptkit.CoordinationStopped || callbackCount != 2 {
		t.Fatalf("coordinate=%#v, %v", result, err)
	}
	if failed, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", ContinueFrom: initial.Exchange, Validation: validation, Observer: all.Observe}); err == nil || failed != nil {
		t.Fatalf("failed branch=%#v, %v", failed, err)
	}
	validator, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", Observer: selected.Observe})
	if err != nil {
		t.Fatal(err)
	}
	validatorID = validator.RunID
	a, b := all.Snapshot(), selected.Snapshot()
	if calls != 7 || len(a.Members) != 3 || a.Attempts != 6 || a.TransientRetries != 1 || a.RepairCalls != 2 || a.Usage.PromptTokens.Value != 14 || !a.Usage.PromptTokens.Known || !a.TerminalCoverageComplete {
		t.Fatalf("all summary=%#v calls=%d", a, calls)
	}
	if len(b.Members) != 2 || b.Attempts != 4 || b.Usage.PromptTokens.Value != 13 || !b.TerminalCoverageComplete {
		t.Fatalf("selected summary=%#v", b)
	}
	ids := map[string]bool{}
	for _, member := range b.Members {
		ids[member.RunID] = true
	}
	if !ids[branchID] || !ids[validatorID] || ids[initial.RunID] {
		t.Fatalf("selected members=%#v", b.Members)
	}
}

func TestExecutionGroupPreparedSilenceAndPanicCoverage(t *testing.T) {
	engine := newSimpleEngine(t, "message", promptkit.WithLLMClient(&fakeLLMClient{response: &promptkit.GenerateResponse{Content: "ok"}}))
	ctx := context.Background()
	var group promptkit.ExecutionGroup
	prepared, err := engine.PrepareExecution(ctx, promptkit.RunRequest{PromptID: "prompt", Observer: group.Observe})
	if err != nil {
		t.Fatal(err)
	}
	prepared.Discard()
	_, _ = engine.RunPrepared(ctx, prepared)
	if len(group.Snapshot().Members) != 0 {
		t.Fatal("preparation or failed claim enrolled a run")
	}
	func() {
		defer func() {
			if recover() == nil {
				t.Error("expected panic")
			}
		}()
		_, _ = engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", Observer: promptkit.ComposeObservers(group.Observe, func(event promptkit.ExecutionEvent) {
			_ = group.Snapshot()
			if event.Kind == promptkit.ExecutionStarted {
				panic("stop")
			}
		})})
	}()
	result := group.Snapshot()
	if len(result.Members) != 1 || result.TerminalCoverageComplete || result.Members[0].TerminalObserved || result.Attempts != 0 {
		t.Fatalf("panic summary=%#v", result)
	}
}

func TestExecutionGroupPreparedAndCallbackMembership(t *testing.T) {
	engine := newSimpleEngine(t, "message", promptkit.WithLLMClient(&fakeLLMClient{response: &promptkit.GenerateResponse{Content: "ok"}}))
	ctx := context.Background()
	var continuations, nested promptkit.ExecutionGroup
	initial, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", CaptureExchange: true})
	if err != nil {
		t.Fatal(err)
	}
	prepared, err := engine.PrepareExecution(ctx, promptkit.RunRequest{PromptID: "prompt", Observer: nested.Observe})
	if err != nil {
		t.Fatal(err)
	}
	preparedResult, err := engine.RunPrepared(ctx, prepared)
	if err != nil {
		t.Fatal(err)
	}
	var callbackID string
	_, err = engine.Coordinate(ctx, promptkit.CoordinateRequest{Start: initial.Exchange, AdditionalTurns: 1, Observer: continuations.Observe, Next: func(ctx context.Context, turn promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
		if turn.Result == nil {
			return &promptkit.RunRequest{PromptID: "prompt"}, nil
		}
		callbackResult, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "prompt", Observer: nested.Observe})
		if err == nil {
			callbackID = callbackResult.RunID
		}
		return nil, err
	}})
	if err != nil {
		t.Fatal(err)
	}
	a, b := continuations.Snapshot(), nested.Snapshot()
	if len(a.Members) != 1 || len(b.Members) != 2 || !a.TerminalCoverageComplete || !b.TerminalCoverageComplete {
		t.Fatalf("summaries = %#v %#v", a, b)
	}
	ids := map[string]bool{}
	for _, member := range b.Members {
		ids[member.RunID] = true
	}
	if !ids[preparedResult.RunID] || !ids[callbackID] || ids[initial.RunID] {
		t.Fatalf("members=%#v", b.Members)
	}
}