promptkit/internal/profile/filesystem_repository.go

package profile

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"os"
	"strings"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
	"gitea.maximumdirect.net/eric/promptkit/internal/filecatalog"
	"gitea.maximumdirect.net/eric/promptkit/internal/jsonvalue"
	"gopkg.in/yaml.v3"
)

var (
	ErrProfileNotFound     = errors.New("execution profile not found")
	ErrInvalidYAML         = errors.New("invalid YAML format")
	ErrInvalidProfile      = errors.New("invalid execution profile configuration")
	ErrRawAPIKeyNotAllowed = errors.New("raw api_key is not allowed; use api_key_env")
)

type filesystemRepository struct {
	dir string
}

func NewFilesystemRepository(dir string) Repository {
	return &filesystemRepository{dir: dir}
}

func (r *filesystemRepository) GetProfile(ctx context.Context, id string) (*domain.ExecutionProfile, error) {
	return loadProfile(ctx, os.DirFS(r.dir), ".", id)
}

type fsRepository struct {
	fsys fs.FS
	root string
}

func NewFSRepository(fsys fs.FS, root string) Repository {
	return &fsRepository{fsys: fsys, root: root}
}

func (r *fsRepository) GetProfile(ctx context.Context, id string) (*domain.ExecutionProfile, error) {
	return loadProfile(ctx, r.fsys, r.root, id)
}

type overlayRepository struct {
	primary  Repository
	fallback Repository
}

func NewOverlayRepository(primary, fallback Repository) Repository {
	return &overlayRepository{primary: primary, fallback: fallback}
}

func (r *overlayRepository) GetProfile(ctx context.Context, id string) (*domain.ExecutionProfile, error) {
	if r.primary != nil {
		prof, err := r.primary.GetProfile(ctx, id)
		if err == nil {
			return prof, nil
		}
		if !errors.Is(err, ErrProfileNotFound) {
			return nil, err
		}
	}
	if r.fallback == nil {
		return nil, ErrProfileNotFound
	}
	return r.fallback.GetProfile(ctx, id)
}

func loadProfile(ctx context.Context, fsys fs.FS, root string, id string) (*domain.ExecutionProfile, error) {
	id = strings.TrimSpace(id)
	if id == "" {
		return nil, fmt.Errorf("%w: profile id is required", ErrInvalidProfile)
	}
	if fsys == nil {
		return nil, fmt.Errorf("failed to read profile directory: filesystem is nil")
	}

	files, err := filecatalog.FindFSYAMLFiles(ctx, fsys, root)
	if err != nil {
		return nil, fmt.Errorf("failed to read profile directory: %w", err)
	}

	var matches []profileMatch
	for _, fullPath := range files {
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}

		relPath := filecatalog.DisplayPath(root, fullPath)
		data, err := fs.ReadFile(fsys, fullPath)
		if err != nil {
			return nil, fmt.Errorf("failed to read profile file %s: %w", relPath, err)
		}
		metadata, metadataErr := readProfileFileMetadata(data)
		idMatch := metadata.matchesID(id)
		if metadataErr != nil {
			if idMatch {
				return nil, fmt.Errorf("%w: %s: %v", ErrInvalidYAML, relPath, metadataErr)
			}
			continue
		}
		if metadata.hasRawAPIKey {
			if idMatch {
				return nil, fmt.Errorf("%w: %s", ErrRawAPIKeyNotAllowed, relPath)
			}
			continue
		}

		prof, err := decodeProfile(data)
		if err != nil {
			if idMatch {
				return nil, fmt.Errorf("%w: %s: %v", ErrInvalidYAML, relPath, err)
			}
			continue
		}

		prof.ID = strings.TrimSpace(prof.ID)
		if prof.ID != id {
			continue
		}
		prof.BackendID = strings.TrimSpace(prof.BackendID)
		prof.ExtraParams, err = jsonvalue.CopyMap(prof.ExtraParams)
		if err != nil {
			return nil, fmt.Errorf("%w: %s: %v", ErrInvalidProfile, relPath, err)
		}
		if err := validateProfile(prof); err != nil {
			if errors.Is(err, ErrRawAPIKeyNotAllowed) {
				return nil, fmt.Errorf("%w: %s", err, relPath)
			}
			return nil, fmt.Errorf("%w: %s: %v", ErrInvalidProfile, relPath, err)
		}
		matches = append(matches, profileMatch{
			profile: prof,
			path:    relPath,
		})
	}

	if len(matches) > 1 {
		paths := make([]string, 0, len(matches))
		for _, match := range matches {
			paths = append(paths, match.path)
		}
		return nil, fmt.Errorf("%w: duplicate execution profile id %q found in: %s", ErrInvalidProfile, id, strings.Join(paths, ", "))
	}

	if len(matches) == 1 {
		return matches[0].profile, nil
	}

	return nil, ErrProfileNotFound
}

type profileMatch struct {
	profile *domain.ExecutionProfile
	path    string
}

type profileFileMetadata struct {
	ids          []string
	hasRawAPIKey bool
}

func readProfileFileMetadata(data []byte) (profileFileMetadata, error) {
	decoder := yaml.NewDecoder(bytes.NewReader(data))
	var node yaml.Node
	if err := decoder.Decode(&node); err != nil {
		return profileFileMetadata{}, err
	}
	metadata := profileMetadataFromNode(&node)
	documentCount := 1
	for {
		var trailing yaml.Node
		err := decoder.Decode(&trailing)
		if errors.Is(err, io.EOF) {
			if documentCount == 1 {
				return metadata, nil
			}
			return metadata, errors.New("profile file must contain exactly one YAML document")
		}
		if err != nil {
			return metadata, err
		}
		documentCount++
		metadata.merge(profileMetadataFromNode(&trailing))
	}
}

func profileMetadataFromNode(node *yaml.Node) profileFileMetadata {
	if node.Kind != yaml.DocumentNode || len(node.Content) == 0 {
		return profileFileMetadata{}
	}
	mapping := node.Content[0]
	if mapping.Kind != yaml.MappingNode {
		return profileFileMetadata{}
	}

	var metadata profileFileMetadata
	for i := 0; i+1 < len(mapping.Content); i += 2 {
		key := mapping.Content[i]
		value := mapping.Content[i+1]
		switch key.Value {
		case "id":
			metadata.ids = append(metadata.ids, strings.TrimSpace(value.Value))
		case "api_key":
			metadata.hasRawAPIKey = true
		}
	}
	return metadata
}

func (m profileFileMetadata) matchesID(id string) bool {
	for _, candidate := range m.ids {
		if candidate == id {
			return true
		}
	}
	return false
}

func (m *profileFileMetadata) merge(other profileFileMetadata) {
	m.ids = append(m.ids, other.ids...)
	m.hasRawAPIKey = m.hasRawAPIKey || other.hasRawAPIKey
}

func decodeProfile(data []byte) (*domain.ExecutionProfile, error) {
	var prof domain.ExecutionProfile
	decoder := yaml.NewDecoder(bytes.NewReader(data))
	decoder.KnownFields(true)
	if err := decoder.Decode(&prof); err != nil {
		return nil, err
	}
	if err := requireYAMLStreamEnd(decoder); err != nil {
		return nil, err
	}
	return &prof, nil
}

func requireYAMLStreamEnd(decoder *yaml.Decoder) error {
	var trailing yaml.Node
	err := decoder.Decode(&trailing)
	if errors.Is(err, io.EOF) {
		return nil
	}
	if err != nil {
		return err
	}
	return errors.New("profile file must contain exactly one YAML document")
}

func validateProfile(p *domain.ExecutionProfile) error {
	if strings.TrimSpace(p.ID) == "" {
		return errors.New("id is required")
	}
	if strings.TrimSpace(p.BackendID) == "" && strings.TrimSpace(p.Endpoint) == "" {
		return errors.New("backend or endpoint is required")
	}
	if strings.TrimSpace(p.Model) == "" {
		return errors.New("model is required")
	}

	return domain.ValidateExecutionTargetSettings(domain.ExecutionTarget{
		Temperature:    p.Temperature,
		MaxTokens:      p.MaxTokens,
		TopP:           p.TopP,
		TimeoutSeconds: p.TimeoutSeconds,
	})
}