Merge pull request #978 from trheyi/main

Add upload management to neo package for enhanced file handling

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{

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
+	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"`    // Chunk size of the file, Optional, default is 2M
+	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"` // Allowed types of the file, Optional, default is all
+	AllowedTypes []string               `json:"allowed_types,omitempty" yaml:"allowed_types,omitempty"` // Allowed types of the file, Optional, default is all
-	Driver       string                 `json:"driver,omitempty"`        // Driver, Optional, default is local
+	Driver       string                 `json:"driver,omitempty" yaml:"driver,omitempty"`               // Driver, Optional, default is local
-	Options      map[string]interface{} `json:"options,omitempty"`       // Options, Optional
+	Options      map[string]interface{} `json:"options,omitempty" yaml:"options,omitempty"`             // Options, Optional
 }
 
 type allowedType struct {

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")

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
+	Chat      *attachment.ManagerOption `json:"chat,omitempty" yaml:"chat,omitempty"`           // Chat conversation upload setting, if not set use the local and root path is `/attachments`.
-	Options      map[string]interface{} `json:"options" yaml:"options"`                                 // the options of the upload, it is used to configure the upload driver.
+	Knowledge *attachment.ManagerOption `json:"knowledge,omitempty" yaml:"knowledge,omitempty"` // Knowledge base upload setting, if not set use the chat upload setting.
-	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.
 }
 
 // Knowledge base Settings