Add upload management to neo package for enhanced file handling
- Introduced initUpload function to initialize upload settings for chat and knowledge attachments. - Updated Upload struct to include specific settings for chat and knowledge uploads. - Implemented RegisterDefault and ReplaceEnv methods in the attachment manager for better configuration management. - Enhanced error handling during the upload initialization process to ensure robust functionality.
This commit is contained in:
Max
parent
96d2f22e77
commit
a9b6eed025
4 changed files with 127 additions and 15 deletions
|
|
@ -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{
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
66
neo/load.go
66
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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue