notarius/internal/framework/prompt/registry.go

package prompt

import (
	"crypto/sha256"
	"embed"
	"encoding/hex"
	"fmt"
	"path"
	"sort"
	"strings"
	"text/template"
)

//go:embed assets/**
var embeddedAssets embed.FS

const (
	SourceBuiltin       = "builtin"
	VersionV1           = "v1"
	TestGenericPromptID = "test.generic"
)

// Metadata describes a registered prompt asset.
type Metadata struct {
	PromptID      string `json:"prompt_id"`
	PromptVersion string `json:"prompt_version"`
	PromptSource  string `json:"prompt_source"`
	EmbeddedPath  string `json:"embedded_path"`
	SHA256        string `json:"sha256"`
}

// DiagnosticsMap returns prompt metadata without rendered prompt text.
func (m Metadata) DiagnosticsMap() map[string]any {
	return map[string]any{
		"prompt_id":      m.PromptID,
		"prompt_version": m.PromptVersion,
		"prompt_source":  m.PromptSource,
		"embedded_path":  m.EmbeddedPath,
		"sha256":         m.SHA256,
	}
}

type definition struct {
	id          string
	version     string
	embeddedDir string
	systemPath  string
	userPath    string
}

type compiledPrompt struct {
	systemTmpl *template.Template
	userTmpl   *template.Template
	metadata   Metadata
}

var promptRegistry map[string]compiledPrompt
var sharedHardening string

func init() {
	var err error
	sharedHardening, err = readAsset("assets/shared/prompt_hardening.md")
	if err != nil {
		panic(err)
	}

	defs := []definition{
		{
			id:          TestGenericPromptID,
			version:     VersionV1,
			embeddedDir: "assets/test/generic",
			systemPath:  "assets/test/generic/system.md",
			userPath:    "assets/test/generic/user.md",
		},
	}

	promptRegistry = make(map[string]compiledPrompt, len(defs))
	for _, def := range defs {
		compiled, compileErr := compilePrompt(def)
		if compileErr != nil {
			panic(compileErr)
		}
		promptRegistry[def.id] = compiled
	}
}

// LookupMetadata returns metadata for the requested prompt ID.
func LookupMetadata(promptID string) (Metadata, bool) {
	compiled, ok := promptRegistry[strings.TrimSpace(promptID)]
	if !ok {
		return Metadata{}, false
	}
	return compiled.metadata, true
}

// MustLookupMetadata returns metadata for the requested prompt ID and panics when missing.
func MustLookupMetadata(promptID string) Metadata {
	metadata, ok := LookupMetadata(promptID)
	if !ok {
		panic(fmt.Sprintf("unknown prompt id %q", promptID))
	}
	return metadata
}

// RegisteredMetadata returns all prompt metadata sorted by prompt ID.
func RegisteredMetadata() []Metadata {
	ids := make([]string, 0, len(promptRegistry))
	for id := range promptRegistry {
		ids = append(ids, id)
	}
	sort.Strings(ids)

	out := make([]Metadata, 0, len(ids))
	for _, id := range ids {
		out = append(out, promptRegistry[id].metadata)
	}
	return out
}

// HardeningText returns the shared hardening instructions available to templates.
func HardeningText() string {
	return sharedHardening
}

func readAsset(assetPath string) (string, error) {
	content, err := embeddedAssets.ReadFile(assetPath)
	if err != nil {
		return "", fmt.Errorf("read embedded prompt asset %q: %w", assetPath, err)
	}
	return string(content), nil
}

func compilePrompt(def definition) (compiledPrompt, error) {
	if strings.TrimSpace(def.id) == "" {
		return compiledPrompt{}, fmt.Errorf("prompt id must not be empty")
	}
	if strings.TrimSpace(def.version) == "" {
		return compiledPrompt{}, fmt.Errorf("prompt version must not be empty")
	}
	if strings.TrimSpace(def.embeddedDir) == "" {
		return compiledPrompt{}, fmt.Errorf("prompt embedded path must not be empty")
	}

	systemSource, err := readAsset(def.systemPath)
	if err != nil {
		return compiledPrompt{}, err
	}
	userSource, err := readAsset(def.userPath)
	if err != nil {
		return compiledPrompt{}, err
	}

	funcs := template.FuncMap{
		"hardening": func() string { return sharedHardening },
	}
	systemTmpl, err := template.New(path.Base(def.systemPath)).Option("missingkey=error").Funcs(funcs).Parse(systemSource)
	if err != nil {
		return compiledPrompt{}, fmt.Errorf("parse embedded system prompt %q: %w", def.systemPath, err)
	}
	userTmpl, err := template.New(path.Base(def.userPath)).Option("missingkey=error").Funcs(funcs).Parse(userSource)
	if err != nil {
		return compiledPrompt{}, fmt.Errorf("parse embedded user prompt %q: %w", def.userPath, err)
	}

	hashInput := systemSource + "\n\n" + userSource
	hash := sha256.Sum256([]byte(hashInput))
	metadata := Metadata{
		PromptID:      strings.TrimSpace(def.id),
		PromptVersion: strings.TrimSpace(def.version),
		PromptSource:  SourceBuiltin,
		EmbeddedPath:  strings.TrimSpace(def.embeddedDir),
		SHA256:        "sha256:" + hex.EncodeToString(hash[:]),
	}

	return compiledPrompt{
		systemTmpl: systemTmpl,
		userTmpl:   userTmpl,
		metadata:   metadata,
	}, nil
}