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

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

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

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)
+}

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