opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Enhance file upload functionality and authentication in neo package

Browse source 
- 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.
This commit is contained in:
Max 2025-05-31 16:36:02 +08:00
parent 43686be6e4
commit 5b99f3bf44
7 changed files with 274 additions and 66 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

98
neo/api.go
View file

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

58
neo/attachment/manager.go
View file

@ -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

16
neo/attachment/types.go
View file

@ -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

80
neo/load.go
View file

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

65
neo/neo.go
View file

@ -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

2
neo/types.go
View file

@ -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

21
widgets/app/app.go
View file

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