promptkit/assembled_workflow_benchmark_test.go

package promptkit_test

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"testing"
	"testing/fstest"

	"gitea.maximumdirect.net/eric/ConfigKit"
	"gitea.maximumdirect.net/eric/promptkit"
)

var (
	assembledPreparedSink *promptkit.PreparedRun
	assembledResultSink   *promptkit.RunResult
	assembledDetailsSink  promptkit.PreparedRun
	assembledHandleSink   *promptkit.PreparedExecution
)

type assembledWorkflowClient struct{ content string }

func (client assembledWorkflowClient) Generate(
	context.Context,
	promptkit.GenerateRequest,
) (*promptkit.GenerateResponse, error) {
	return &promptkit.GenerateResponse{Content: client.content}, nil
}

type assembledWorkflowFixture struct {
	engine          *promptkit.Engine
	request         promptkit.RunRequest
	expectedMessage string
	response        string
}

func assembledWorkflowSchema(b *testing.B) []byte {
	b.Helper()
	properties := make(map[string]any, 256)
	for index := range 256 {
		properties[fmt.Sprintf("property-%03d", index)] = map[string]any{"type": "string"}
	}
	data, err := json.Marshal(map[string]any{
		"$schema":              "https://json-schema.org/draft/2020-12/schema",
		"type":                 "object",
		"properties":           properties,
		"additionalProperties": false,
	})
	if err != nil {
		b.Fatal(err)
	}
	return data
}

func assembledWorkflowExtras(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 newAssembledWorkflowFixture(b *testing.B, name string) assembledWorkflowFixture {
	b.Helper()
	message := "Hello {{.name}}"
	response := "ok"
	output := map[string]any{"format": "text", "validation_mode": "none"}
	profile := map[string]any{"endpoint": "http://example.test/v1", "model": "benchmark-model"}
	options := []promptkit.Option{promptkit.WithLLMClient(assembledWorkflowClient{content: response})}
	var expectedMessage string

	switch name {
	case "small":
		expectedMessage = "Hello Ada"
	case "large_message":
		prefix := strings.Repeat("x", 256*1024)
		message = prefix + "{{.name}}"
		expectedMessage = prefix + "Ada"
	case "mutable_payload":
		message = `{"subject":"{{.name}}"}`
		expectedMessage = `{"subject":"Ada"}`
		response = "{}"
		profile["extra_params"] = assembledWorkflowExtras(64)
		output = map[string]any{
			"format": "json", "validation_mode": "json_schema", "schema_path": "workflow-schema.json",
		}
		options[0] = promptkit.WithLLMClient(assembledWorkflowClient{content: response})
		options = append(options, promptkit.WithSchemaFS(fstest.MapFS{
			"workflow-schema.json": &fstest.MapFile{Data: assembledWorkflowSchema(b)},
		}, "."))
	default:
		b.Fatalf("unknown assembled workflow fixture %q", name)
	}

	config := map[string]any{
		"profiles": map[string]any{"profile": profile},
		"prompts": map[string]any{"workflow": map[string]any{
			"version":         "v1",
			"default_profile": "profile",
			"messages":        []any{map[string]any{"role": "user", "content": message}},
			"output":          output,
		}},
	}
	engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &configkit.Input{
		Source: configkit.Mapping("assembled workflow benchmark", config),
	}}}, options...)
	if err != nil {
		b.Fatal(err)
	}
	request := promptkit.RunRequest{
		PromptID: "workflow", Vars: map[string]string{"name": "Ada"},
	}
	return assembledWorkflowFixture{
		engine: engine, request: request, expectedMessage: expectedMessage, response: response,
	}
}

func verifyAssembledWorkflowFixture(b *testing.B, fixture assembledWorkflowFixture) {
	b.Helper()
	ctx := context.Background()
	prepared, err := fixture.engine.Prepare(ctx, fixture.request)
	if err != nil {
		b.Fatal(err)
	}
	if prepared.PromptID != "workflow" || prepared.PromptVersion == nil || *prepared.PromptVersion != "v1" || len(prepared.Messages) != 1 ||
		prepared.Messages[0].Content != fixture.expectedMessage {
		b.Fatalf("prepared fixture output is invalid: %+v", prepared)
	}
	handle, err := fixture.engine.PrepareExecution(ctx, fixture.request)
	if err != nil {
		b.Fatal(err)
	}
	details := handle.Details()
	if details.PromptID != "workflow" || len(details.Messages) != 1 || details.Messages[0].Content != fixture.expectedMessage {
		handle.Discard()
		b.Fatalf("prepared execution details are invalid: %+v", details)
	}
	handle.Discard()
	result, err := fixture.engine.Run(ctx, fixture.request)
	if err != nil {
		b.Fatal(err)
	}
	if result.RawOutput != fixture.response {
		b.Fatalf("run output=%q, want %q", result.RawOutput, fixture.response)
	}
	if fixture.response == "{}" && result.Validation.Status != promptkit.ValidationPassed {
		b.Fatalf("structured response validation=%+v, want passed", result.Validation)
	}
}

func BenchmarkEnginePreparation(b *testing.B) {
	for _, name := range []string{"small", "large_message", "mutable_payload"} {
		b.Run(name, func(b *testing.B) {
			fixture := newAssembledWorkflowFixture(b, name)
			verifyAssembledWorkflowFixture(b, fixture)
			ctx := context.Background()
			b.Run("prepare", func(b *testing.B) {
				b.ReportAllocs()
				b.ResetTimer()
				for b.Loop() {
					prepared, err := fixture.engine.Prepare(ctx, fixture.request)
					if err != nil {
						b.Fatal(err)
					}
					assembledPreparedSink = prepared
				}
				b.StopTimer()
				assembledPreparedSink = nil
			})
			b.Run("prepared_execution_creation_and_discard", func(b *testing.B) {
				b.ReportAllocs()
				b.ResetTimer()
				for b.Loop() {
					handle, err := fixture.engine.PrepareExecution(ctx, fixture.request)
					if err != nil {
						b.Fatal(err)
					}
					handle.Discard()
					assembledHandleSink = handle
				}
				b.StopTimer()
				assembledHandleSink = nil
			})
		})
	}
}

func BenchmarkEngineRun(b *testing.B) {
	for _, name := range []string{"small", "large_message", "mutable_payload"} {
		b.Run(name, func(b *testing.B) {
			fixture := newAssembledWorkflowFixture(b, name)
			verifyAssembledWorkflowFixture(b, fixture)
			ctx := context.Background()
			b.ReportAllocs()
			b.ResetTimer()
			for b.Loop() {
				result, err := fixture.engine.Run(ctx, fixture.request)
				if err != nil {
					b.Fatal(err)
				}
				assembledResultSink = result
			}
			b.StopTimer()
			assembledResultSink = nil
		})
	}
}

func BenchmarkPreparedExecutionDetails(b *testing.B) {
	for _, name := range []string{"small", "large_message", "mutable_payload"} {
		b.Run(name, func(b *testing.B) {
			fixture := newAssembledWorkflowFixture(b, name)
			verifyAssembledWorkflowFixture(b, fixture)
			handle, err := fixture.engine.PrepareExecution(context.Background(), fixture.request)
			if err != nil {
				b.Fatal(err)
			}
			defer handle.Discard()
			initial := handle.Details()
			if len(initial.Messages) != 1 || initial.Messages[0].Content != fixture.expectedMessage {
				b.Fatalf("prepared details fixture is invalid: %+v", initial)
			}
			b.ReportAllocs()
			b.ResetTimer()
			for b.Loop() {
				assembledDetailsSink = handle.Details()
			}
			b.StopTimer()
			assembledDetailsSink = promptkit.PreparedRun{}
		})
	}
}