From 5b99f3bf4411d07159883adeb1b6c6b04802f9fe Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 31 May 2025 16:36:02 +0800 Subject: [PATCH] Enhance file upload functionality and authentication in neo package - Updated the API to support dynamic storage paths for file uploads, allowing for differentiated handling of chat, knowledge, and asset uploads. - Implemented comprehensive authentication checks to ensure only authorized users can upload files. - Introduced new upload options to handle original filenames and improved error responses for better user feedback. - Enhanced CORS headers to accommodate new upload requirements and improve API security. - Initialized authentication settings in the Load function to streamline user management. --- neo/api.go | 98 ++++++++++++++++++++++++++++++++++----- neo/attachment/manager.go | 58 +++++++++++++++++++++-- neo/attachment/types.go | 16 +++---- neo/load.go | 80 ++++++++++++++++++++++++++++++++ neo/neo.go | 65 ++++++++++---------------- neo/types.go | 2 +- widgets/app/app.go | 21 +++++++++ 7 files changed, 274 insertions(+), 66 deletions(-) diff --git a/neo/api.go b/neo/api.go index 04c2c70d..aae3da1f 100644 --- a/neo/api.go +++ b/neo/api.go @@ -4,6 +4,7 @@ import ( "fmt" "io" "net/url" + "os" "path/filepath" "strconv" "strings" @@ -16,6 +17,7 @@ import ( "github.com/yaoapp/gou/process" "github.com/yaoapp/yao/helper" "github.com/yaoapp/yao/neo/assistant" + "github.com/yaoapp/yao/neo/attachment" chatctx "github.com/yaoapp/yao/neo/context" "github.com/yaoapp/yao/neo/message" "github.com/yaoapp/yao/neo/store" @@ -36,7 +38,7 @@ func (neo *DSL) API(router *gin.Engine, path string) error { router.OPTIONS(path+"/chats", neo.optionsHandler) router.OPTIONS(path+"/chats/:id", neo.optionsHandler) router.OPTIONS(path+"/history", neo.optionsHandler) - router.OPTIONS(path+"/upload", neo.optionsHandler) + router.OPTIONS(path+"/upload/:storage", neo.optionsHandler) router.OPTIONS(path+"/download", neo.optionsHandler) router.OPTIONS(path+"/mentions", neo.optionsHandler) router.OPTIONS(path+"/generate", neo.optionsHandler) @@ -122,7 +124,7 @@ func (neo *DSL) API(router *gin.Engine, path string) error { // Upload file example: // curl -X POST 'http://localhost:5099/api/__yao/neo/upload?chat_id=chat_123&token=xxx' \ // -F 'file=@/path/to/file.txt' - router.POST(path+"/upload", append(middlewares, neo.handleUpload)...) + router.POST(path+"/upload/:storage", append(middlewares, neo.handleUpload)...) // Download file example: // curl -X GET 'http://localhost:5099/api/__yao/neo/download?file_id=file_123&disposition=attachment&token=xxx' \ @@ -177,19 +179,91 @@ func (neo *DSL) handleUpload(c *gin.Context) { sid = uuid.New().String() } - // Set the context - ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "") - defer cancel() - - // Upload the file - file, err := neo.Upload(ctx, c) + uid, err := neo.UserOrGuestID(sid) if err != nil { - c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.JSON(401, gin.H{"message": fmt.Sprintf("Unauthorized, %s", err.Error()), "code": 401}) c.Done() return } - c.JSON(200, file) + if uid == nil || uid == "" { + c.JSON(401, gin.H{"message": "Unauthorized", "code": 401}) + c.Done() + return + } + + // Storage name must be chat, knowledge or assets + storage := c.Param("storage") + if storage != "chat" && storage != "knowledge" && storage != "assets" { + c.JSON(400, gin.H{"message": "Invalid storage", "code": 400}) + c.Done() + return + } + + // Get the manager + var manager, ok = attachment.Managers[storage] + if !ok { + c.JSON(400, gin.H{"message": "Invalid storage: " + storage, "code": 400}) + c.Done() + return + } + + // Get Option from form data + var option attachment.UploadOption + err = c.ShouldBind(&option) + if err != nil { + c.JSON(400, gin.H{"message": err.Error(), "code": 400}) + c.Done() + return + } + + // Validate the option with the storage + option.UserID = fmt.Sprintf("%v", uid) + + if storage == "chat" { + if option.ChatID == "" { + c.JSON(400, gin.H{"message": "chat_id is required", "code": 400}) + c.Done() + return + } + } + + // Get the file + file, err := c.FormFile("file") + if err != nil { + c.JSON(400, gin.H{"message": err.Error(), "code": 400}) + c.Done() + return + } + + // Open the file + reader, err := file.Open() + if err != nil { + c.JSON(400, gin.H{"message": err.Error(), "code": 400}) + c.Done() + return + } + defer func() { + reader.Close() + os.Remove(file.Filename) + }() + + // Convert the header to a FileHeader + header, err := attachment.ToFileHeader(file.Header) + if err != nil { + c.JSON(400, gin.H{"message": err.Error(), "code": 400}) + c.Done() + return + } + + // Upload the file + res, err := manager.Upload(c.Request.Context(), header, reader, option) + if err != nil { + c.JSON(400, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + c.JSON(200, map[string]interface{}{"data": res}) c.Done() } @@ -402,7 +476,7 @@ func (neo *DSL) corsMiddleware(allowsMap map[string]bool) gin.HandlerFunc { // Set CORS headers c.Header("Access-Control-Allow-Origin", origin) c.Header("Access-Control-Allow-Credentials", "true") - c.Header("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, Accept, Origin, Cache-Control, X-Requested-With") + c.Header("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, Accept, Origin, Cache-Control, X-Requested-With, Content-Sync, Content-Uid, Content-Range") c.Header("Access-Control-Allow-Methods", "GET, POST, OPTIONS") if c.Request.Method == "OPTIONS" { @@ -420,7 +494,7 @@ func (neo *DSL) optionsHandler(c *gin.Context) { if origin != "" { c.Header("Access-Control-Allow-Origin", origin) c.Header("Access-Control-Allow-Methods", "GET, POST, DELETE, OPTIONS") - c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization, Accept") + c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization, Accept, Content-Sync, Content-Uid, Content-Range") c.Header("Access-Control-Allow-Credentials", "true") c.Header("Access-Control-Max-Age", "86400") // 24 hours } diff --git a/neo/attachment/manager.go b/neo/attachment/manager.go index c3c85aa1..84cda782 100644 --- a/neo/attachment/manager.go +++ b/neo/attachment/manager.go @@ -8,6 +8,8 @@ import ( "fmt" "io" "mime" + "mime/multipart" + "net/textproto" "os" "path/filepath" "strconv" @@ -36,6 +38,25 @@ func Register(name string, driver string, option ManagerOption) (*Manager, error return manager, nil } +// ToFileHeader converts a multipart.FileHeader or textproto.MIMEHeader to a FileHeader +func ToFileHeader(header interface{}) (*FileHeader, error) { + + switch header := header.(type) { + case *multipart.FileHeader: + return &FileHeader{ + FileHeader: header, + }, nil + case textproto.MIMEHeader: + return &FileHeader{ + FileHeader: &multipart.FileHeader{ + Header: header, + }, + }, nil + default: + return nil, fmt.Errorf("invalid header type: %T", header) + } +} + // RegisterDefault registers a default attachment manager func RegisterDefault(name string) (*Manager, error) { @@ -60,6 +81,22 @@ func RegisterDefault(name string) (*Manager, error) { "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "application/vnd.openxmlformats-officedocument.presentationml.presentation", "application/vnd.openxmlformats-officedocument.presentationml.slideshow", + ".md", + ".txt", + ".csv", + ".xls", + ".xlsx", + ".ppt", + ".pptx", + ".doc", + ".docx", + ".mdx", + ".m4a", + ".mp3", + ".mp4", + ".wav", + ".webm", + ".yao", }, } return Register(name, option.Driver, option) @@ -363,7 +400,14 @@ func (manager Manager) makeFile(file *FileHeader, option UploadOption) (*File, e // Get the content type contentType := file.Header.Get("Content-Type") - extension := filepath.Ext(file.Filename) + + // Use original filename if provided, otherwise use the file header filename + filename := file.Filename + if option.OriginalFilename != "" { + filename = option.OriginalFilename + } + + extension := filepath.Ext(filename) // Get the extension from the content type if not available from filename if extension == "" { @@ -400,7 +444,7 @@ func (manager Manager) makeFile(file *FileHeader, option UploadOption) (*File, e // Validate allowed types if !manager.allowed(contentType, extension) { - return nil, fmt.Errorf("%s type %s is not allowed", file.Filename, contentType) + return nil, fmt.Errorf("%s type %s is not allowed", filename, contentType) } // Generate file ID @@ -411,7 +455,7 @@ func (manager Manager) makeFile(file *FileHeader, option UploadOption) (*File, e return &File{ ID: id, - Filename: file.Filename, + Filename: filename, // Use the correct filename (original or from header) ContentType: contentType, Bytes: int(file.Size), CreatedAt: int(time.Now().Unix()), @@ -445,6 +489,12 @@ func (manager Manager) allowed(contentType string, extension string) bool { // generateFileID generates a file ID with proper namespace func (manager Manager) generateFileID(file *FileHeader, extension string, option UploadOption) (string, error) { filename := file.Filename + + // Use original filename if provided for better file identification + if option.OriginalFilename != "" { + filename = option.OriginalFilename + } + if file.IsChunk() { filename = file.UID() } @@ -464,7 +514,7 @@ func (manager Manager) generateFileID(file *FileHeader, extension string, option path = filepath.Join(path, option.AssistantID) } - return filepath.Join(path, hash[:2], hash[2:4], hash, extension), nil + return filepath.Join(path, hash[:2], hash[2:4], hash) + extension, nil } // getSize converts the size to bytes diff --git a/neo/attachment/types.go b/neo/attachment/types.go index 5005a4ea..847c3852 100644 --- a/neo/attachment/types.go +++ b/neo/attachment/types.go @@ -77,14 +77,14 @@ type allowedType struct { // UploadOption the upload option type UploadOption struct { - CompressImage bool `json:"compress_image,omitempty"` // Compress the file, Optional, default is true - CompressSize int `json:"compress_size,omitempty"` // Compress the file size, Optional, default is 1920, if compress_image is true, the file size will be compressed to the compress_size - Gzip bool `json:"gzip,omitempty"` // Gzip the file, Optional, default is false - Knowledge bool `json:"knowledge,omitempty"` // Push to knowledge base, Optional, default is false - ChatID string `json:"chat_id,omitempty"` // Chat ID, Optional - AssistantID string `json:"assistant_id,omitempty"` // Assistant ID, Optional - UserID string `json:"user_id,omitempty"` // User ID, Optional - ChunkSize string `json:"chunk_size,omitempty"` // Chunk size of the file, Optional, default is 2M, default is 2M + CompressImage bool `json:"compress_image,omitempty" form:"compress_image"` // Compress the file, Optional, default is true + CompressSize int `json:"compress_size,omitempty" form:"compress_size"` // Compress the file size, Optional, default is 1920, if compress_image is true, the file size will be compressed to the compress_size + Gzip bool `json:"gzip,omitempty" form:"gzip"` // Gzip the file, Optional, default is false + Knowledge bool `json:"knowledge,omitempty" form:"knowledge"` // Push to knowledge base, Optional, default is false + ChatID string `json:"chat_id,omitempty" form:"chat_id"` // Chat ID, Optional + AssistantID string `json:"assistant_id,omitempty" form:"assistant_id"` // Assistant ID, Optional + UserID string `json:"user_id,omitempty"` // User ID, Optional + OriginalFilename string `json:"original_filename,omitempty" form:"original_filename"` // Original filename sent separately to avoid encoding issues } // FileHeader the file header diff --git a/neo/load.go b/neo/load.go index b388b247..6d6d5394 100644 --- a/neo/load.go +++ b/neo/load.go @@ -6,6 +6,7 @@ import ( "github.com/yaoapp/gou/application" "github.com/yaoapp/gou/connector" + "github.com/yaoapp/gou/model" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/neo/assistant" "github.com/yaoapp/yao/neo/attachment" @@ -77,6 +78,12 @@ func Load(cfg config.Config) error { return err } + // Initialize Auth + err = initAuth() + if err != nil { + return err + } + // Initialize Upload err = initUpload() if err != nil { @@ -92,6 +99,79 @@ func Load(cfg config.Config) error { return nil } +// initAuth initialize the auth +func initAuth() error { + if Neo.AuthSetting == nil { + Neo.AuthSetting = &Auth{ + Models: &AuthModels{User: "admin.user", Guest: "guest"}, + Fields: &AuthFields{ID: "id", Roles: "roles", Permission: "permission"}, + SessionFields: &AuthSessionFields{ID: "user_id", Roles: "user_roles", Guest: "guest_id"}, + } + } + + if Neo.AuthSetting.Models == nil { + Neo.AuthSetting.Models = &AuthModels{User: "admin.user", Guest: "guest"} + } + + if Neo.AuthSetting.Fields == nil { + Neo.AuthSetting.Fields = &AuthFields{ID: "id", Roles: "roles", Permission: "permission"} + } + + if Neo.AuthSetting.SessionFields == nil { + Neo.AuthSetting.SessionFields = &AuthSessionFields{ID: "user_id", Roles: "user_roles", Guest: "guest_id"} + } + + if Neo.AuthSetting.Models.User == "" { + Neo.AuthSetting.Models.User = "admin.user" + } + + if Neo.AuthSetting.Models.Guest == "" { + Neo.AuthSetting.Models.Guest = "guest" + } + + if Neo.AuthSetting.Fields.Roles == "" { + Neo.AuthSetting.Fields.Roles = "roles" + } + + if Neo.AuthSetting.Fields.Permission == "" { + Neo.AuthSetting.Fields.Permission = "permission" + } + + if Neo.AuthSetting.Fields.ID == "" { + Neo.AuthSetting.Fields.ID = "id" + } + + if Neo.AuthSetting.Fields.ID == "" { + Neo.AuthSetting.Fields.ID = "id" + } + + if Neo.AuthSetting.SessionFields.ID == "" { + Neo.AuthSetting.SessionFields.ID = "user_id" + } + + if Neo.AuthSetting.SessionFields.Roles == "" { + Neo.AuthSetting.SessionFields.Roles = "user_roles" + } + + if Neo.AuthSetting.SessionFields.Guest == "" { + Neo.AuthSetting.SessionFields.Guest = "guest_id" + } + + // Validate User Model and Fields + if !model.Exists(Neo.AuthSetting.Models.User) { + return fmt.Errorf("model %s not found", Neo.AuthSetting.Models.User) + } + user := model.Select(Neo.AuthSetting.Models.User) + shouldHave := []string{Neo.AuthSetting.Fields.ID, Neo.AuthSetting.Fields.Roles, Neo.AuthSetting.Fields.Permission} + for _, name := range shouldHave { + if _, has := user.Columns[name]; !has { + return fmt.Errorf("model %s should have column %s", Neo.AuthSetting.Models.User, name) + } + } + + return nil +} + // initUpload initialize the upload func initUpload() error { diff --git a/neo/neo.go b/neo/neo.go index 8b63a7ed..4702a8a0 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -2,10 +2,9 @@ package neo import ( "fmt" - "os" - "strings" "github.com/gin-gonic/gin" + "github.com/yaoapp/gou/session" "github.com/yaoapp/yao/neo/assistant" chatctx "github.com/yaoapp/yao/neo/context" ) @@ -32,47 +31,31 @@ func (neo *DSL) Select(id string) (assistant.API, error) { return assistant.Get(id) } -// Upload upload a file -func (neo *DSL) Upload(ctx chatctx.Context, c *gin.Context) (*assistant.File, error) { - // Get the file - tmpfile, err := c.FormFile("file") +// UserID get the user id from the session +func (neo *DSL) UserID(sid string) (interface{}, error) { + fieldID := neo.AuthSetting.SessionFields.ID + return session.Global().ID(sid).Get(fieldID) +} + +// GuestID get the guest id from the session +func (neo *DSL) GuestID(sid string) (interface{}, error) { + fieldGuest := neo.AuthSetting.SessionFields.Guest + return session.Global().ID(sid).Get(fieldGuest) +} + +// UserRoles get the user roles from the session +func (neo *DSL) UserRoles(sid string) (interface{}, error) { + fieldRoles := neo.AuthSetting.SessionFields.Roles + return session.Global().ID(sid).Get(fieldRoles) +} + +// UserOrGuestID get the user id or guest id from the session +func (neo *DSL) UserOrGuestID(sid string) (interface{}, error) { + userID, err := neo.UserID(sid) if err != nil { - return nil, err + return neo.GuestID(sid) } - - reader, err := tmpfile.Open() - if err != nil { - return nil, err - } - defer func() { - reader.Close() - os.Remove(tmpfile.Filename) - }() - - // Get option from form data option_xxx - option := map[string]interface{}{} - for key := range c.Request.Form { - if strings.HasPrefix(key, "option_") { - option[strings.TrimPrefix(key, "option_")] = c.PostForm(key) - } - } - - // Default use the assistant in context - ast := neo.Assistant - if ctx.ChatID == "" { - if ctx.AssistantID == "" { - return nil, fmt.Errorf("assistant_id is required") - } - ast, err = neo.Select(ctx.AssistantID) - if err != nil { - return nil, err - } - } - fmt.Println(ast) - - return nil, nil - - // return ast.Upload(ctx, tmpfile, reader, option) + return userID, nil } // Download downloads a file diff --git a/neo/types.go b/neo/types.go index fa160e7d..1b9c99f9 100644 --- a/neo/types.go +++ b/neo/types.go @@ -68,7 +68,7 @@ type AuthModels struct { type AuthSessionFields struct { ID string `json:"id,omitempty" yaml:"id,omitempty"` // the field name of the user id, default is user_id Roles string `json:"roles,omitempty" yaml:"roles,omitempty"` // the field name of the user roles, default is user_roles. the value must be an JSON array string. - Guest string `json:"guest,omitempty" yaml:"guest,omitempty"` // the field name of the guest user, default is guest + Guest string `json:"guest,omitempty" yaml:"guest,omitempty"` // the field name of the guest user, default is guest_id } // AuthFields the auth field diff --git a/widgets/app/app.go b/widgets/app/app.go index 2a865365..aff516e1 100644 --- a/widgets/app/app.go +++ b/widgets/app/app.go @@ -558,7 +558,28 @@ func processXgen(process *process.Process) interface{} { "placeholder": ast.GetPlaceholder(lang), } } + + // Available connectors agent["connectors"] = connector.AIConnectors + + // Available storages + agent["storages"] = map[string]interface{}{ + "chat": map[string]interface{}{ + "max_size": neo.Neo.UploadSetting.Chat.MaxSize, + "chunk_size": neo.Neo.UploadSetting.Chat.ChunkSize, + "allowed_types": neo.Neo.UploadSetting.Chat.AllowedTypes, + }, + "assets": map[string]interface{}{ + "max_size": neo.Neo.UploadSetting.Assets.MaxSize, + "chunk_size": neo.Neo.UploadSetting.Assets.ChunkSize, + "allowed_types": neo.Neo.UploadSetting.Assets.AllowedTypes, + }, + "knowledge": map[string]interface{}{ + "max_size": neo.Neo.UploadSetting.Knowledge.MaxSize, + "chunk_size": neo.Neo.UploadSetting.Knowledge.ChunkSize, + "allowed_types": neo.Neo.UploadSetting.Knowledge.AllowedTypes, + }, + } } xgenSetting := map[string]interface{}{