yao/neo/assistant/local/file.go

package local

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"io"
	"mime/multipart"
	"path/filepath"
	"strings"
	"time"

	"github.com/yaoapp/gou/fs"
	"github.com/yaoapp/yao/neo/assistant"
)

// AllowedFileTypes the allowed file types
var AllowedFileTypes = map[string]string{
	"application/json":   "json",
	"application/pdf":    "pdf",
	"application/msword": "doc",
	"application/vnd.openxmlformats-officedocument.wordprocessingml.document":   "docx",
	"application/vnd.oasis.opendocument.text":                                   "odt",
	"application/vnd.ms-excel":                                                  "xls",
	"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":         "xlsx",
	"application/vnd.ms-powerpoint":                                             "ppt",
	"application/vnd.openxmlformats-officedocument.presentationml.presentation": "pptx",
}

// MaxSize 20M max file size
var MaxSize int64 = 20 * 1024 * 1024

// Upload the file
func (ast *Local) Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*assistant.File, error) {

	// check file size
	if file.Size > MaxSize {
		return nil, fmt.Errorf("file size %d exceeds the maximum size of %d", file.Size, MaxSize)
	}

	contentType := file.Header.Get("Content-Type")
	if !ast.allowed(contentType) {
		return nil, fmt.Errorf("file type %s not allowed", contentType)
	}

	data, err := fs.Get("data")
	if err != nil {
		return nil, err
	}

	ext := filepath.Ext(file.Filename)
	id, err := ast.id(file.Filename, ext)
	if err != nil {
		return nil, err
	}

	filename := fmt.Sprintf("%s%s", id, ext)
	_, err = data.Write(filename, reader, 0644)
	if err != nil {
		return nil, err
	}

	return &assistant.File{
		ID:          filename,
		Filename:    filename,
		ContentType: contentType,
		Bytes:       int(file.Size),
		CreatedAt:   int(time.Now().Unix()),
	}, nil
}

func (ast *Local) id(temp string, ext string) (string, error) {
	date := time.Now().Format("20060102")
	hash := fmt.Sprintf("%x", sha256.Sum256([]byte(temp)))[:8]
	return fmt.Sprintf("/__assistants/%s/%s/%s%s", ast.ID, date, hash, ext), nil
}

func (ast *Local) allowed(contentType string) bool {
	if _, ok := AllowedFileTypes[contentType]; ok {
		return true
	}
	// text/* // image/* // audio/* // video/*
	if strings.HasPrefix(contentType, "text/") || strings.HasPrefix(contentType, "image/") || strings.HasPrefix(contentType, "audio/") || strings.HasPrefix(contentType, "video/") {
		return true
	}
	return false
}

// Download downloads a file
func (ast *Local) Download(ctx context.Context, fileID string) (*assistant.FileResponse, error) {

	// Get the data filesystem
	data, err := fs.Get("data")
	if err != nil {
		return nil, fmt.Errorf("get filesystem error: %s", err.Error())
	}

	// Check if file exists
	exists, err := data.Exists(fileID)
	if err != nil {
		return nil, fmt.Errorf("check file error: %s", err.Error())
	}
	if !exists {
		return nil, fmt.Errorf("file %s not found", fileID)
	}

	// Open the file
	reader, err := data.ReadCloser(fileID)
	if err != nil {
		return nil, err
	}

	// Get content type and extension
	ext := filepath.Ext(fileID)

	// Get content type from mime type
	contentType := "application/octet-stream"
	if v, err := data.MimeType(fileID); err == nil {
		contentType = v
	}

	for mimeType, extension := range AllowedFileTypes {
		if "."+extension == ext {
			contentType = mimeType
			break
		}
	}

	return &assistant.FileResponse{
		Reader:      reader,
		ContentType: contentType,
		Extension:   ext,
	}, nil
}

// ReadBase64 reads a file and returns its base64 encoded content
func (ast *Local) ReadBase64(ctx context.Context, fileID string) (string, error) {
	// Get the data filesystem
	data, err := fs.Get("data")
	if err != nil {
		return "", fmt.Errorf("get filesystem error: %s", err.Error())
	}

	// Check if file exists
	exists, err := data.Exists(fileID)
	if err != nil {
		return "", fmt.Errorf("check file error: %s", err.Error())
	}
	if !exists {
		return "", fmt.Errorf("file %s not found", fileID)
	}

	// Read file content
	content, err := data.ReadFile(fileID)
	if err != nil {
		return "", fmt.Errorf("read file error: %s", err.Error())
	}

	// Encode to base64
	return base64.StdEncoding.EncodeToString(content), nil
}