promptkit/examples/go-library/prepare/main.go

package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"

	"go.maximumdirect.net/configkit"
	"go.maximumdirect.net/promptkit"
)

type summary struct {
	PromptID        string              `json:"prompt_id"`
	SelectedProfile string              `json:"selected_profile"`
	Model           string              `json:"model"`
	MessageCount    int                 `json:"message_count"`
	ProfileTimeout  int                 `json:"profile_timeout_seconds"`
	Namespaces      []string            `json:"namespaces"`
	Collections     []collectionSummary `json:"collections"`
	ScopedPrompts   []string            `json:"scoped_prompts"`
	ReviewPrompts   []string            `json:"review_prompts"`
}

type collectionSummary struct {
	Namespace      string  `json:"namespace"`
	ReleaseVersion *string `json:"release_version"`
}

func main() {
	if err := run(context.Background(), "examples/go-library/prepare", os.Stdout); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func run(ctx context.Context, exampleDir string, out io.Writer) error {
	support, err := promptkit.CollectionInput(
		"support",
		configkit.File(filepath.Join(exampleDir, "collections", "support", "collection.yaml")),
	)
	if err != nil {
		return err
	}
	review, err := promptkit.CollectionInput(
		"review",
		configkit.File(filepath.Join(exampleDir, "collections", "review", "collection.yaml")),
	)
	if err != nil {
		return err
	}
	override, err := promptkit.CollectionInput(
		"support",
		configkit.File(filepath.Join(exampleDir, "collections", "support", "operator-overrides.yaml")),
	)
	if err != nil {
		return err
	}
	engine, err := promptkit.NewEngine(promptkit.Config{Sources: promptkit.ConfigurationSources{
		Underlays: []configkit.Underlay{
			{Inputs: []configkit.Input{support, review}},
			{Inputs: []configkit.Input{override}},
		},
	}})
	if err != nil {
		return err
	}
	scoped, err := engine.WithNamespaces(promptkit.Namespaces{
		Prompt: "support", Profile: "support", Backend: "support",
	})
	if err != nil {
		return err
	}
	view, err := engine.InspectConfiguration()
	if err != nil {
		return err
	}
	if _, found := view.Explain(configkit.Keys("collections", "support", "profiles").Index(1).Key("timeout_seconds")); !found {
		return fmt.Errorf("collection profile override was not found")
	}
	requirements, err := scoped.InspectRequirements(ctx, promptkit.RequirementInspectionRequest{
		PromptID: "prepare",
	})
	if err != nil {
		return err
	}
	if !requirements.Complete || len(requirements.Inputs) != 1 ||
		requirements.Inputs[0].Name != "note" ||
		requirements.Inputs[0].Kind != promptkit.RequirementRequired {
		return fmt.Errorf("required note input was not discovered before preparation")
	}

	request := promptkit.RunRequest{
		PromptID:  "prepare",
		ProfileID: "fast",
		Inputs: map[string]promptkit.ArtifactRef{
			"note": promptkit.Inline("Ada finished the migration review."),
		},
	}
	if err := scoped.ValidateRequest(ctx, request); err != nil {
		return err
	}
	prepared, err := scoped.Prepare(ctx, request)
	if err != nil {
		return err
	}
	profile, err := scoped.InspectProfile(ctx, "fast")
	if err != nil {
		return err
	}
	if profile.EffectiveModelParams.TimeoutSeconds != 30 {
		return fmt.Errorf(
			"later profile override was not applied: timeout is %d seconds",
			profile.EffectiveModelParams.TimeoutSeconds,
		)
	}
	namespaces, err := scoped.ListNamespaces(ctx)
	if err != nil {
		return err
	}
	collections, err := scoped.ListCollections(ctx)
	if err != nil {
		return err
	}
	scopedPrompts, err := scoped.ListPrompts(ctx)
	if err != nil {
		return err
	}
	reviewPrompts, err := scoped.ListPrompts(ctx, promptkit.NamespaceFilter{
		Mode: promptkit.NamespaceNamed, Name: "review",
	})
	if err != nil {
		return err
	}
	scopedPromptIDs := make([]string, len(scopedPrompts))
	for index := range scopedPrompts {
		scopedPromptIDs[index] = scopedPrompts[index].ID
	}
	reviewPromptIDs := make([]string, len(reviewPrompts))
	for index := range reviewPrompts {
		reviewPromptIDs[index] = reviewPrompts[index].ID
	}
	collectionSummaries := make([]collectionSummary, len(collections))
	for index := range collections {
		collectionSummaries[index] = collectionSummary{
			Namespace: collections[index].Namespace, ReleaseVersion: collections[index].ReleaseVersion,
		}
	}

	encoder := json.NewEncoder(out)
	encoder.SetIndent("", "  ")
	return encoder.Encode(summary{
		PromptID:        prepared.PromptID,
		SelectedProfile: prepared.SelectedProfileID,
		Model:           prepared.EffectiveModelParams.Model,
		MessageCount:    len(prepared.Messages),
		ProfileTimeout:  profile.EffectiveModelParams.TimeoutSeconds,
		Namespaces:      namespaces,
		Collections:     collectionSummaries,
		ScopedPrompts:   scopedPromptIDs,
		ReviewPrompts:   reviewPromptIDs,
	})
}