Merge pull request #980 from trheyi/main

Enhance file upload functionality and authentication in neo package

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
+	uid, err := neo.UserOrGuestID(sid)
-	ctx, cancel := chatctx.NewWithCancel(sid, c.Query("chat_id"), "")
-	defer cancel()
-
-	// Upload the file
-	file, err := neo.Upload(ctx, c)
 	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
 	}

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

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
+	CompressImage    bool   `json:"compress_image,omitempty" form:"compress_image"`       // 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
+	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"`           // Gzip the file, Optional, default is false
+	Gzip             bool   `json:"gzip,omitempty" form:"gzip"`                           // Gzip the file, Optional, default is false
-	Knowledge     bool   `json:"knowledge,omitempty"`      // Push to knowledge base, 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"`        // Chat ID, Optional
+	ChatID           string `json:"chat_id,omitempty" form:"chat_id"`                     // Chat ID, Optional
-	AssistantID   string `json:"assistant_id,omitempty"`   // Assistant ID, Optional
+	AssistantID      string `json:"assistant_id,omitempty" form:"assistant_id"`           // 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
+	OriginalFilename string `json:"original_filename,omitempty" form:"original_filename"` // Original filename sent separately to avoid encoding issues
 }
 
 // FileHeader the file header

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 {
 

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
+// UserID get the user id from the session
-func (neo *DSL) Upload(ctx chatctx.Context, c *gin.Context) (*assistant.File, error) {
+func (neo *DSL) UserID(sid string) (interface{}, error) {
-	// Get the file
+	fieldID := neo.AuthSetting.SessionFields.ID
-	tmpfile, err := c.FormFile("file")
+	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)
 	}
-
+	return userID, nil
-	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)
 }
 
 // Download downloads a file

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

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{}{