notarius/internal/framework/prompt/registry.go

package prompt

import (
	"crypto/sha256"
	"embed"
	"encoding/hex"
	"fmt"
	"io/fs"
	"path"
	"reflect"
	"sort"
	"strings"
	"text/template"
	"text/template/parse"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

//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
	ReferenceSlots []contracts.ReferenceSlot
}

// Bundle is a compiled system/user prompt pair.
type Bundle struct {
	systemTmpl     *template.Template
	userTmpl       *template.Template
	metadata       Metadata
	referenceSlots map[string]contracts.ReferenceSlot
}

// 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 },
		"reference":    func(string) (string, error) { return "", nil },
		"hasreference": func(string) (bool, error) { return false, nil },
	}
	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)
	}
	referenceSlots := referenceSlotMap(def.ReferenceSlots)
	if err := validateTemplateReferenceSlots(systemTmpl, referenceSlots); err != nil {
		return nil, fmt.Errorf("validate embedded system prompt %q: %w", systemPath, err)
	}
	if err := validateTemplateReferenceSlots(userTmpl, referenceSlots); err != nil {
		return nil, fmt.Errorf("validate 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,
		referenceSlots: referenceSlots,
	}, nil
}

func referenceSlotMap(slots []contracts.ReferenceSlot) map[string]contracts.ReferenceSlot {
	if len(slots) == 0 {
		return nil
	}
	out := make(map[string]contracts.ReferenceSlot, len(slots))
	for _, slot := range slots {
		name := strings.TrimSpace(slot.Name)
		if name == "" {
			continue
		}
		slot.Name = name
		slot.AcceptedMediaTypes = append([]string(nil), slot.AcceptedMediaTypes...)
		out[name] = slot
	}
	return out
}

func validateTemplateReferenceSlots(tmpl *template.Template, declared map[string]contracts.ReferenceSlot) error {
	if tmpl == nil || tmpl.Tree == nil || tmpl.Tree.Root == nil {
		return nil
	}
	return validateReferenceNodes(tmpl.Tree.Root, declared)
}

func validateReferenceNodes(node parse.Node, declared map[string]contracts.ReferenceSlot) error {
	if node == nil || reflect.ValueOf(node).IsNil() {
		return nil
	}
	switch typed := node.(type) {
	case *parse.ListNode:
		for _, child := range typed.Nodes {
			if err := validateReferenceNodes(child, declared); err != nil {
				return err
			}
		}
	case *parse.ActionNode:
		return validateReferencePipeline(typed.Pipe, declared)
	case *parse.IfNode:
		if err := validateReferencePipeline(typed.Pipe, declared); err != nil {
			return err
		}
		if err := validateReferenceNodes(typed.List, declared); err != nil {
			return err
		}
		return validateReferenceNodes(typed.ElseList, declared)
	case *parse.RangeNode:
		if err := validateReferencePipeline(typed.Pipe, declared); err != nil {
			return err
		}
		if err := validateReferenceNodes(typed.List, declared); err != nil {
			return err
		}
		return validateReferenceNodes(typed.ElseList, declared)
	case *parse.WithNode:
		if err := validateReferencePipeline(typed.Pipe, declared); err != nil {
			return err
		}
		if err := validateReferenceNodes(typed.List, declared); err != nil {
			return err
		}
		return validateReferenceNodes(typed.ElseList, declared)
	case *parse.TemplateNode:
		return nil
	}
	return nil
}

func validateReferencePipeline(pipe *parse.PipeNode, declared map[string]contracts.ReferenceSlot) error {
	if pipe == nil {
		return nil
	}
	for _, cmd := range pipe.Cmds {
		if err := validateReferenceCommand(cmd, declared); err != nil {
			return err
		}
	}
	return nil
}

func validateReferenceCommand(cmd *parse.CommandNode, declared map[string]contracts.ReferenceSlot) error {
	if cmd == nil || len(cmd.Args) == 0 {
		return nil
	}
	for _, arg := range cmd.Args[1:] {
		if nested, ok := arg.(*parse.PipeNode); ok {
			if err := validateReferencePipeline(nested, declared); err != nil {
				return err
			}
		}
	}
	identifier, ok := cmd.Args[0].(*parse.IdentifierNode)
	if !ok {
		return nil
	}
	if identifier.Ident != "reference" && identifier.Ident != "hasreference" {
		return nil
	}
	if len(cmd.Args) != 2 {
		return fmt.Errorf("%s requires one string slot name", identifier.Ident)
	}
	slotArg, ok := cmd.Args[1].(*parse.StringNode)
	if !ok {
		return fmt.Errorf("%s requires a string literal slot name", identifier.Ident)
	}
	slotName := strings.TrimSpace(slotArg.Text)
	if slotName == "" {
		return fmt.Errorf("%s slot name must not be empty", identifier.Ident)
	}
	if _, ok := declared[slotName]; !ok {
		return fmt.Errorf("%s slot %q is not declared", identifier.Ident, slotName)
	}
	return 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
}