diff --git a/neo/attachment/manager.go b/neo/attachment/manager.go index 63421328..f2ee753a 100644 --- a/neo/attachment/manager.go +++ b/neo/attachment/manager.go @@ -8,11 +8,13 @@ import ( "fmt" "io" "mime" + "os" "path/filepath" "strconv" "strings" "time" + "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/neo/attachment/local" "github.com/yaoapp/yao/neo/attachment/s3" ) @@ -23,11 +25,6 @@ var Managers = map[string]*Manager{} // Register registers a global attachment manager func Register(name string, driver string, option ManagerOption) (*Manager, error) { - // Check if the manager already exists - if _, ok := Managers[name]; ok { - return nil, fmt.Errorf("manager %s already exists", name) - } - // Create a new manager manager, err := New(option) if err != nil { @@ -39,6 +36,57 @@ func Register(name string, driver string, option ManagerOption) (*Manager, error return manager, nil } +// RegisterDefault registers a default attachment manager +func RegisterDefault(name string) (*Manager, error) { + + option := ManagerOption{ + Driver: "local", + Options: map[string]interface{}{"path": filepath.Join(config.Conf.DataRoot, "attachments")}, + MaxSize: "50M", + ChunkSize: "2M", + AllowedTypes: []string{ + "text/*", + "image/*", + "video/*", + "audio/*", + "application/x-zip-compressed", + "application/x-tar", + "application/x-gzip", + "application/yao", + "application/zip", + "application/pdf", + "application/json", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + "application/vnd.openxmlformats-officedocument.presentationml.slideshow", + }, + } + return Register(name, option.Driver, option) +} + +// ReplaceEnv replaces the environment variables in the options +func (option *ManagerOption) ReplaceEnv(root string) { + if option.Options != nil { + // Replace the environment variables in the options + for k, v := range option.Options { + if iv, ok := v.(string); ok { + if strings.HasPrefix(iv, "$ENV.") { + iv = os.ExpandEnv(fmt.Sprintf("${%s}", strings.TrimPrefix(iv, "$ENV."))) + option.Options[k] = iv + } + + // Path + if k == "path" { + iv = strings.TrimPrefix(iv, "/") + option.Options[k] = filepath.Join(root, iv) + } + + } + } + } +} + // New creates a new attachment manager func New(option ManagerOption) (*Manager, error) { manager := &Manager{ diff --git a/neo/attachment/types.go b/neo/attachment/types.go index c53559fc..5005a4ea 100644 --- a/neo/attachment/types.go +++ b/neo/attachment/types.go @@ -63,11 +63,11 @@ type Storage interface { // ManagerOption the manager option type ManagerOption struct { - MaxSize string `json:"max_size,omitempty"` // Max size of the file, Optional, default is 20M - ChunkSize string `json:"chunk_size,omitempty"` // Chunk size of the file, Optional, default is 2M - AllowedTypes []string `json:"allowed_types,omitempty"` // Allowed types of the file, Optional, default is all - Driver string `json:"driver,omitempty"` // Driver, Optional, default is local - Options map[string]interface{} `json:"options,omitempty"` // Options, Optional + MaxSize string `json:"max_size,omitempty" yaml:"max_size,omitempty"` // Max size of the file, Optional, default is 20M + ChunkSize string `json:"chunk_size,omitempty" yaml:"chunk_size,omitempty"` // Chunk size of the file, Optional, default is 2M + AllowedTypes []string `json:"allowed_types,omitempty" yaml:"allowed_types,omitempty"` // Allowed types of the file, Optional, default is all + Driver string `json:"driver,omitempty" yaml:"driver,omitempty"` // Driver, Optional, default is local + Options map[string]interface{} `json:"options,omitempty" yaml:"options,omitempty"` // Options, Optional } type allowedType struct { diff --git a/neo/load.go b/neo/load.go index 642fce01..959f96bd 100644 --- a/neo/load.go +++ b/neo/load.go @@ -8,6 +8,7 @@ import ( "github.com/yaoapp/gou/connector" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/neo/assistant" + "github.com/yaoapp/yao/neo/attachment" "github.com/yaoapp/yao/neo/i18n" "github.com/yaoapp/yao/neo/store" ) @@ -76,6 +77,12 @@ func Load(cfg config.Config) error { return err } + // Initialize Upload + err = initUpload() + if err != nil { + return err + } + // Initialize Assistant err = initAssistant() if err != nil { @@ -85,6 +92,65 @@ func Load(cfg config.Config) error { return nil } +// initUpload initialize the upload +func initUpload() error { + + if Neo.UploadSetting == nil { + _, err := attachment.RegisterDefault("chat") + if err != nil { + return err + } + _, err = attachment.RegisterDefault("knowledge") + if err != nil { + return err + } + return nil + } + + // If the chat upload setting is not set, use the default chat upload setting. + if Neo.UploadSetting.Chat == nil { + _, err := attachment.RegisterDefault("chat") + if err != nil { + return err + } + } + + // Use the chat upload setting for knowledge upload, if the knowledge upload setting is not set. + if Neo.UploadSetting.Knowledge == nil { + if Neo.UploadSetting.Chat == nil { + _, err := attachment.RegisterDefault("knowledge") + if err != nil { + return err + } + } else { + _, err := attachment.Register("knowledge", Neo.UploadSetting.Chat.Driver, *Neo.UploadSetting.Chat) + if err != nil { + return err + } + } + } + + // Use custom chat upload setting + if Neo.UploadSetting.Chat != nil { + Neo.UploadSetting.Chat.ReplaceEnv(config.Conf.DataRoot) + _, err := attachment.Register("chat", Neo.UploadSetting.Chat.Driver, *Neo.UploadSetting.Chat) // Register the chat upload manager + if err != nil { + return err + } + } + + // Use custom knowledge upload setting + if Neo.UploadSetting.Knowledge != nil { + Neo.UploadSetting.Knowledge.ReplaceEnv(config.Conf.DataRoot) + _, err := attachment.Register("knowledge", Neo.UploadSetting.Knowledge.Driver, *Neo.UploadSetting.Knowledge) + if err != nil { + return err + } + } + + return nil +} + // initGlobalI18n initialize the global i18n func initGlobalI18n() error { locales, err := i18n.GetLocales("neo") diff --git a/neo/types.go b/neo/types.go index 96a7ca06..ef711487 100644 --- a/neo/types.go +++ b/neo/types.go @@ -3,6 +3,7 @@ package neo import ( "github.com/gin-gonic/gin" "github.com/yaoapp/yao/neo/assistant" + "github.com/yaoapp/yao/neo/attachment" "github.com/yaoapp/yao/neo/rag" "github.com/yaoapp/yao/neo/store" "github.com/yaoapp/yao/neo/vision" @@ -80,11 +81,8 @@ type AuthFields struct { // Upload the upload setting // =============================== type Upload struct { - Driver string `json:"driver" yaml:"driver"` // local, s3, default is local - Options map[string]interface{} `json:"options" yaml:"options"` // the options of the upload, it is used to configure the upload driver. - Compression bool `json:"compression,omitempty" yaml:"compression,omitempty"` // Compress the image/video to a smaller size, if the image/video is too large, it will be compressed to a smaller size. - ChunkSize string `json:"chunk_size,omitempty" yaml:"chunk_size,omitempty"` // the chunk size of the file, if the file is too large, it will be chunked into smaller chunks. - AllowedTypes []string `json:"allowed_types,omitempty" yaml:"allowed_types,omitempty"` // the allowed types of the file, if the file is not in the allowed types, it will be rejected. + Chat *attachment.ManagerOption `json:"chat,omitempty" yaml:"chat,omitempty"` // Chat conversation upload setting, if not set use the local and root path is `/attachments`. + Knowledge *attachment.ManagerOption `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Knowledge base upload setting, if not set use the chat upload setting. } // Knowledge base Settings