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" {
				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("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().RenderOutput(func(writer io.Writer) error {
		output, err := encodeSmokeReport(acquired.Value)
		if err != nil {
			return err
		}
		_, err = writer.Write(output)
		return err
	})
	if operationErr != nil {
		return definitionFailure(smokeFailureMessage(input.kind), acquired.Explanation, operationErr)
	}
	return commandkit.Success()
}

func runSmokeAggregate[I, S any](invocation commandkit.Invocation, suite Suite[I, S], input *smokeCommandInput[I], state S) commandkit.Outcome {
	kinds := [...]promptkit.SmokeTestKind{
		promptkit.SmokeTestReachability,
		promptkit.SmokeTestAuthentication,
		promptkit.SmokeTestTextInference,
		promptkit.SmokeTestStructuredOutput,
	}
	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().RenderOutput(func(writer io.Writer) error {
		output, err := encodeSmokeResults(results)
		if err != nil {
			return err
		}
		_, err = writer.Write(output)
		return err
	})
	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
}