promptkit/internal/artifact/reader.go

package artifact

import (
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"mime"
	"os"
	"path/filepath"

	"gitea.maximumdirect.net/eric/promptkit/internal/defaults"
	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
)

var (
	ErrUnsupportedRefType = errors.New("unsupported artifact reference type")
	ErrMissingFilePath    = errors.New("missing file path for file artifact")
	ErrUnsupportedFile    = errors.New("file artifact path is not a regular file")
)

const fileReadChunkSize = 64 * 1024

// Reader resolves artifact references into actual artifacts.
type Reader interface {
	Read(ctx context.Context, ref domain.ArtifactRef) (*domain.Artifact, error)
}

// CompositeReader routes artifact resolution based on the reference type.
type CompositeReader struct {
	inlineReader *inlineReader
	fileReader   Reader
}

func NewCompositeReader() Reader {
	return &CompositeReader{
		inlineReader: &inlineReader{},
		fileReader:   &fileReader{open: openArtifactFile},
	}
}

func (c *CompositeReader) Read(ctx context.Context, ref domain.ArtifactRef) (*domain.Artifact, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}

	switch ref.Type {
	case domain.ArtifactRefInline:
		return c.inlineReader.Read(ctx, ref)
	case domain.ArtifactRefFile:
		return c.fileReader.Read(ctx, ref)
	default:
		return nil, fmt.Errorf("%w: %s", ErrUnsupportedRefType, ref.Type)
	}
}

type inlineReader struct{}

func (r *inlineReader) Read(ctx context.Context, ref domain.ArtifactRef) (*domain.Artifact, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}

	body := []byte(ref.Body)
	return &domain.Artifact{
		ContentType: defaults.ContentTypeTextPlain,
		Body:        body,
		Size:        int64(len(body)),
		Hash:        fmt.Sprintf("%x", sha256.Sum256(body)),
		URI:         ref.URI,
	}, nil
}

type artifactFile interface {
	Read([]byte) (int, error)
	Stat() (os.FileInfo, error)
	Close() error
}

type fileReader struct {
	open func(string) (artifactFile, error)
}

func (r *fileReader) Read(ctx context.Context, ref domain.ArtifactRef) (*domain.Artifact, error) {
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	default:
	}

	if ref.URI == "" {
		return nil, ErrMissingFilePath
	}

	return readFileArtifact(ctx, ref.URI, r.open)
}

func openArtifactFile(path string) (artifactFile, error) {
	return os.Open(path)
}

func readFileArtifact(ctx context.Context, path string, open func(string) (artifactFile, error)) (*domain.Artifact, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	info, err := os.Stat(path)
	if err != nil {
		return nil, fmt.Errorf("failed to read file %s: %w", path, err)
	}
	if !info.Mode().IsRegular() {
		return nil, fmt.Errorf("%w: %s", ErrUnsupportedFile, path)
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	file, err := open(path)
	if err != nil {
		return nil, fmt.Errorf("failed to read file %s: %w", path, err)
	}
	defer file.Close()

	openedInfo, err := file.Stat()
	if err != nil {
		return nil, fmt.Errorf("failed to inspect opened file %s: %w", path, err)
	}
	if !openedInfo.Mode().IsRegular() {
		return nil, fmt.Errorf("%w: %s", ErrUnsupportedFile, path)
	}

	data := make([]byte, 0)
	chunk := make([]byte, fileReadChunkSize)
	for {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		n, readErr := file.Read(chunk)
		if n > 0 {
			data = append(data, chunk[:n]...)
		}
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		if errors.Is(readErr, io.EOF) {
			break
		}
		if readErr != nil {
			return nil, fmt.Errorf("failed to read file %s: %w", path, readErr)
		}
	}

	contentType := mime.TypeByExtension(filepath.Ext(path))
	if contentType == "" {
		contentType = defaults.ContentTypeTextPlain
	}
	hash := fmt.Sprintf("%x", sha256.Sum256(data))
	if err := ctx.Err(); err != nil {
		return nil, err
	}

	return &domain.Artifact{
		Name:        filepath.Base(path),
		ContentType: contentType,
		Body:        data,
		URI:         path,
		Size:        int64(len(data)),
		Hash:        hash,
	}, nil
}