commandkit/integrations/ckprompt/examples/composed/main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"strconv"

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

type campaign struct {
	PromptID      string `yaml:"prompt_id"`
	PromptVersion string `yaml:"prompt_version"`
	ProfileID     string `yaml:"profile_id"`
	Model         string `yaml:"model"`
}

type appConfiguration struct {
	DefaultCampaign string              `yaml:"default_campaign"`
	Campaigns       map[string]campaign `yaml:"campaigns"`
}

type inspectionInput struct {
	Campaign commandkit.Supplied[string]
	Limit    commandkit.Supplied[int]
}

type applicationCacheKey struct {
	selection ckconfig.Selection
	campaign  commandkit.Supplied[string]
}

// runCache belongs to one invocation and is shared only where the application
// chooses to pass it to its loaders.
type runCache struct {
	applications map[applicationCacheKey]*configkit.Result[appConfiguration]
	llmClient    promptkit.LLMClient
}

type exampleLLMClient struct {
	calls *int
}

func (client *exampleLLMClient) Generate(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
	if client.calls != nil {
		(*client.calls)++
	}
	return &promptkit.GenerateResponse{Content: "Example response."}, nil
}

type restoreInput struct {
	Campaign commandkit.Supplied[string]
	Session  string
}

type session struct {
	ID       string
	Campaign string
	Summary  string
}

func campaignSelector() commandkit.Flag[*inspectionInput] {
	return commandkit.ValueFlag(commandkit.FlagMeta{Name: "campaign", Usage: "Select campaign inputs."}, (*string)(nil), parseString,
		func(input **inspectionInput, value commandkit.Supplied[string]) { (*input).Campaign = value })
}

func loadApplicationConfig(ctx context.Context, selection ckconfig.Selection, input *inspectionInput, cache *runCache) (*configkit.Result[appConfiguration], error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	campaignInput := commandkit.Supplied[string]{}
	if input != nil {
		campaignInput = input.Campaign
	}
	key := applicationCacheKey{selection: selection, campaign: campaignInput}
	if cache != nil {
		if cached, ok := cache.applications[key]; ok {
			return cached, nil
		}
	}
	source := configkit.Mapping("composed application settings", map[string]any{
		"default_campaign": "spring",
		"campaigns": map[string]any{
			"spring": map[string]any{"prompt_id": "briefing", "prompt_version": "1", "profile_id": "spring", "model": "spring-model"},
			"autumn": map[string]any{"prompt_id": "briefing", "prompt_version": "1", "profile_id": "autumn", "model": "autumn-model"},
		},
	})
	if selection.Config.Set {
		source = configkit.File(selection.Config.Value)
	}
	options := configkit.Options[appConfiguration]{}
	if selection.Overlay.Set {
		options.Overlay = selection.Overlay.Value
		options.Filesystem = &configkit.FilesystemContext{Root: ".", BaseDir: "."}
	}
	if campaignInput.Set {
		options.Overrides = []configkit.Override{{Path: configkit.Keys("default_campaign"), Value: campaignInput.Value, Flag: "--campaign"}}
	}
	result, err := configkit.Load(source, options)
	if err == nil && result != nil && cache != nil {
		if cache.applications == nil {
			cache.applications = make(map[applicationCacheKey]*configkit.Result[appConfiguration])
		}
		cache.applications[key] = result
	}
	return result, err
}

func loadPromptEngine(ctx context.Context, selection ckconfig.Selection, input *inspectionInput, cache *runCache) (*promptkit.Engine, error) {
	application, err := loadApplicationConfig(ctx, selection, input, cache)
	if err != nil {
		return nil, err
	}
	name, chosen, err := selectedCampaign(application.Config, input.Campaign)
	if err != nil {
		return nil, err
	}
	engineConfig := configkit.Mapping("campaign "+name, map[string]any{
		"collections": map[string]any{
			"field-notes": map[string]any{
				"release_version": "1",
				"backends":        map[string]any{"demo": map[string]any{"endpoint": "https://example.invalid/v1"}},
				"profiles":        map[string]any{"assistant": map[string]any{"backend": "field-notes/demo", "model": "example-model"}},
				"prompts": map[string]any{"briefing": map[string]any{
					"version":         "1",
					"default_profile": "field-notes/assistant",
					"inputs":          []any{map[string]any{"name": "topic", "required": true, "content_type": "text/plain", "description": "Topic to brief"}},
					"messages":        []any{map[string]any{"role": "user", "content": "Write a concise briefing about {{.topic}}."}},
					"output":          map[string]any{"format": "text", "validation_mode": "none"},
				}},
			},
		},
		"profiles": map[string]any{chosen.ProfileID: map[string]any{
			"endpoint": "https://example.invalid/v1",
			"model":    chosen.Model,
			"extra_params": map[string]any{
				"campaign": name,
			},
		}},
		"prompts": map[string]any{chosen.PromptID: map[string]any{
			"version":         chosen.PromptVersion,
			"default_profile": chosen.ProfileID,
			"inputs":          []any{map[string]any{"name": "topic", "required": true, "content_type": "text/plain", "description": "Topic to brief"}},
			"messages":        []any{map[string]any{"role": "user", "content": "Write a concise briefing about {{.topic}}. {{if .mode}}{{input \"notes\"}}{{end}}"}},
			"output":          map[string]any{"format": "text", "validation_mode": "none"},
		}},
	})
	var client promptkit.LLMClient = &exampleLLMClient{}
	if cache != nil && cache.llmClient != nil {
		client = cache.llmClient
	}
	return promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{Main: &configkit.Input{Source: engineConfig}}}, promptkit.WithLLMClient(client))
}

func selectedCampaign(configuration appConfiguration, supplied commandkit.Supplied[string]) (string, campaign, error) {
	name := configuration.DefaultCampaign
	if supplied.Set {
		name = supplied.Value
	}
	chosen, ok := configuration.Campaigns[name]
	if !ok {
		return "", campaign{}, fmt.Errorf("unknown campaign %q", name)
	}
	return name, chosen, nil
}

func buildValidationRequest(ctx context.Context, seed ckprompt.ValidationRequest, input *inspectionInput, cache *runCache) (promptkit.RunRequest, error) {
	application, err := loadApplicationConfig(ctx, seed.Selection, input, cache)
	if err != nil {
		return promptkit.RunRequest{}, err
	}
	name, _, err := selectedCampaign(application.Config, input.Campaign)
	if err != nil {
		return promptkit.RunRequest{}, err
	}

	request := seed.Request
	vars := make(map[string]string, len(request.Vars)+2)
	for key, value := range request.Vars {
		vars[key] = value
	}
	if _, supplied := vars["topic"]; !supplied {
		vars["topic"] = name + " campaign planning"
	}
	if _, supplied := vars["mode"]; !supplied {
		vars["mode"] = ""
	}
	request.Vars = vars
	return request, nil
}

func newApp() (*commandkit.App, error) {
	return newAppWithLLMClient(&exampleLLMClient{})
}

func newAppWithLLMClient(client promptkit.LLMClient) (*commandkit.App, error) {
	return newAppWithEngineLoader(client, loadPromptEngine)
}

func newAppWithEngineLoader(client promptkit.LLMClient, engineLoader func(context.Context, ckconfig.Selection, *inspectionInput, *runCache) (*promptkit.Engine, error)) (*commandkit.App, error) {
	if client == nil {
		client = &exampleLLMClient{}
	}
	limitDefault := 0
	limitFlag := commandkit.ValueFlag(commandkit.FlagMeta{Name: "limit", Usage: "Override the application item limit."}, &limitDefault, strconv.Atoi,
		func(input **inspectionInput, value commandkit.Supplied[int]) { (*input).Limit = value })
	inputFactory := func() *inspectionInput { return &inspectionInput{} }
	stateFactory := func() *runCache { return &runCache{llmClient: client} }

	appScope := ckconfig.FromConfigKit("app", "Application settings.", loadApplicationConfig)
	engineScope := ckprompt.ConfigScope(engineLoader)
	configCommand := ckconfig.Command(ckconfig.Suite[*inspectionInput, *runCache]{
		NewInput: inputFactory,
		NewState: stateFactory,
		CheckInput: func(scope string, _ ckconfig.Selection, input *inspectionInput) error {
			if scope == "app" && input.Campaign.Set && !input.Limit.Set {
				return errors.New("application campaign inspection requires --limit")
			}
			return nil
		},
		Selectors: []ckconfig.Selector[*inspectionInput]{
			{Binding: campaignSelector(), Scopes: []string{"app", "promptkit"}},
			{Binding: limitFlag, Scopes: []string{"app"}},
		},
		Scopes: []ckconfig.Scope[*inspectionInput, *runCache]{appScope, engineScope},
	})

	definitions := ckprompt.FromEngine(engineLoader)
	definitions.NewInput = inputFactory
	definitions.NewState = stateFactory
	definitions.Selectors = []commandkit.Flag[*inspectionInput]{campaignSelector()}
	definitions = ckprompt.WithRequestValidation(definitions, engineLoader, buildValidationRequest)
	definitions = ckprompt.WithSmokeTesting(definitions, engineLoader, nil)
	definitionCommand := ckprompt.Command(definitions)

	store := map[string]session{
		"session-42": {ID: "session-42", Campaign: "spring", Summary: "A synthetic planning session."},
	}
	restore := commandkit.NewCommand(commandkit.CommandSpec[restoreInput]{
		Name:    "restore",
		Summary: "Restore a saved session.",
		Flags: []commandkit.Flag[restoreInput]{commandkit.ValueFlag(commandkit.FlagMeta{Name: "campaign", Usage: "Campaign owning the session."}, (*string)(nil), parseString,
			func(input *restoreInput, value commandkit.Supplied[string]) { input.Campaign = value })},
		Arguments: []commandkit.Argument[restoreInput]{commandkit.ValueArgument(commandkit.ArgumentMeta{Name: "ID", Usage: "Session identifier."}, true, parseString,
			func(input *restoreInput, value commandkit.Supplied[string]) { input.Session = value.Value })},
		Run: func(invocation commandkit.Invocation, input restoreInput) commandkit.Outcome {
			if err := invocation.Context.Err(); err != nil {
				return commandkit.Failure(err)
			}
			application, err := loadApplicationConfig(invocation.Context, ckconfig.Selection{}, nil, nil)
			if err != nil {
				return commandkit.Failure(err)
			}
			name, _, err := selectedCampaign(application.Config, input.Campaign)
			if err != nil {
				return commandkit.Failure(err)
			}
			record, ok := store[input.Session]
			if !ok {
				return commandkit.Outcome{Kind: commandkit.OutcomeDomain, Status: 4, Domain: "session_not_found", Cause: errors.New("saved session was not found")}
			}
			if record.Campaign != name {
				return commandkit.Outcome{Kind: commandkit.OutcomeDomain, Status: 3, Domain: "campaign_mismatch", Cause: fmt.Errorf("session belongs to campaign %q", record.Campaign)}
			}
			if _, err := fmt.Fprintf(invocation.Streams.Out, "restored %s: %s\n", record.ID, record.Summary); err != nil {
				return commandkit.Failure(err)
			}
			return commandkit.Success()
		},
	})
	declarations := []commandkit.Declaration{
		configCommand,
		definitionCommand,
		commandkit.Group("session", "Manage saved sessions.", restore),
		commandkit.Alias("list collections", "promptkit list collections"),
		commandkit.Alias("promptkit list bundles", "promptkit list collections"),
		commandkit.Alias("pk", "promptkit"),
		commandkit.Alias("smoke", "promptkit test"),
		commandkit.FlagAliasFor("promptkit show prompt", "--fmt", "--format"),
		commandkit.FlagAlias("--hlep", "--help"),
	}
	return commandkit.New(commandkit.Identity{Name: "composed-example", Version: "1.0"}, declarations...)
}

func parseString(raw string) (string, error) { return raw, nil }

func main() {
	app, err := newApp()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	outcome := app.Execute(context.Background(), os.Args[1:], commandkit.Streams{Out: os.Stdout, Err: os.Stderr})
	if outcome.Status != 0 {
		os.Exit(outcome.Status)
	}
}