276 lines
7.1 KiB
Go
276 lines
7.1 KiB
Go
package llm
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"io"
|
|
"io/fs"
|
|
"path"
|
|
"strings"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/readonlyfs"
|
|
"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 := make(map[string][]byte)
|
|
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] = data
|
|
return nil
|
|
}); err != nil {
|
|
return nil, fmt.Errorf("walk asset root %s: %w", source.Root, err)
|
|
}
|
|
}
|
|
return readonlyfs.New(out)
|
|
}
|
|
|
|
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
|
|
}
|