commandkit/integrations/ckprompt/smoke_command.go

package ckprompt

import (
	"errors"
	"fmt"
	"io"

	"go.maximumdirect.net/ckconfig"
	commandkit "go.maximumdirect.net/commandkit"
	promptkit "go.maximumdirect.net/promptkit"
)

type smokeCommandInput[I any] struct {
	selection ckconfig.Selection
	kind      string
	request   promptkit.SmokeTestRequest
	execution ExecutionInput
	resources ResourceInput
	captured  commandkit.FlagValues[I]
	selected  I
}

func newSmokeTestCommand[I, S any](suite Suite[I, S]) commandkit.Command {
	flags := commandkit.CaptureFlags(suite.Selectors, func(input **smokeCommandInput[I]) *commandkit.FlagValues[I] {
		if input == nil || *input == nil {
			return nil
		}
		return &(*input).captured
	})
	flags = append(flags,
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "config", Usage: "Configuration input path.", Completion: commandkit.CompleteFiles()}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				(*input).selection.Config = value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "overlay", Usage: "Configuration overlay path.", Completion: commandkit.CompleteFiles()}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				(*input).selection.Overlay = value
			}),
	)

	resources := NewResourceFlags(func(input **smokeCommandInput[I]) *ResourceInput { return &(*input).resources })
	execution := NewExecutionFlags(func(input **smokeCommandInput[I]) *ExecutionInput { return &(*input).execution })
	flags = append(flags, resources.Backend, resources.Prompt, resources.Profile)
	flags = append(flags, execution.All()...)

	return commandkit.NewCommand(commandkit.CommandSpec[*smokeCommandInput[I]]{
		Name:    "test",
		Summary: "Run one PromptKit smoke test.",
		NewInput: func() *smokeCommandInput[I] {
			return &smokeCommandInput[I]{}
		},
		Flags: flags,
		Arguments: []commandkit.Argument[*smokeCommandInput[I]]{
			commandkit.ValueArgument(commandkit.ArgumentMeta{Name: "kind", Usage: "Native smoke-test kind.", Completion: smokeKindCompletion()}, true, parseDefinitionString,
				func(input **smokeCommandInput[I], value commandkit.Supplied[string]) { (*input).kind = value.Value }),
		},
		CheckInput: func(input *smokeCommandInput[I]) error {
			if input == nil {
				return errors.New("PromptKit smoke-test input must not be nil")
			}
			input.request.BackendID = input.resources.Backend.Value
			input.request.PromptID = input.resources.Prompt.Value
			input.request.ProfileID = input.resources.Profile.Value
			input.request.Execution = input.execution.Execution
			input.request.TransientRetries = input.execution.TransientRetries
			if input.kind != "all" {
				input.request.Kind = promptkit.SmokeTestKind(input.kind)
			}
			selected, err := prepareSuiteInput(suite, input.selection, input.captured)
			if err != nil {
				return err
			}
			input.selected = selected
			return nil
		},
		Run: func(invocation commandkit.Invocation, input *smokeCommandInput[I]) commandkit.Outcome {
			return runSmokeTest(invocation, suite, input)
		},
	})
}

func runSmokeTest[I, S any](invocation commandkit.Invocation, suite Suite[I, S], input *smokeCommandInput[I]) commandkit.Outcome {
	if err := promptKitContextError(invocation.Context); err != nil {
		return definitionFailure("PromptKit smoke testing was canceled.", commandkit.Explanation{}, err)
	}
	var state S
	if suite.NewState != nil {
		state = suite.NewState()
	}
	if err := promptKitContextError(invocation.Context); err != nil {
		return definitionFailure("PromptKit smoke testing was canceled.", commandkit.Explanation{}, err)
	}
	if input.kind == "all" {
		return runSmokeAggregate(invocation, suite, input, state)
	}

	request := SmokeRequest{Selection: input.selection, Request: input.request}
	acquired, operationErr := suite.SmokeTest(invocation.Context, request, input.selected, state)
	queuePromptDiagnostics(invocation.Reporter(), acquired.Diagnostics)
	if acquired.Value == nil && operationErr == nil {
		operationErr = errors.New("PromptKit smoke test returned neither a report nor an error")
	}
	if err := promptKitContextError(invocation.Context); err != nil && !errors.Is(operationErr, err) {
		if operationErr == nil {
			operationErr = err
		} else {
			operationErr = errors.Join(operationErr, err)
		}
	}
	if acquired.Value == nil {
		return definitionFailure(smokeFailureMessage(input.kind), acquired.Explanation, operationErr)
	}

	_ = invocation.Reporter().Report(commandkit.Report{Value: smokeReportFromNative(acquired.Value)})
	if operationErr != nil {
		return definitionFailure(smokeFailureMessage(input.kind), acquired.Explanation, operationErr)
	}
	return commandkit.Success()
}

func aggregateSmokeKinds() [4]promptkit.SmokeTestKind {
	return [...]promptkit.SmokeTestKind{
		promptkit.SmokeTestReachability,
		promptkit.SmokeTestAuthentication,
		promptkit.SmokeTestTextInference,
		promptkit.SmokeTestStructuredOutput,
	}
}

func runSmokeAggregate[I, S any](invocation commandkit.Invocation, suite Suite[I, S], input *smokeCommandInput[I], state S) commandkit.Outcome {
	kinds := aggregateSmokeKinds()
	results := make([]smokeResult, 0, len(kinds))
	var operationErrors []error
	var explanations []commandkit.Explanation

	for _, kind := range kinds {
		if err := promptKitContextError(invocation.Context); err != nil {
			if len(results) == 0 {
				return definitionFailure("PromptKit aggregate smoke testing was canceled.", commandkit.Explanation{}, err)
			}
			operationErrors = append(operationErrors, err)
			break
		}

		seed := smokeRequestForKind(input.request, kind)
		acquired, operationErr := suite.SmokeTest(invocation.Context, SmokeRequest{Selection: input.selection, Request: seed}, input.selected, state)
		if acquired.Value == nil && operationErr == nil {
			operationErr = errors.New("PromptKit smoke test returned neither a report nor an error")
		}
		if err := promptKitContextError(invocation.Context); err != nil && !errors.Is(operationErr, err) {
			if operationErr == nil {
				operationErr = err
			} else {
				operationErr = errors.Join(operationErr, err)
			}
		}

		queuePromptDiagnostics(invocation.Reporter(), acquired.Diagnostics)
		results = append(results, smokeResult{Kind: kind, Report: acquired.Value, Err: operationErr})
		if operationErr != nil {
			operationErrors = append(operationErrors, operationErr)
			if acquired.Explanation.Message != "" || len(acquired.Explanation.Details) > 0 {
				explanations = append(explanations, acquired.Explanation)
			}
		}
		if err := promptKitContextError(invocation.Context); err != nil {
			break
		}
	}

	if len(results) == 0 {
		return definitionFailure("PromptKit aggregate smoke testing did not produce any results.", combineSmokeExplanations(explanations), errors.Join(operationErrors...))
	}
	_ = invocation.Reporter().Report(commandkit.Report{Value: smokeResultsReport(results)})
	if len(operationErrors) > 0 {
		return definitionFailure("One or more PromptKit smoke tests failed.", combineSmokeExplanations(explanations), errors.Join(operationErrors...))
	}
	return commandkit.Success()
}

func smokeFailureMessage(kind string) string {
	return fmt.Sprintf("PromptKit %s smoke testing failed.", kind)
}

func combineSmokeExplanations(explanations []commandkit.Explanation) commandkit.Explanation {
	var selected []commandkit.Explanation
	for _, explanation := range explanations {
		if explanation.Message != "" || len(explanation.Details) > 0 {
			selected = append(selected, explanation)
		}
	}
	if len(selected) == 0 {
		return commandkit.Explanation{}
	}
	if len(selected) == 1 {
		return selected[0]
	}

	combined := commandkit.Explanation{Message: selected[0].Message}
	combined.Details = append(combined.Details, selected[0].Details...)
	for _, explanation := range selected[1:] {
		if explanation.Message != "" {
			message := explanation.Message
			combined.Details = append(combined.Details, func(writer io.Writer) error {
				return writeSmokeExplanationMessage(writer, message)
			})
		}
		combined.Details = append(combined.Details, explanation.Details...)
	}
	return combined
}

func writeSmokeExplanationMessage(writer io.Writer, message string) error {
	_, err := io.WriteString(writer, message+"\n")
	return err
}

func smokeRequestForKind(seed promptkit.SmokeTestRequest, kind promptkit.SmokeTestKind) promptkit.SmokeTestRequest {
	request := seed
	request.Kind = kind
	if seed.Execution != nil {
		execution := *seed.Execution
		execution.Temperature = copyFloat64(seed.Execution.Temperature)
		execution.MaxTokens = copyInt(seed.Execution.MaxTokens)
		execution.TopP = copyFloat64(seed.Execution.TopP)
		execution.TimeoutSeconds = copyInt(seed.Execution.TimeoutSeconds)
		execution.ReasoningEffort = copyString(seed.Execution.ReasoningEffort)
		request.Execution = &execution
	}
	request.TransientRetries = copyInt(seed.TransientRetries)
	return request
}

func copyFloat64(value *float64) *float64 {
	if value == nil {
		return nil
	}
	copy := *value
	return &copy
}

func copyInt(value *int) *int {
	if value == nil {
		return nil
	}
	copy := *value
	return &copy
}

func copyString(value *string) *string {
	if value == nil {
		return nil
	}
	copy := *value
	return &copy
}