notarius/internal/framework/prompt/registry.go

package prompt

import (
	"crypto/sha256"
	"embed"
	"encoding/hex"
	"fmt"
	"io/fs"
	"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,
	}
}

// Definition identifies a caller-owned system/user prompt bundle.
type Definition struct {
	PromptID     string
	Version      string
	EmbeddedPath string
	SystemPath   string
	UserPath     string
}

// Bundle is a compiled system/user prompt pair.
type Bundle struct {
	systemTmpl *template.Template
	userTmpl   *template.Template
	metadata   Metadata
}

// Metadata returns metadata for the compiled prompt bundle.
func (b *Bundle) Metadata() Metadata {
	if b == nil {
		return Metadata{}
	}
	return b.metadata
}

var promptRegistry map[string]*Bundle
var sharedHardening string

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

	defs := []Definition{
		{
			PromptID:     TestGenericPromptID,
			Version:      VersionV1,
			EmbeddedPath: "assets/test/generic",
			SystemPath:   "assets/test/generic/system.md",
			UserPath:     "assets/test/generic/user.md",
		},
	}

	promptRegistry = make(map[string]*Bundle, len(defs))
	for _, def := range defs {
		compiled, compileErr := LoadBundle(embeddedAssets, def)
		if compileErr != nil {
			panic(compileErr)
		}
		promptRegistry[compiled.metadata.PromptID] = 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
}

// LoadBundle compiles a system/user prompt bundle from a caller-owned filesystem.
func LoadBundle(fsys fs.FS, def Definition) (*Bundle, error) {
	promptID := strings.TrimSpace(def.PromptID)
	version := strings.TrimSpace(def.Version)
	embeddedPath := strings.TrimSpace(def.EmbeddedPath)
	systemPath := strings.TrimSpace(def.SystemPath)
	userPath := strings.TrimSpace(def.UserPath)
	if promptID == "" {
		return nil, fmt.Errorf("prompt id must not be empty")
	}
	if version == "" {
		return nil, fmt.Errorf("prompt version must not be empty")
	}
	if embeddedPath == "" {
		return nil, fmt.Errorf("prompt embedded path must not be empty")
	}

	systemSource, err := readPromptAsset(fsys, systemPath)
	if err != nil {
		return nil, err
	}
	userSource, err := readPromptAsset(fsys, userPath)
	if err != nil {
		return nil, err
	}

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

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

	return &Bundle{
		systemTmpl: systemTmpl,
		userTmpl:   userTmpl,
		metadata:   metadata,
	}, nil
}

func readPromptAsset(fsys fs.FS, assetPath string) (string, error) {
	if strings.TrimSpace(assetPath) == "" {
		return "", fmt.Errorf("prompt asset path must not be empty")
	}
	content, err := fs.ReadFile(fsys, assetPath)
	if err != nil {
		return "", fmt.Errorf("read embedded prompt asset %q: %w", assetPath, err)
	}
	return string(content), nil
}