Add file download functionality to Neo API
- 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.
This commit is contained in:
Max
parent
3fbcd4a5e1
commit
b34f97ac2a
5 changed files with 175 additions and 9 deletions
47
neo/api.go
47
neo/api.go
|
|
@ -2,7 +2,9 @@ package neo
|
|||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/url"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
|
@ -28,6 +30,7 @@ func (neo *DSL) API(router *gin.Engine, path string) error {
|
|||
router.OPTIONS(path+"/chats", neo.optionsHandler)
|
||||
router.OPTIONS(path+"/history", neo.optionsHandler)
|
||||
router.OPTIONS(path+"/upload", neo.optionsHandler)
|
||||
router.OPTIONS(path+"/download", neo.optionsHandler)
|
||||
|
||||
// Register endpoints with middlewares
|
||||
router.GET(path, append(middlewares, neo.handleChat)...)
|
||||
|
|
@ -36,6 +39,7 @@ func (neo *DSL) API(router *gin.Engine, path string) error {
|
|||
router.GET(path+"/chats", append(middlewares, neo.handleChatList)...)
|
||||
router.GET(path+"/history", append(middlewares, neo.handleChatHistory)...)
|
||||
router.POST(path+"/upload", append(middlewares, neo.handleUpload)...)
|
||||
router.GET(path+"/download", append(middlewares, neo.handleDownload)...)
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -135,6 +139,49 @@ func (neo *DSL) handleChatHistory(c *gin.Context) {
|
|||
c.Done()
|
||||
}
|
||||
|
||||
// handleDownload handles the download request
|
||||
func (neo *DSL) handleDownload(c *gin.Context) {
|
||||
sid := c.GetString("__sid")
|
||||
if sid == "" {
|
||||
c.JSON(400, gin.H{"message": "sid is required", "code": 400})
|
||||
c.Done()
|
||||
return
|
||||
}
|
||||
|
||||
fileID := c.Query("file_id")
|
||||
if fileID == "" {
|
||||
c.JSON(400, gin.H{"message": "file_id is required", "code": 400})
|
||||
c.Done()
|
||||
return
|
||||
}
|
||||
|
||||
// Set the context
|
||||
ctx, cancel := NewContextWithCancel(sid, c.Query("chat_id"), "")
|
||||
defer cancel()
|
||||
|
||||
// Download the file
|
||||
fileResponse, err := neo.Download(ctx, c)
|
||||
if err != nil {
|
||||
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
|
||||
c.Done()
|
||||
return
|
||||
}
|
||||
defer fileResponse.Reader.Close()
|
||||
|
||||
// Set response headers
|
||||
c.Header("Content-Type", fileResponse.ContentType)
|
||||
if disposition := c.Query("disposition"); disposition == "attachment" {
|
||||
c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=%s", filepath.Base(fileID)+fileResponse.Extension))
|
||||
}
|
||||
|
||||
// Copy the file content to response
|
||||
_, err = io.Copy(c.Writer, fileResponse.Reader)
|
||||
if err != nil {
|
||||
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// getCorsHandlers returns CORS middleware handlers
|
||||
func (neo *DSL) getCorsHandlers() ([]gin.HandlerFunc, error) {
|
||||
if len(neo.Allows) == 0 {
|
||||
|
|
|
|||
|
|
@ -49,7 +49,7 @@ func (ast *Base) Upload(ctx context.Context, file *multipart.FileHeader, reader
|
|||
}
|
||||
|
||||
ext := filepath.Ext(file.Filename)
|
||||
id, err := ast.id(file.Filename)
|
||||
id, err := ast.id(file.Filename, ext)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
|
@ -61,7 +61,7 @@ func (ast *Base) Upload(ctx context.Context, file *multipart.FileHeader, reader
|
|||
}
|
||||
|
||||
return &assistant.File{
|
||||
ID: strings.ReplaceAll(id, "/", "_"),
|
||||
ID: filename,
|
||||
Filename: filename,
|
||||
ContentType: contentType,
|
||||
Bytes: int(file.Size),
|
||||
|
|
@ -69,10 +69,10 @@ func (ast *Base) Upload(ctx context.Context, file *multipart.FileHeader, reader
|
|||
}, nil
|
||||
}
|
||||
|
||||
func (ast *Base) id(temp string) (string, error) {
|
||||
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", ast.ID, date, hash), nil
|
||||
return fmt.Sprintf("/__assistants/%s/%s/%s%s", ast.ID, date, hash, ext), nil
|
||||
}
|
||||
|
||||
func (ast *Base) allowed(contentType string) bool {
|
||||
|
|
@ -85,3 +85,50 @@ func (ast *Base) allowed(contentType string) bool {
|
|||
}
|
||||
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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -49,19 +49,19 @@ func (ast *OpenAI) Upload(ctx context.Context, file *multipart.FileHeader, reade
|
|||
}
|
||||
|
||||
ext := filepath.Ext(file.Filename)
|
||||
id, err := ast.id(file.Filename)
|
||||
id, err := ast.id(file.Filename, ext)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
filename := fmt.Sprintf("%s%s", id, ext)
|
||||
filename := id
|
||||
_, err = data.Write(filename, reader, 0644)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &assistant.File{
|
||||
ID: strings.ReplaceAll(id, "/", "_"),
|
||||
ID: filename,
|
||||
Filename: filename,
|
||||
ContentType: contentType,
|
||||
Bytes: int(file.Size),
|
||||
|
|
@ -69,10 +69,10 @@ func (ast *OpenAI) Upload(ctx context.Context, file *multipart.FileHeader, reade
|
|||
}, nil
|
||||
}
|
||||
|
||||
func (ast *OpenAI) id(temp string) (string, error) {
|
||||
func (ast *OpenAI) 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", ast.ID, date, hash), nil
|
||||
return fmt.Sprintf("/__assistants/%s/%s/%s%s", ast.ID, date, hash, ext), nil
|
||||
}
|
||||
|
||||
func (ast *OpenAI) allowed(contentType string) bool {
|
||||
|
|
@ -97,3 +97,43 @@ func (ast *OpenAI) FileContent() {}
|
|||
|
||||
// FileInfo get the information of a file
|
||||
func (ast *OpenAI) FileInfo() {}
|
||||
|
||||
// Download downloads a file
|
||||
func (ast *OpenAI) 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
|
||||
}
|
||||
|
||||
return &assistant.FileResponse{
|
||||
Reader: reader,
|
||||
ContentType: contentType,
|
||||
Extension: ext,
|
||||
}, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import (
|
|||
type API interface {
|
||||
Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error
|
||||
Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error)
|
||||
Download(ctx context.Context, fileID string) (*FileResponse, error)
|
||||
}
|
||||
|
||||
// Prompt a prompt
|
||||
|
|
@ -46,3 +47,10 @@ type File struct {
|
|||
Filename string `json:"filename"`
|
||||
ContentType string `json:"content_type"`
|
||||
}
|
||||
|
||||
// FileResponse represents a file download response
|
||||
type FileResponse struct {
|
||||
Reader io.ReadCloser
|
||||
ContentType string
|
||||
Extension string
|
||||
}
|
||||
|
|
|
|||
24
neo/neo.go
24
neo/neo.go
|
|
@ -94,6 +94,30 @@ func (neo *DSL) Upload(ctx Context, c *gin.Context) (*assistant.File, error) {
|
|||
return ast.Upload(ctx, tmpfile, reader, option)
|
||||
}
|
||||
|
||||
// Download downloads a file
|
||||
func (neo *DSL) Download(ctx Context, c *gin.Context) (*assistant.FileResponse, error) {
|
||||
// Get file_id from query string
|
||||
fileID := c.Query("file_id")
|
||||
if fileID == "" {
|
||||
return nil, fmt.Errorf("file_id is required")
|
||||
}
|
||||
|
||||
// Get assistant_id from context or query
|
||||
res, err := neo.HookCreate(ctx, []map[string]interface{}{}, c)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Select Assistant
|
||||
ast, err := neo.selectAssistant(res.AssistantID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Download file using the assistant
|
||||
return ast.Download(ctx.Context, fileID)
|
||||
}
|
||||
|
||||
// chat chat with AI
|
||||
func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]interface{}, c *gin.Context) error {
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue