From 2903602c6bee9546aa1b2b109f1fb632578039ed Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 2 Jan 2025 16:55:21 +0800 Subject: [PATCH] Refactor Neo API initialization and enhance assistant management - Introduced new methods for initializing RAG and store, improving modularity and clarity in the Load function. - Refactored the assistant initialization process to streamline the loading of built-in assistants and set RAG configurations. - Moved the defaultAssistant method from neo.go to load.go for better organization and accessibility. - Enhanced error handling in store initialization to support multiple connector types, including Redis and Mongo. - Updated the assistant struct to include RAG settings, improving the overall assistant management capabilities. --- neo/assistant/assistant.go | 14 +++++ neo/assistant/types.go | 8 +++ neo/load.go | 104 ++++++++++++++++++++++++++++++------- neo/neo.go | 14 ----- 4 files changed, 108 insertions(+), 32 deletions(-) diff --git a/neo/assistant/assistant.go b/neo/assistant/assistant.go index ee57d890..2460a44d 100644 --- a/neo/assistant/assistant.go +++ b/neo/assistant/assistant.go @@ -9,6 +9,7 @@ import ( jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/fs" + "github.com/yaoapp/gou/rag/driver" v8 "github.com/yaoapp/gou/runtime/v8" "github.com/yaoapp/yao/neo/store" "github.com/yaoapp/yao/share" @@ -18,6 +19,7 @@ import ( // loaded the loaded assistant var loaded = NewCache(200) // 200 is the default capacity var storage store.Store = nil +var rag *RAG = nil // LoadBuiltIn load the built-in assistants func LoadBuiltIn() error { @@ -85,6 +87,18 @@ func SetStorage(s store.Store) { storage = s } +// SetRAG set the RAG engine +// e: the RAG engine +// u: the RAG file uploader +// v: the RAG vectorizer +func SetRAG(e driver.Engine, u driver.FileUpload, v driver.Vectorizer) { + rag = &RAG{ + Engine: e, + Uploader: u, + Vectorizer: v, + } +} + // SetCache set the cache func SetCache(capacity int) { ClearCache() diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 977cf5e9..17f17bf0 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -5,6 +5,7 @@ import ( "io" "mime/multipart" + "github.com/yaoapp/gou/rag/driver" v8 "github.com/yaoapp/gou/runtime/v8" ) @@ -16,6 +17,13 @@ type API interface { ReadBase64(ctx context.Context, fileID string) (string, error) } +// RAG the RAG interface +type RAG struct { + Engine driver.Engine + Uploader driver.FileUpload + Vectorizer driver.Vectorizer +} + // Prompt a prompt type Prompt struct { Role string `json:"role"` diff --git a/neo/load.go b/neo/load.go index 34cf3704..cfafc65c 100644 --- a/neo/load.go +++ b/neo/load.go @@ -1,10 +1,12 @@ package neo import ( + "fmt" "path/filepath" "github.com/fatih/color" "github.com/yaoapp/gou/application" + "github.com/yaoapp/gou/connector" "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/neo/assistant" @@ -15,20 +17,6 @@ import ( // Neo the neo AI assistant var Neo *DSL -// initRAG initialize the RAG instance -func (neo *DSL) initRAG() { - if neo.RAGSetting.Engine.Driver == "" { - return - } - instance, err := rag.New(neo.RAGSetting) - if err != nil { - color.Red("[Neo] Failed to initialize RAG: %v", err) - log.Error("[Neo] Failed to initialize RAG: %v", err) - return - } - neo.RAG = instance -} - // Load load AIGC func Load(cfg config.Config) error { @@ -60,7 +48,7 @@ func Load(cfg config.Config) error { Neo = &setting // Store Setting - err = Neo.createStore() + err = Neo.initStore() if err != nil { return err } @@ -68,13 +56,79 @@ func Load(cfg config.Config) error { // Initialize RAG Neo.initRAG() - // Load Built-in Assistants - assistant.SetStorage(Neo.Store) - err = assistant.LoadBuiltIn() + // Initialize Assistant + err = Neo.initAssistant() if err != nil { return err } + return nil +} + +// initRAG initialize the RAG instance +func (neo *DSL) initRAG() { + if neo.RAGSetting.Engine.Driver == "" { + return + } + instance, err := rag.New(neo.RAGSetting) + if err != nil { + color.Red("[Neo] Failed to initialize RAG: %v", err) + log.Error("[Neo] Failed to initialize RAG: %v", err) + return + } + + neo.RAG = instance +} + +// initStore initialize the store +func (neo *DSL) initStore() error { + + var err error + if neo.StoreSetting.Connector == "default" || neo.StoreSetting.Connector == "" { + neo.Store, err = store.NewXun(neo.StoreSetting) + return err + } + + // other connector + conn, err := connector.Select(neo.StoreSetting.Connector) + if err != nil { + return err + } + + if conn.Is(connector.DATABASE) { + neo.Store, err = store.NewXun(neo.StoreSetting) + return err + + } else if conn.Is(connector.REDIS) { + neo.Store = store.NewRedis() + return nil + + } else if conn.Is(connector.MONGO) { + neo.Store = store.NewMongo() + return nil + } + + return fmt.Errorf("%s store connector %s not support", neo.ID, neo.StoreSetting.Connector) +} + +// initAssistant initialize the assistant +func (neo *DSL) initAssistant() error { + + // Set Storage + assistant.SetStorage(Neo.Store) + + // Assistant RAG + if Neo.RAG != nil { + assistant.SetRAG(Neo.RAG.Engine(), Neo.RAG.FileUpload(), Neo.RAG.Vectorizer()) + } + + // Load Built-in Assistants + err := assistant.LoadBuiltIn() + if err != nil { + return err + } + + // Default Assistant defaultAssistant, err := Neo.defaultAssistant() if err != nil { return err @@ -83,3 +137,17 @@ func Load(cfg config.Config) error { Neo.Assistant = defaultAssistant.API return nil } + +// defaultAssistant get the default assistant +func (neo *DSL) defaultAssistant() (*assistant.Assistant, error) { + if neo.Use != "" { + return assistant.Get(neo.Use) + } + + name := neo.Name + if name == "" { + name = "Neo" + } + + return assistant.GetByConnector(neo.Connector, name) +} diff --git a/neo/neo.go b/neo/neo.go index e5cf0399..f072dc1c 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -346,20 +346,6 @@ func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]inter } } -// defaultAssistant get the default assistant -func (neo *DSL) defaultAssistant() (*assistant.Assistant, error) { - if neo.Use != "" { - return assistant.Get(neo.Use) - } - - name := neo.Name - if name == "" { - name = "Neo" - } - - return assistant.GetByConnector(neo.Connector, name) -} - // updateAssistantList update the assistant list func (neo *DSL) updateAssistantList(list []assistant.Assistant) { lock.Lock()