promptkit/internal/usecase/prepared_snapshot_benchmark_test.go

package usecase

import (
	"context"
	"fmt"
	"reflect"
	"runtime"
	"testing"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
	"gitea.maximumdirect.net/eric/promptkit/internal/validate"
)

var preparedSnapshotHandleSink *PreparedExecution

const preparedSnapshotRetainedBatchSize = 4

type preparedSnapshotBenchmarkClient struct{}

func (preparedSnapshotBenchmarkClient) Generate(
	context.Context,
	domain.GenerateRequest,
) (*domain.GenerateResponse, error) {
	return &domain.GenerateResponse{Content: "ok"}, nil
}

type preparedSnapshotBenchmarkValidation struct{ schema any }

func (validation *preparedSnapshotBenchmarkValidation) Validate(
	context.Context,
	*domain.Artifact,
) (domain.ValidationResult, error) {
	return domain.ValidationResult{Status: domain.ValidationPassed, IsValid: true}, nil
}

func (validation *preparedSnapshotBenchmarkValidation) SchemaDocument() any {
	return validation.schema
}

type preparedSnapshotBenchmarkValidator struct{ schema any }

func (validator preparedSnapshotBenchmarkValidator) Validate(
	context.Context,
	*domain.Artifact,
	domain.OutputContract,
) (domain.ValidationResult, error) {
	return domain.ValidationResult{Status: domain.ValidationPassed, IsValid: true}, nil
}

func (validator preparedSnapshotBenchmarkValidator) PrepareValidation(
	context.Context,
	domain.OutputContract,
) (validate.PreparedValidation, error) {
	return &preparedSnapshotBenchmarkValidation{schema: validator.schema}, nil
}

func preparedSnapshotBenchmarkSchema(properties int) map[string]any {
	fields := make(map[string]any, properties)
	for index := range properties {
		fields[fmt.Sprintf("property-%03d", index)] = map[string]any{"type": "string"}
	}
	return map[string]any{
		"type": "object", "properties": fields, "additionalProperties": false,
	}
}

func preparedSnapshotBenchmarkExtras(entries int) map[string]any {
	providers := make(map[string]any, entries)
	for index := range entries {
		providers[fmt.Sprintf("provider-%03d", index)] = map[string]any{
			"enabled": true,
			"labels":  []any{fmt.Sprintf("tier-%d", index%4), fmt.Sprintf("region-%d", index%3)},
			"limits":  map[string]any{"burst": index + 1, "period": "minute"},
			"routes":  []any{map[string]any{"weight": index + 1, "fallback": true}},
		}
	}
	return map[string]any{"provider_options": providers}
}

func newPreparedSnapshotBenchmarkRunner(
	mutable bool,
) (*Runner, domain.RunRequest, any) {
	messageCount := 1
	extraEntries := 0
	schemaProperties := 1
	if mutable {
		messageCount = 8
		extraEntries = 64
		schemaProperties = 256
	}

	messages := make([]domain.RenderedMessage, messageCount)
	for index := range messages {
		messages[index] = domain.RenderedMessage{
			Role: domain.RoleUser, Content: fmt.Sprintf("rendered message %d", index),
			CacheControl: &domain.CacheControl{Type: domain.CacheControlEphemeral, TTL: "1h"},
		}
	}
	renderer := &fakeRenderer{rendered: &domain.RenderedPrompt{
		SessionID: "snapshot-session", Messages: messages,
	}}

	profile := defaultExecutionProfile()
	if extraEntries > 0 {
		profile.ResolvedTarget.ExtraParams = preparedSnapshotBenchmarkExtras(extraEntries)
	}
	profileRepository := &fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": profile}}
	definition := promptDef(domain.FormatJSON, domain.ValidationJSONSchema, 0)
	definition.Validation.SchemaPath = "snapshot-schema.json"
	schema := preparedSnapshotBenchmarkSchema(schemaProperties)
	validator := preparedSnapshotBenchmarkValidator{schema: schema}
	runner := NewRunner(
		&fakePromptRepo{def: definition},
		profileRepository,
		staticArtifactReader{},
		renderer,
		preparedSnapshotBenchmarkClient{},
		validator,
		nil,
	)
	request := domain.RunRequest{
		PromptID: "p", ProfileID: "exec",
		Inputs: singleInputRef(), Vars: map[string]string{"name": "Ada"},
		AppendedMessages: []domain.RenderedMessage{{
			Role: domain.RoleAssistant, Content: "appended message",
			CacheControl: &domain.CacheControl{Type: domain.CacheControlEphemeral, TTL: "1h"},
		}},
	}
	return runner, request, schema
}

func checkPreparedSnapshotBenchmarkFixture(b *testing.B, runner *Runner, request domain.RunRequest, schema any) {
	b.Helper()
	prepared, err := runner.PrepareExecution(context.Background(), request)
	if err != nil {
		b.Fatal(err)
	}
	details := prepared.Details()
	if details == nil || details.PromptID != request.PromptID || len(details.Messages) < 2 ||
		details.Messages[0].Content != "rendered message 0" || details.Messages[len(details.Messages)-1].Content != "appended message" ||
		details.InputHashes["transcript"] == "" || details.StructuredOutput == nil || details.StructuredOutput.JSONSchema == nil {
		prepared.Discard()
		b.Fatalf("prepared snapshot fixture is incomplete: %#v", details)
	}
	if got := details.StructuredOutput.JSONSchema.Schema; !reflect.DeepEqual(got, schema) {
		prepared.Discard()
		b.Fatalf("prepared schema differs from fixture: %#v", got)
	}
	result, err := runner.RunPrepared(context.Background(), prepared)
	if err != nil || result == nil || result.RawOutput != "ok" || !result.Validation.IsValid {
		b.Fatalf("completed fixture=(%+v, %v)", result, err)
	}
	completedDetails := prepared.Details()
	if completedDetails == nil || completedDetails.Messages[0].Content != "rendered message 0" ||
		!reflect.DeepEqual(completedDetails.StructuredOutput.JSONSchema.Schema, schema) {
		b.Fatalf("completed handle did not retain its details: %#v", completedDetails)
	}
}

func BenchmarkPreparedSnapshot(b *testing.B) {
	for _, fixture := range []struct {
		name    string
		mutable bool
	}{
		{name: "small"},
		{name: "mutable_payload", mutable: true},
	} {
		b.Run(fixture.name, func(b *testing.B) {
			runner, request, schema := newPreparedSnapshotBenchmarkRunner(fixture.mutable)
			checkPreparedSnapshotBenchmarkFixture(b, runner, request, schema)
			b.ReportAllocs()
			b.ResetTimer()
			for b.Loop() {
				prepared, err := runner.PrepareExecution(context.Background(), request)
				if err != nil {
					b.Fatal(err)
				}
				prepared.Discard()
				preparedSnapshotHandleSink = prepared
			}
			b.StopTimer()
			preparedSnapshotHandleSink = nil
		})
	}
}

func BenchmarkRetainedPreparedHandles(b *testing.B) {
	for _, state := range []string{"ready", "completed", "discarded"} {
		b.Run(state, func(b *testing.B) {
			runner, request, schema := newPreparedSnapshotBenchmarkRunner(true)
			checkPreparedSnapshotBenchmarkFixture(b, runner, request, schema)
			retained := make([]*PreparedExecution, preparedSnapshotRetainedBatchSize)
			runtime.GC()
			var setupMemory runtime.MemStats
			runtime.ReadMemStats(&setupMemory)
			runtime.KeepAlive(runner)
			runtime.KeepAlive(request)
			runtime.KeepAlive(retained)

			retainHandle := func() *PreparedExecution {
				prepared, err := runner.PrepareExecution(context.Background(), request)
				if err != nil {
					b.Fatal(err)
				}
				switch state {
				case "completed":
					result, err := runner.RunPrepared(context.Background(), prepared)
					if err != nil || result == nil || result.RawOutput != "ok" || !result.Validation.IsValid {
						b.Fatalf("complete retained handle=(%+v, %v)", result, err)
					}
				case "discarded":
					prepared.Discard()
				}
				return prepared
			}

			b.ReportAllocs()
			b.ResetTimer()
			for iteration := 0; iteration < b.N; iteration++ {
				retained[iteration%len(retained)] = retainHandle()
			}
			b.StopTimer()
			for index, prepared := range retained {
				if prepared == nil {
					retained[index] = retainHandle()
				}
			}
			runtime.GC()
			var retainedMemory runtime.MemStats
			runtime.ReadMemStats(&retainedMemory)
			b.ReportMetric(float64(int64(retainedMemory.HeapAlloc)-int64(setupMemory.HeapAlloc)), "retained-B")
			runtime.KeepAlive(runner)
			runtime.KeepAlive(request)
			runtime.KeepAlive(retained)
		})
	}
}