commandkit/integrations/ckprompt/smoke_command.go

package ckprompt

import (
	"errors"
	"fmt"
	"io"
	"strconv"

	commandkit "gitea.maximumdirect.net/eric/CommandKit"
	"gitea.maximumdirect.net/eric/CommandKit/integrations/ckconfig"
	promptkit "gitea.maximumdirect.net/eric/promptkit"
)

type smokeCommandInput[I any] struct {
	selection ckconfig.Selection
	kind      string
	request   promptkit.SmokeTestRequest
	format    string
	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."}, (*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."}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				(*input).selection.Overlay = value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "format", Usage: "Smoke report format: yaml."}, stringDefault("yaml"), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) { (*input).format = value.Value }),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "backend", Usage: "PromptKit backend ID."}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				(*input).request.BackendID = value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "prompt", Usage: "PromptKit prompt ID."}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				(*input).request.PromptID = value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "profile", Usage: "PromptKit profile ID."}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				(*input).request.ProfileID = value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "endpoint", Usage: "Execution endpoint override."}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				if !value.Set {
					return
				}
				execution := smokeExecution(input)
				execution.Endpoint = value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "model", Usage: "Execution model override."}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				if !value.Set {
					return
				}
				execution := smokeExecution(input)
				execution.Model = value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "temperature", Usage: "Execution temperature override."}, (*float64)(nil), parseSmokeFloat,
			func(input **smokeCommandInput[I], value commandkit.Supplied[float64]) {
				if !value.Set {
					return
				}
				execution := smokeExecution(input)
				execution.Temperature = &value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "max-tokens", Usage: "Execution maximum tokens override."}, (*int)(nil), strconv.Atoi,
			func(input **smokeCommandInput[I], value commandkit.Supplied[int]) {
				if !value.Set {
					return
				}
				execution := smokeExecution(input)
				execution.MaxTokens = &value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "top-p", Usage: "Execution top-p override."}, (*float64)(nil), parseSmokeFloat,
			func(input **smokeCommandInput[I], value commandkit.Supplied[float64]) {
				if !value.Set {
					return
				}
				execution := smokeExecution(input)
				execution.TopP = &value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "timeout-seconds", Usage: "Execution timeout override."}, (*int)(nil), strconv.Atoi,
			func(input **smokeCommandInput[I], value commandkit.Supplied[int]) {
				if !value.Set {
					return
				}
				execution := smokeExecution(input)
				execution.TimeoutSeconds = &value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "service-tier", Usage: "Execution service tier."}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				if !value.Set {
					return
				}
				execution := smokeExecution(input)
				execution.ServiceTier = value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "reasoning-effort", Usage: "Execution reasoning effort override."}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				if !value.Set {
					return
				}
				execution := smokeExecution(input)
				execution.ReasoningEffort = &value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "api-key-env", Usage: "Execution API key environment variable."}, (*string)(nil), parseDefinitionString,
			func(input **smokeCommandInput[I], value commandkit.Supplied[string]) {
				if !value.Set {
					return
				}
				execution := smokeExecution(input)
				execution.APIKeyEnv = value.Value
			}),
		commandkit.ValueFlag(commandkit.FlagMeta{Name: "transient-retries", Usage: "Inference transient retry count."}, (*int)(nil), strconv.Atoi,
			func(input **smokeCommandInput[I], value commandkit.Supplied[int]) {
				if value.Set {
					(*input).request.TransientRetries = &value.Value
				}
			}),
	)

	return commandkit.NewCommand(commandkit.CommandSpec[*smokeCommandInput[I]]{
		Name:    "test",
		Summary: "Run one PromptKit smoke test.",
		NewInput: func() *smokeCommandInput[I] {
			return &smokeCommandInput[I]{format: "yaml"}
		},
		Flags: flags,
		Arguments: []commandkit.Argument[*smokeCommandInput[I]]{
			commandkit.ValueArgument(commandkit.ArgumentMeta{Name: "kind", Usage: "Native smoke-test kind."}, 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")
			}
			if err := checkSmokeFormat(input.format); err != nil {
				return err
			}
			if input.kind == "all" {
				return errors.New("PromptKit aggregate smoke tests are not available")
			}
			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 smokeExecution[I any](input **smokeCommandInput[I]) *promptkit.ExecutionTargetOverride {
	if (*input).request.Execution == nil {
		(*input).request.Execution = &promptkit.ExecutionTargetOverride{}
	}
	return (*input).request.Execution
}

func parseSmokeFloat(value string) (float64, error) {
	return strconv.ParseFloat(value, 64)
}

func checkSmokeFormat(format string) error {
	if format != "yaml" {
		return fmt.Errorf("unsupported smoke report format %q; expected yaml", format)
	}
	return nil
}

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(err, nil)
	}
	var state S
	if suite.NewState != nil {
		state = suite.NewState()
	}
	if err := promptKitContextError(invocation.Context); err != nil {
		return definitionFailure(err, nil)
	}

	request := SmokeRequest{Selection: input.selection, Request: input.request}
	acquired, operationErr := suite.SmokeTest(invocation.Context, request, 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)
		}
	}
	if acquired.Value == nil {
		return definitionFailure(operationErr, acquired.Diagnostics)
	}

	output, encodeErr := encodeSmokeReport(acquired.Value)
	if encodeErr != nil {
		return smokeDeliveryFailure(operationErr, encodeErr, acquired.Diagnostics, false)
	}
	return writeSmokeReport(invocation, output, acquired.Diagnostics, operationErr)
}

func writeSmokeReport(invocation commandkit.Invocation, output []byte, diagnostics []commandkit.Diagnostic, operationErr error) commandkit.Outcome {
	written, deliveryErr := invocation.Streams.Out.Write(output)
	if deliveryErr == nil && written != len(output) {
		deliveryErr = io.ErrShortWrite
	}
	if deliveryErr != nil {
		return smokeDeliveryFailure(operationErr, deliveryErr, diagnostics, true)
	}
	if operationErr != nil {
		return definitionFailure(operationErr, diagnostics)
	}
	outcome := commandkit.Success()
	outcome.Diagnostics = diagnostics
	return outcome
}

func smokeDeliveryFailure(operationErr, deliveryErr error, diagnostics []commandkit.Diagnostic, reportStarted bool) commandkit.Outcome {
	cause := operationErr
	if cause == nil {
		cause = deliveryErr
	}
	outcome := definitionFailure(cause, diagnostics)
	outcome.DeliveryError = deliveryErr
	outcome.ReportIncomplete = reportStarted
	return outcome
}