b34f97ac2a
- Introduced a new endpoint for downloading files, enhancing the API's capabilities. - Implemented the handleDownload method to manage download requests, including validation for required parameters (sid and file_id). - Updated the Download method in the assistant interface to retrieve files based on file_id, ensuring proper error handling and response structure. - Enhanced file upload logic to include file extensions in generated IDs for better file management. - Added necessary response headers for file downloads, improving user experience during file retrieval.
134 lines
3.5 KiB
Go
134 lines
3.5 KiB
Go
package base
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"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 *Base) 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 *Base) 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 *Base) 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 *Base) 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
|
|
}
|