Add Scriptorium prompt assets

2026-07-05 18:13:10 +00:00
parent 0fc740470f
commit de6689bc1d
21 changed files with 1123 additions and 8 deletions
--- a/internal/framework/llm/asset_registry.go
+++ b/internal/framework/llm/asset_registry.go
@@ -0,0 +1,333 @@
+package llm
+
+import (
+	"bytes"
+	"crypto/sha256"
+	"encoding/hex"
+	"fmt"
+	"io"
+	"io/fs"
+	"path"
+	"sort"
+	"strings"
+	"time"
+
+	"gitea.maximumdirect.net/eric/scriptorium"
+)
+
+type AssetSource struct {
+	FS   fs.FS
+	Root string
+}
+
+type AssetRegistry struct {
+	prompts []AssetSource
+	schemas []AssetSource
+}
+
+type AssetHashPart struct {
+	FS   fs.FS
+	Path string
+}
+
+func NewAssetRegistry() *AssetRegistry {
+	return &AssetRegistry{}
+}
+
+func (r *AssetRegistry) RegisterPromptFS(fsys fs.FS, root string) error {
+	if r == nil {
+		return fmt.Errorf("asset registry must not be nil")
+	}
+	source, err := newAssetSource(fsys, root)
+	if err != nil {
+		return fmt.Errorf("register prompt assets: %w", err)
+	}
+	r.prompts = append(r.prompts, source)
+	return nil
+}
+
+func (r *AssetRegistry) RegisterSchemaFS(fsys fs.FS, root string) error {
+	if r == nil {
+		return fmt.Errorf("asset registry must not be nil")
+	}
+	source, err := newAssetSource(fsys, root)
+	if err != nil {
+		return fmt.Errorf("register schema assets: %w", err)
+	}
+	r.schemas = append(r.schemas, source)
+	return nil
+}
+
+func (r *AssetRegistry) PromptFS() (fs.FS, error) {
+	if r == nil {
+		return nil, fmt.Errorf("asset registry must not be nil")
+	}
+	return flattenAssetSources(r.prompts)
+}
+
+func (r *AssetRegistry) SchemaFS() (fs.FS, error) {
+	if r == nil {
+		return nil, fmt.Errorf("asset registry must not be nil")
+	}
+	return flattenAssetSources(r.schemas)
+}
+
+func (r *AssetRegistry) ScriptoriumOptions() ([]scriptorium.Option, error) {
+	promptFS, err := r.PromptFS()
+	if err != nil {
+		return nil, fmt.Errorf("prepare prompt assets: %w", err)
+	}
+	schemaFS, err := r.SchemaFS()
+	if err != nil {
+		return nil, fmt.Errorf("prepare schema assets: %w", err)
+	}
+	return []scriptorium.Option{
+		scriptorium.WithPromptFS(promptFS, "."),
+		scriptorium.WithSchemaFS(schemaFS, "."),
+	}, nil
+}
+
+func HashAssets(parts []AssetHashPart) (string, error) {
+	if len(parts) == 0 {
+		return "", fmt.Errorf("asset hash requires at least one part")
+	}
+	hash := sha256.New()
+	for _, part := range parts {
+		cleanPath, err := cleanAssetPath(part.Path)
+		if err != nil {
+			return "", fmt.Errorf("hash asset %q: %w", part.Path, err)
+		}
+		data, err := fs.ReadFile(part.FS, cleanPath)
+		if err != nil {
+			return "", fmt.Errorf("read hash asset %s: %w", cleanPath, err)
+		}
+		if _, err := io.WriteString(hash, cleanPath); err != nil {
+			return "", err
+		}
+		if _, err := hash.Write([]byte{0}); err != nil {
+			return "", err
+		}
+		if _, err := hash.Write(data); err != nil {
+			return "", err
+		}
+		if _, err := hash.Write([]byte{0}); err != nil {
+			return "", err
+		}
+	}
+	return "sha256:" + hex.EncodeToString(hash.Sum(nil)), nil
+}
+
+func newAssetSource(fsys fs.FS, root string) (AssetSource, error) {
+	if fsys == nil {
+		return AssetSource{}, fmt.Errorf("filesystem must not be nil")
+	}
+	cleanRoot, err := cleanAssetRoot(root)
+	if err != nil {
+		return AssetSource{}, err
+	}
+	return AssetSource{FS: fsys, Root: cleanRoot}, nil
+}
+
+func flattenAssetSources(sources []AssetSource) (fs.FS, error) {
+	out := assetMapFS{}
+	for _, source := range sources {
+		if err := fs.WalkDir(source.FS, source.Root, func(name string, entry fs.DirEntry, walkErr error) error {
+			if walkErr != nil {
+				return walkErr
+			}
+			if entry.IsDir() {
+				return nil
+			}
+			rel := name
+			if source.Root != "." {
+				rel = strings.TrimPrefix(name, source.Root+"/")
+			}
+			rel, err := cleanAssetPath(rel)
+			if err != nil {
+				return err
+			}
+			if _, exists := out[rel]; exists {
+				return fmt.Errorf("duplicate asset path %q", rel)
+			}
+			data, err := fs.ReadFile(source.FS, name)
+			if err != nil {
+				return err
+			}
+			out[rel] = append([]byte(nil), data...)
+			return nil
+		}); err != nil {
+			return nil, fmt.Errorf("walk asset root %s: %w", source.Root, err)
+		}
+	}
+	return out, nil
+}
+
+func cleanAssetRoot(root string) (string, error) {
+	trimmed := strings.TrimSpace(root)
+	if trimmed == "" || trimmed == "." {
+		return ".", nil
+	}
+	return cleanAssetPath(trimmed)
+}
+
+func cleanAssetPath(name string) (string, error) {
+	trimmed := strings.TrimSpace(name)
+	if trimmed == "" {
+		return "", fmt.Errorf("path must not be empty")
+	}
+	cleaned := path.Clean(strings.TrimPrefix(trimmed, "/"))
+	if cleaned == "." || !fs.ValidPath(cleaned) {
+		return "", fmt.Errorf("invalid path %q", name)
+	}
+	return cleaned, nil
+}
+
+type assetMapFS map[string][]byte
+
+func (m assetMapFS) Open(name string) (fs.File, error) {
+	cleaned, err := cleanOpenPath(name)
+	if err != nil {
+		return nil, &fs.PathError{Op: "open", Path: name, Err: err}
+	}
+	if data, ok := m[cleaned]; ok {
+		return &assetFile{
+			reader: bytes.NewReader(data),
+			info:   assetFileInfo{name: path.Base(cleaned), size: int64(len(data))},
+		}, nil
+	}
+	entries := m.dirEntries(cleaned)
+	if entries != nil {
+		return &assetDir{name: path.Base(cleaned), entries: entries}, nil
+	}
+	return nil, &fs.PathError{Op: "open", Path: name, Err: fs.ErrNotExist}
+}
+
+func (m assetMapFS) ReadFile(name string) ([]byte, error) {
+	cleaned, err := cleanOpenPath(name)
+	if err != nil {
+		return nil, &fs.PathError{Op: "readfile", Path: name, Err: err}
+	}
+	data, ok := m[cleaned]
+	if !ok {
+		return nil, &fs.PathError{Op: "readfile", Path: name, Err: fs.ErrNotExist}
+	}
+	return append([]byte(nil), data...), nil
+}
+
+func (m assetMapFS) ReadDir(name string) ([]fs.DirEntry, error) {
+	cleaned, err := cleanOpenPath(name)
+	if err != nil {
+		return nil, &fs.PathError{Op: "readdir", Path: name, Err: err}
+	}
+	entries := m.dirEntries(cleaned)
+	if entries == nil {
+		return nil, &fs.PathError{Op: "readdir", Path: name, Err: fs.ErrNotExist}
+	}
+	return entries, nil
+}
+
+func (m assetMapFS) dirEntries(dir string) []fs.DirEntry {
+	children := map[string]assetDirEntry{}
+	prefix := ""
+	if dir != "." {
+		prefix = dir + "/"
+	}
+	for name, data := range m {
+		if !strings.HasPrefix(name, prefix) {
+			continue
+		}
+		rest := strings.TrimPrefix(name, prefix)
+		if rest == "" {
+			continue
+		}
+		childName, _, hasSlash := strings.Cut(rest, "/")
+		entry := assetDirEntry{name: childName, dir: hasSlash}
+		if !hasSlash {
+			entry.size = int64(len(data))
+		}
+		children[childName] = entry
+	}
+	if len(children) == 0 {
+		return nil
+	}
+	names := make([]string, 0, len(children))
+	for name := range children {
+		names = append(names, name)
+	}
+	sort.Strings(names)
+	entries := make([]fs.DirEntry, 0, len(names))
+	for _, name := range names {
+		entries = append(entries, children[name])
+	}
+	return entries
+}
+
+func cleanOpenPath(name string) (string, error) {
+	if name == "." {
+		return ".", nil
+	}
+	return cleanAssetPath(name)
+}
+
+type assetFile struct {
+	reader *bytes.Reader
+	info   assetFileInfo
+}
+
+func (f *assetFile) Stat() (fs.FileInfo, error) { return f.info, nil }
+func (f *assetFile) Read(p []byte) (int, error) { return f.reader.Read(p) }
+func (f *assetFile) Close() error               { return nil }
+
+type assetDir struct {
+	name    string
+	offset  int
+	entries []fs.DirEntry
+}
+
+func (d *assetDir) Stat() (fs.FileInfo, error) { return assetFileInfo{name: d.name, dir: true}, nil }
+func (d *assetDir) Read([]byte) (int, error)   { return 0, fmt.Errorf("cannot read directory") }
+func (d *assetDir) Close() error               { return nil }
+func (d *assetDir) ReadDir(n int) ([]fs.DirEntry, error) {
+	if d.offset >= len(d.entries) {
+		return nil, io.EOF
+	}
+	end := len(d.entries)
+	if n > 0 && d.offset+n < end {
+		end = d.offset + n
+	}
+	out := append([]fs.DirEntry(nil), d.entries[d.offset:end]...)
+	d.offset = end
+	return out, nil
+}
+
+type assetDirEntry struct {
+	name string
+	dir  bool
+	size int64
+}
+
+func (e assetDirEntry) Name() string      { return e.name }
+func (e assetDirEntry) IsDir() bool       { return e.dir }
+func (e assetDirEntry) Type() fs.FileMode { return e.InfoMode().Type() }
+func (e assetDirEntry) Info() (fs.FileInfo, error) {
+	return assetFileInfo{name: e.name, dir: e.dir, size: e.size}, nil
+}
+func (e assetDirEntry) InfoMode() fs.FileMode {
+	if e.dir {
+		return fs.ModeDir | 0o555
+	}
+	return 0o444
+}
+
+type assetFileInfo struct {
+	name string
+	dir  bool
+	size int64
+}
+
+func (i assetFileInfo) Name() string       { return i.name }
+func (i assetFileInfo) Size() int64        { return i.size }
+func (i assetFileInfo) Mode() fs.FileMode  { return assetDirEntry{dir: i.dir}.InfoMode() }
+func (i assetFileInfo) ModTime() time.Time { return time.Time{} }
+func (i assetFileInfo) IsDir() bool        { return i.dir }
+func (i assetFileInfo) Sys() any           { return nil }
--- a/internal/framework/llm/asset_registry_test.go
+++ b/internal/framework/llm/asset_registry_test.go
@@ -0,0 +1,167 @@
+package llm
+
+import (
+	"context"
+	"strings"
+	"testing"
+	"testing/fstest"
+	"time"
+
+	"gitea.maximumdirect.net/eric/scriptorium"
+)
+
+func TestAssetRegistryCombinesPromptAndSchemaSources(t *testing.T) {
+	registry := NewAssetRegistry()
+	mustRegisterPromptFS(t, registry, fstest.MapFS{
+		"prompts/test.yaml":           {Data: []byte(validPromptYAML("schemas/out.json"))},
+		"prompts/messages/user.tmpl":  {Data: []byte(`Input: {{ input "transcript" }}`)},
+		"prompts/messages/task.tmpl":  {Data: []byte("Return JSON.")},
+		"schemas/ignored/schema.json": {Data: []byte(`{"type":"object"}`)},
+	}, "prompts")
+	mustRegisterSchemaFS(t, registry, fstest.MapFS{
+		"root/schemas/out.json": {Data: []byte(`{"type":"object"}`)},
+	}, "root")
+
+	engine := newAssetTestEngine(t, registry)
+	prepared, err := engine.Prepare(context.Background(), scriptorium.RunRequest{
+		PromptID:  "asset.test",
+		ProfileID: "asset-test-profile",
+		Inputs: map[string]scriptorium.ArtifactRef{
+			"transcript": scriptorium.Inline(`{"ok":true}`),
+		},
+	})
+	if err != nil {
+		t.Fatalf("Prepare() error = %v, want nil", err)
+	}
+	if got := len(prepared.Messages); got != 2 {
+		t.Fatalf("message count = %d, want 2", got)
+	}
+	if prepared.OutputContract.SchemaPath != "schemas/out.json" {
+		t.Fatalf("schema path = %q, want schemas/out.json", prepared.OutputContract.SchemaPath)
+	}
+}
+
+func TestAssetRegistryPrepareFailsForMissingPromptAsset(t *testing.T) {
+	registry := NewAssetRegistry()
+	mustRegisterPromptFS(t, registry, fstest.MapFS{
+		"test.yaml": {Data: []byte(validPromptYAML("out.json"))},
+	}, ".")
+	mustRegisterSchemaFS(t, registry, fstest.MapFS{
+		"out.json": {Data: []byte(`{"type":"object"}`)},
+	}, ".")
+
+	engine := newAssetTestEngine(t, registry)
+	_, err := engine.Prepare(context.Background(), scriptorium.RunRequest{
+		PromptID:  "asset.test",
+		ProfileID: "asset-test-profile",
+		Inputs: map[string]scriptorium.ArtifactRef{
+			"transcript": scriptorium.Inline(`{"ok":true}`),
+		},
+	})
+	if err == nil || !strings.Contains(err.Error(), "content_file") {
+		t.Fatalf("Prepare() error = %v, want missing content_file error", err)
+	}
+}
+
+func TestAssetRegistryPrepareFailsForMissingSchemaAsset(t *testing.T) {
+	registry := NewAssetRegistry()
+	mustRegisterPromptFS(t, registry, fstest.MapFS{
+		"test.yaml":          {Data: []byte(validPromptYAML("missing.json"))},
+		"messages/user.tmpl": {Data: []byte(`Input: {{ input "transcript" }}`)},
+		"messages/task.tmpl": {Data: []byte("Return JSON.")},
+	}, ".")
+	mustRegisterSchemaFS(t, registry, fstest.MapFS{
+		"present.json": {Data: []byte(`{"type":"object"}`)},
+	}, ".")
+
+	engine := newAssetTestEngine(t, registry)
+	_, err := engine.Prepare(context.Background(), scriptorium.RunRequest{
+		PromptID:  "asset.test",
+		ProfileID: "asset-test-profile",
+		Inputs: map[string]scriptorium.ArtifactRef{
+			"transcript": scriptorium.Inline(`{"ok":true}`),
+		},
+	})
+	if err == nil || !strings.Contains(err.Error(), "missing.json") {
+		t.Fatalf("Prepare() error = %v, want missing schema error", err)
+	}
+}
+
+func TestAssetRegistryRejectsDuplicateAssetPaths(t *testing.T) {
+	registry := NewAssetRegistry()
+	mustRegisterPromptFS(t, registry, fstest.MapFS{"one/prompt.yaml": {Data: []byte("id: one")}}, "one")
+	mustRegisterPromptFS(t, registry, fstest.MapFS{"two/prompt.yaml": {Data: []byte("id: two")}}, "two")
+
+	_, err := registry.PromptFS()
+	if err == nil || !strings.Contains(err.Error(), "duplicate asset path") {
+		t.Fatalf("PromptFS() error = %v, want duplicate path error", err)
+	}
+}
+
+func TestHashAssetsOmitsRawAssetContent(t *testing.T) {
+	hash, err := HashAssets([]AssetHashPart{{
+		FS:   fstest.MapFS{"prompt.md": {Data: []byte("secret prompt text")}},
+		Path: "prompt.md",
+	}})
+	if err != nil {
+		t.Fatalf("HashAssets() error = %v, want nil", err)
+	}
+	if !strings.HasPrefix(hash, "sha256:") {
+		t.Fatalf("hash = %q, want sha256-prefixed value", hash)
+	}
+	if strings.Contains(hash, "secret prompt text") {
+		t.Fatalf("hash leaked asset content")
+	}
+}
+
+func newAssetTestEngine(t *testing.T, registry *AssetRegistry) *scriptorium.Engine {
+	t.Helper()
+	options, err := registry.ScriptoriumOptions()
+	if err != nil {
+		t.Fatalf("ScriptoriumOptions() error = %v, want nil", err)
+	}
+	options = append(options, scriptorium.WithProfiles(scriptorium.OpenAICompatibleProfile(scriptorium.OpenAICompatibleProfileConfig{
+		ID:       "asset-test-profile",
+		Endpoint: "http://127.0.0.1:1/v1",
+		Model:    "asset-test-model",
+	})))
+	engine, err := scriptorium.NewEngine(scriptorium.Config{Timeout: time.Second}, options...)
+	if err != nil {
+		t.Fatalf("NewEngine() error = %v, want nil", err)
+	}
+	return engine
+}
+
+func mustRegisterPromptFS(t *testing.T, registry *AssetRegistry, fsys fstest.MapFS, root string) {
+	t.Helper()
+	if err := registry.RegisterPromptFS(fsys, root); err != nil {
+		t.Fatalf("RegisterPromptFS() error = %v, want nil", err)
+	}
+}
+
+func mustRegisterSchemaFS(t *testing.T, registry *AssetRegistry, fsys fstest.MapFS, root string) {
+	t.Helper()
+	if err := registry.RegisterSchemaFS(fsys, root); err != nil {
+		t.Fatalf("RegisterSchemaFS() error = %v, want nil", err)
+	}
+}
+
+func validPromptYAML(schemaPath string) string {
+	return `id: asset.test
+version: "v1"
+inputs:
+  - name: transcript
+    required: true
+    content_type: application/json
+messages:
+  - role: user
+    content_file: ./messages/user.tmpl
+  - role: user
+    content_file: ./messages/task.tmpl
+output:
+  format: json
+  validation_mode: json_schema
+  schema_path: ` + schemaPath + `
+  repair_attempts: 0
+`
+}