diff --git a/neo/api.go b/neo/api.go index ab687947..5e1b837c 100644 --- a/neo/api.go +++ b/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 { diff --git a/neo/assistant/base/file.go b/neo/assistant/base/file.go index 8b7738f4..fd4f3029 100644 --- a/neo/assistant/base/file.go +++ b/neo/assistant/base/file.go @@ -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 +} diff --git a/neo/assistant/openai/file.go b/neo/assistant/openai/file.go index 7405d7eb..44892d49 100644 --- a/neo/assistant/openai/file.go +++ b/neo/assistant/openai/file.go @@ -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 +} diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 3e338cb0..03930582 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -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 +} diff --git a/neo/neo.go b/neo/neo.go index 6a87be8a..50ca4c3b 100644 --- a/neo/neo.go +++ b/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 {