431 lines
11 KiB
Go
431 lines
11 KiB
Go
package llm
|
|
|
|
import (
|
|
"bytes"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"io"
|
|
"io/fs"
|
|
"path"
|
|
"sort"
|
|
"strings"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit"
|
|
)
|
|
|
|
type AssetSource struct {
|
|
FS fs.FS
|
|
Root string
|
|
}
|
|
|
|
type AssetRegistry struct {
|
|
prompts []AssetSource
|
|
schemas []AssetSource
|
|
fallbackProfiles []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
|
|
}
|
|
|
|
// RegisterFallbackProfileFS registers profile assets that PromptKit uses only
|
|
// when an operator-configured source does not provide a matching profile.
|
|
func (r *AssetRegistry) RegisterFallbackProfileFS(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 fallback profile assets: %w", err)
|
|
}
|
|
r.fallbackProfiles = append(r.fallbackProfiles, 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) FallbackProfileFS() (fs.FS, error) {
|
|
if r == nil {
|
|
return nil, fmt.Errorf("asset registry must not be nil")
|
|
}
|
|
return flattenAssetSources(r.fallbackProfiles)
|
|
}
|
|
|
|
// FallbackProfileDigest returns a deterministic, non-secret identity for the
|
|
// flattened fallback profile assets.
|
|
func (r *AssetRegistry) FallbackProfileDigest() (string, error) {
|
|
_, digest, _, err := r.fallbackProfileAssets()
|
|
return digest, err
|
|
}
|
|
|
|
func (r *AssetRegistry) PromptKitOptions() ([]promptkit.Option, error) {
|
|
options, _, err := r.promptKitOptions()
|
|
return options, err
|
|
}
|
|
|
|
func (r *AssetRegistry) promptKitOptions() ([]promptkit.Option, string, 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)
|
|
}
|
|
options := []promptkit.Option{
|
|
promptkit.WithPromptFS(promptFS, "."),
|
|
promptkit.WithSchemaFS(schemaFS, "."),
|
|
}
|
|
fallbackFS, fallbackDigest, hasFallback, err := r.fallbackProfileAssets()
|
|
if err != nil {
|
|
return nil, "", err
|
|
}
|
|
if hasFallback {
|
|
options = append(options, promptkit.WithFallbackProfileFS(fallbackFS, "."))
|
|
}
|
|
return options, fallbackDigest, nil
|
|
}
|
|
|
|
func (r *AssetRegistry) promptKitFallbackProfileOption() (promptkit.Option, bool, error) {
|
|
fallbackFS, _, hasFallback, err := r.fallbackProfileAssets()
|
|
if err != nil {
|
|
return nil, false, err
|
|
}
|
|
if !hasFallback {
|
|
return nil, false, nil
|
|
}
|
|
return promptkit.WithFallbackProfileFS(fallbackFS, "."), true, nil
|
|
}
|
|
|
|
func (r *AssetRegistry) fallbackProfileAssets() (fs.FS, string, bool, error) {
|
|
if r == nil {
|
|
return nil, "", false, fmt.Errorf("asset registry must not be nil")
|
|
}
|
|
if len(r.fallbackProfiles) == 0 {
|
|
empty := sha256.Sum256([]byte("notarius:fallback-profile-assets:empty"))
|
|
return nil, "sha256:" + hex.EncodeToString(empty[:]), false, nil
|
|
}
|
|
fallbackFS, err := r.FallbackProfileFS()
|
|
if err != nil {
|
|
return nil, "", false, fmt.Errorf("prepare fallback profile assets: %w", err)
|
|
}
|
|
digest, err := hashAssetFilesystem(fallbackFS)
|
|
if err != nil {
|
|
return nil, "", false, err
|
|
}
|
|
return fallbackFS, digest, true, 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 hashAssetFilesystem(fsys fs.FS) (string, error) {
|
|
var parts []AssetHashPart
|
|
err := fs.WalkDir(fsys, ".", func(name string, entry fs.DirEntry, walkErr error) error {
|
|
if walkErr != nil {
|
|
return walkErr
|
|
}
|
|
if entry.IsDir() {
|
|
return nil
|
|
}
|
|
parts = append(parts, AssetHashPart{FS: fsys, Path: name})
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return "", fmt.Errorf("walk assets for digest: %w", err)
|
|
}
|
|
if len(parts) > 0 {
|
|
return HashAssets(parts)
|
|
}
|
|
empty := sha256.Sum256([]byte("notarius:fallback-profile-assets:empty"))
|
|
return "sha256:" + hex.EncodeToString(empty[:]), 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 {
|
|
if dir == "." {
|
|
return []fs.DirEntry{}
|
|
}
|
|
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 }
|